Check primitive Typed Jerboa expressions

ober

551df56d072a552a8b9f8faa55e088250e7a3efd

diff --git a/docs/typed-jerboa.md b/docs/typed-jerboa.md
index 22f5bfa..9125647 100644
--- a/docs/typed-jerboa.md
+++ b/docs/typed-jerboa.md
@@ -139,14 +139,17 @@ Current landing:
   into a first AST layer.
 - `(jerboa typed checker)` performs the first backend-neutral validation pass:
   duplicate names, export resolution, type-reference resolution, compound type
-  arities, duplicate fields, duplicate params, and duplicate variant cases.
+  arities, duplicate fields, duplicate params, duplicate variant cases, and a
+  small expression subset for return checking.
 - `support/typecheck.ss` and the `make typecheck` / `make typed-test` targets
   run the parser/checker over typed source files without invoking Rust.
 - The parser currently recognizes explicit `export` forms, `type` declarations,
   immutable and `mut` record fields, variants including nullary cases, and
   `def` forms with typed parameters and `:` or `->` return markers.
-- This is still a front-end milestone. It does not yet typecheck function
-  bodies, resolve imports, lower to typed core IR, or emit Rust/LLVM.
+- This is still a front-end milestone. Function-body checking currently covers
+  literals, variables, `begin`, simple `let`, and `if`; calls and richer forms
+  are reported as unsupported. It does not yet resolve imports, lower to typed
+  core IR, or emit Rust/LLVM.
 
 ## Surface Syntax
 
@@ -735,7 +738,8 @@ Minimum excluded features:
 - Resolve names. Initial declaration/export/type-reference validation landed
   as `(jerboa typed checker)`.
 - Build module environment.
-- Check primitive expressions. Not yet landed.
+- Check primitive expressions. Initial literal/variable/`begin`/`let`/`if`
+  return checking landed.
 - Check function applications.
 - Check records and variants.
 - Check match exhaustiveness.
diff --git a/lib/jerboa/typed/checker.ss b/lib/jerboa/typed/checker.ss
index 29904bf..a3177c1 100644
--- a/lib/jerboa/typed/checker.ss
+++ b/lib/jerboa/typed/checker.ss
@@ -1,8 +1,8 @@
 #!chezscheme
 ;;; (jerboa typed checker) -- first Typed Jerboa validation pass
 ;;;
-;;; This checker validates names and type references over the parser AST.
-;;; It intentionally does not typecheck function bodies yet.
+;;; This checker validates names and type references over the parser AST, plus
+;;; a small expression subset for function return checking.
 
 (library (jerboa typed checker)
   (export
@@ -114,6 +114,20 @@
     (or (memq name builtin-type-names)
         (memq name type-names)))
 
+  (def (lookup-name name env)
+    (let ([entry (assq name env)])
+      (and entry (cdr entry))))
+
+  (def (type-assignable? actual expected)
+    (or (equal? actual expected)
+        (and (eq? actual 'Nat) (eq? expected 'Int))))
+
+  (def (param-env params)
+    (map (lambda (param)
+           (cons (typed-param-name param)
+                 (typed-param-type param)))
+         params))
+
   (def (check-compound-type type type-names)
     (let* ([head (car type)]
            [args (cdr type)]
@@ -151,6 +165,9 @@
                "invalid type expression"
                type))]))
 
+  (def (valid-type? type type-names)
+    (null? (check-type type type-names)))
+
   (def (check-fields fields type-names)
     (append
       (duplicate-errors 'duplicate-field
@@ -183,7 +200,159 @@
       (append-map
         (lambda (param) (check-type (typed-param-type param) type-names))
         (typed-def-params def))
-      (check-type (typed-def-return-type def) type-names)))
+      (check-type (typed-def-return-type def) type-names)
+      (check-def-body def type-names)))
+
+  (def (infer-body exprs env type-names)
+    (cond
+      [(null? exprs) (values 'Unit '())]
+      [else
+       (let loop ([rest exprs] [errors '()] [last-type 'Unit])
+         (if (null? rest)
+           (values last-type (reverse errors))
+           (let-values ([(expr-type expr-errors)
+                         (infer-expression (car rest) env type-names)])
+             (loop (cdr rest)
+                   (append (reverse expr-errors) errors)
+                   expr-type))))]))
+
+  (def (extend-env name type env)
+    (cons (cons name type) env))
+
+  (def (binding-name binding)
+    (and (pair? binding) (car binding)))
+
+  (def (check-let-binding-shape binding)
+    (cond
+      [(and (pair? binding)
+            (pair? (cdr binding))
+            (null? (cddr binding))
+            (symbol? (car binding)))
+       '()]
+      [else
+       (list (make-check-error 'bad-let-binding
+               "expected let binding (name expr)"
+               binding))]))
+
+  (def (infer-let bindings body env type-names expr)
+    (cond
+      [(not (list? bindings))
+       (values #f
+         (list (make-check-error 'bad-let-binding
+                 "let bindings must be a list"
+                 expr)))]
+      [else
+       (let ([shape-errors (append-map check-let-binding-shape bindings)]
+             [duplicate-name-errors
+              (duplicate-errors 'duplicate-local
+                "duplicate local binding"
+                (map binding-name bindings))])
+         (if (or (not (null? shape-errors))
+                 (not (null? duplicate-name-errors)))
+           (values #f (append duplicate-name-errors shape-errors))
+           (let loop ([rest bindings] [local-env env] [errors '()])
+             (if (null? rest)
+               (let-values ([(body-type body-errors)
+                             (infer-body body local-env type-names)])
+                 (values body-type (append (reverse errors) body-errors)))
+               (let* ([binding (car rest)]
+                      [name (car binding)]
+                      [value-expr (cadr binding)])
+                 (let-values ([(value-type value-errors)
+                               (infer-expression value-expr env type-names)])
+                   (loop (cdr rest)
+                         (if value-type
+                           (extend-env name value-type local-env)
+                           local-env)
+                         (append (reverse value-errors) errors))))))))]))
+
+  (def (infer-if args env type-names expr)
+    (if (not (= (length args) 3))
+      (values #f
+        (list (make-check-error 'bad-if
+                "if expects condition, then branch, and else branch"
+                expr)))
+      (let-values ([(cond-type cond-errors)
+                    (infer-expression (car args) env type-names)]
+                   [(then-type then-errors)
+                    (infer-expression (cadr args) env type-names)]
+                   [(else-type else-errors)
+                    (infer-expression (caddr args) env type-names)])
+        (let ([condition-errors
+               (if (and cond-type (not (eq? cond-type 'Bool)))
+                 (list (make-check-error 'condition-type-mismatch
+                         "if condition must be Bool"
+                         cond-type))
+                 '())]
+              [branch-errors
+               (if (and then-type else-type (not (equal? then-type else-type)))
+                 (list (make-check-error 'branch-type-mismatch
+                         "if branches must have the same type"
+                         (list then-type else-type)))
+                 '())])
+          (values
+            (if (null? branch-errors) then-type #f)
+            (append cond-errors then-errors else-errors
+                    condition-errors branch-errors))))))
+
+  (def (infer-expression expr env type-names)
+    (cond
+      [(boolean? expr) (values 'Bool '())]
+      [(char? expr) (values 'Char '())]
+      [(string? expr) (values 'String '())]
+      [(bytevector? expr) (values 'Bytes '())]
+      [(integer? expr)
+       (values (if (>= expr 0) 'Nat 'Int) '())]
+      [(symbol? expr)
+       (let ([type (lookup-name expr env)])
+         (if type
+           (values type '())
+           (values #f
+             (list (make-check-error 'unknown-value
+                     "unknown value in expression"
+                     expr)))))]
+      [(pair? expr)
+       (case (car expr)
+         [(begin)
+          (infer-body (cdr expr) env type-names)]
+         [(let)
+          (if (< (length expr) 3)
+            (values #f
+              (list (make-check-error 'bad-let
+                      "let expects bindings and body"
+                      expr)))
+            (infer-let (cadr expr) (cddr expr) env type-names expr))]
+         [(if)
+          (infer-if (cdr expr) env type-names expr)]
+         [else
+          (values #f
+            (list (make-check-error 'unsupported-expression
+                    "expression form is not in the typed checker subset yet"
+                    expr)))])]
+      [else
+       (values #f
+         (list (make-check-error 'unsupported-expression
+                 "expression is not in the typed checker subset yet"
+                 expr)))]))
+
+  (def (check-def-body def type-names)
+    (let-values ([(actual-type body-errors)
+                  (infer-body
+                    (typed-def-body def)
+                    (param-env (typed-def-params def))
+                    type-names)])
+      (append
+        body-errors
+        (if (and actual-type
+                 (valid-type? (typed-def-return-type def) type-names)
+                 (not (type-assignable? actual-type
+                       (typed-def-return-type def))))
+          (list (make-check-error 'return-type-mismatch
+                  "function body type does not match declared return type"
+                  (list (typed-def-name def)
+                        (typed-def-return-type def)
+                        actual-type)))
+          '()))))
 
   (def (check-declaration decl type-names)
     (cond
diff --git a/tests/test-typed-checker.ss b/tests/test-typed-checker.ss
index f1e2296..d0ab8cf 100644
--- a/tests/test-typed-checker.ss
+++ b/tests/test-typed-checker.ss
@@ -140,6 +140,67 @@
        (def (f (x : Nat)) : Nat x)))
   '(undefined-export))
 
+(test "body can return parameter"
+  (error-kinds
+    '(typed-library (body param)
+       (export id)
+       (def (id (x : String)) : String x)))
+  '())
+
+(test "body return mismatch"
+  (error-kinds
+    '(typed-library (body mismatch)
+       (export f)
+       (def (f (x : String)) : Nat x)))
+  '(return-type-mismatch))
+
+(test "body unknown value"
+  (error-kinds
+    '(typed-library (body unknown-value)
+       (export f)
+       (def (f (x : Nat)) : Nat y)))
+  '(unknown-value))
+
+(test "let infers local type"
+  (error-kinds
+    '(typed-library (body let-ok)
+       (export f)
+       (def (f (x : Nat)) : Nat
+         (let ([y x]) y))))
+  '())
+
+(test "let return mismatch"
+  (error-kinds
+    '(typed-library (body let-mismatch)
+       (export f)
+       (def (f (x : Nat)) : String
+         (let ([y x]) y))))
+  '(return-type-mismatch))
+
+(test "if condition type mismatch"
+  (error-kinds
+    '(typed-library (body if-condition)
+       (export f)
+       (def (f (x : Nat)) : Nat
+         (if x 1 2))))
+  '(condition-type-mismatch))
+
+(test "if branch type mismatch"
+  (error-kinds
+    '(typed-library (body if-branches)
+       (export f)
+       (def (f (x : Bool)) : Nat
+         (if x 1 "bad"))))
+  '(branch-type-mismatch))
+
+(test "unsupported call expression"
+  (error-kinds
+    '(typed-library (body unsupported)
+       (export f)
+       (def (f (x : Nat)) : Nat
+         (+ x 1))))
+  '(unsupported-expression))
+
 (printf "~%Typed checker: ~a passed, ~a failed~%" pass fail)
 (when (> fail 0)
   (exit 1))