Check Typed Jerboa primitive operators

ober

bdf475e13eca6634744c9b87fee1969451777c7d

diff --git a/docs/typed-jerboa.md b/docs/typed-jerboa.md
index 9125647..9908d8e 100644
--- a/docs/typed-jerboa.md
+++ b/docs/typed-jerboa.md
@@ -147,7 +147,8 @@ Current landing:
   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. Function-body checking currently covers
-  literals, variables, `begin`, simple `let`, and `if`; calls and richer forms
+  literals, variables, `begin`, simple `let`, `if`, arithmetic primitives,
+  numeric comparisons, and boolean primitives. General calls and richer forms
   are reported as unsupported. It does not yet resolve imports, lower to typed
   core IR, or emit Rust/LLVM.
 
@@ -738,8 +739,8 @@ Minimum excluded features:
 - Resolve names. Initial declaration/export/type-reference validation landed
   as `(jerboa typed checker)`.
 - Build module environment.
-- Check primitive expressions. Initial literal/variable/`begin`/`let`/`if`
-  return checking landed.
+- Check primitive expressions. Initial literal/variable/`begin`/`let`/`if`,
+  arithmetic, comparison, and boolean primitive 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 a3177c1..de57e64 100644
--- a/lib/jerboa/typed/checker.ss
+++ b/lib/jerboa/typed/checker.ss
@@ -295,6 +295,90 @@
             (append cond-errors then-errors else-errors
                     condition-errors branch-errors))))))
 
+  (def (numeric-type? type)
+    (and type (memq type '(Nat Int Fixnum Float)) #t))
+
+  (def (merge-numeric-types types)
+    (cond
+      [(memq 'Float types) 'Float]
+      [(memq 'Int types) 'Int]
+      [(memq 'Fixnum types) 'Fixnum]
+      [else 'Nat]))
+
+  (def (infer-args args env type-names)
+    (let loop ([rest args] [types '()] [errors '()])
+      (if (null? rest)
+        (values (reverse types) (reverse errors))
+        (let-values ([(type arg-errors)
+                      (infer-expression (car rest) env type-names)])
+          (loop (cdr rest)
+                (cons type types)
+                (append (reverse arg-errors) errors))))))
+
+  (def (operand-type-errors expected-kind types detail)
+    (let loop ([rest types] [out '()])
+      (cond
+        [(null? rest) (reverse out)]
+        [(not (car rest))
+         (loop (cdr rest) out)]
+        [(eq? expected-kind 'numeric)
+         (if (numeric-type? (car rest))
+           (loop (cdr rest) out)
+           (loop (cdr rest)
+                 (cons (make-check-error 'operand-type-mismatch
+                         "expected numeric operand"
+                         detail)
+                       out)))]
+        [(eq? expected-kind 'Bool)
+         (if (eq? (car rest) 'Bool)
+           (loop (cdr rest) out)
+           (loop (cdr rest)
+                 (cons (make-check-error 'operand-type-mismatch
+                         "expected Bool operand"
+                         detail)
+                       out)))]
+        [else (loop (cdr rest) out)])))
+
+  (def (infer-arithmetic op args env type-names expr)
+    (if (null? args)
+      (values #f
+        (list (make-check-error 'bad-primitive-arity
+                "arithmetic primitive needs at least one operand"
+                expr)))
+      (let-values ([(types errors) (infer-args args env type-names)])
+        (let ([operand-errors (operand-type-errors 'numeric types expr)])
+          (values
+            (if (and (null? errors) (null? operand-errors))
+              (merge-numeric-types types)
+              #f)
+            (append errors operand-errors))))))
+
+  (def (infer-comparison op args env type-names expr)
+    (if (not (= (length args) 2))
+      (values #f
+        (list (make-check-error 'bad-primitive-arity
+                "comparison primitive needs exactly two operands"
+                expr)))
+      (let-values ([(types errors) (infer-args args env type-names)])
+        (let ([operand-errors (operand-type-errors 'numeric types expr)])
+          (values
+            (if (and (null? errors) (null? operand-errors)) 'Bool #f)
+            (append errors operand-errors))))))
+
+  (def (infer-boolean op args env type-names expr)
+    (cond
+      [(and (eq? op 'not) (not (= (length args) 1)))
+       (values #f
+         (list (make-check-error 'bad-primitive-arity
+                 "not needs exactly one operand"
+                 expr)))]
+      [else
+       (let-values ([(types errors) (infer-args args env type-names)])
+         (let ([operand-errors (operand-type-errors 'Bool types expr)])
+           (values
+             (if (and (null? errors) (null? operand-errors)) 'Bool #f)
+             (append errors operand-errors))))]))
+
   (def (infer-expression expr env type-names)
     (cond
       [(boolean? expr) (values 'Bool '())]
@@ -324,6 +408,12 @@
             (infer-let (cadr expr) (cddr expr) env type-names expr))]
          [(if)
           (infer-if (cdr expr) env type-names expr)]
+         [(+ - * /)
+          (infer-arithmetic (car expr) (cdr expr) env type-names expr)]
+         [(= < <= > >=)
+          (infer-comparison (car expr) (cdr expr) env type-names expr)]
+         [(not and or)
+          (infer-boolean (car expr) (cdr expr) env type-names expr)]
          [else
           (values #f
             (list (make-check-error 'unsupported-expression
diff --git a/tests/fixtures/typed/valid-split-tree.ss b/tests/fixtures/typed/valid-split-tree.ss
index c8bff21..9d20e8a 100644
--- a/tests/fixtures/typed/valid-split-tree.ss
+++ b/tests/fixtures/typed/valid-split-tree.ss
@@ -11,4 +11,4 @@
     (Noop))
 
   (def (split-size (pane : Pane)) : Nat
-    0))
+    (+ 0 0)))
diff --git a/tests/test-typed-checker.ss b/tests/test-typed-checker.ss
index d0ab8cf..474f95c 100644
--- a/tests/test-typed-checker.ss
+++ b/tests/test-typed-checker.ss
@@ -198,9 +198,57 @@
     '(typed-library (body unsupported)
        (export f)
        (def (f (x : Nat)) : Nat
-         (+ x 1))))
+         (foo x))))
   '(unsupported-expression))
 
+(test "arithmetic primitive returns numeric type"
+  (error-kinds
+    '(typed-library (body arithmetic)
+       (export f)
+       (def (f (x : Nat)) : Nat
+         (+ x 1))))
+  '())
+
+(test "arithmetic primitive rejects nonnumeric operand"
+  (error-kinds
+    '(typed-library (body arithmetic-bad)
+       (export f)
+       (def (f (x : String)) : Nat
+         (+ x 1))))
+  '(operand-type-mismatch))
+
+(test "comparison primitive feeds if condition"
+  (error-kinds
+    '(typed-library (body compare)
+       (export f)
+       (def (f (x : Nat)) : Nat
+         (if (> x 0) x 0))))
+  '())
+
+(test "comparison primitive rejects nonnumeric operand"
+  (error-kinds
+    '(typed-library (body compare-bad)
+       (export f)
+       (def (f (x : String)) : Bool
+         (> x 0))))
+  '(operand-type-mismatch))
+
+(test "boolean primitive returns Bool"
+  (error-kinds
+    '(typed-library (body bool)
+       (export f)
+       (def (f (x : Bool) (y : Bool)) : Bool
+         (and x (not y)))))
+  '())
+
+(test "boolean primitive rejects non-Bool operand"
+  (error-kinds
+    '(typed-library (body bool-bad)
+       (export f)
+       (def (f (x : Nat)) : Bool
+         (not x))))
+  '(operand-type-mismatch))
+
 (printf "~%Typed checker: ~a passed, ~a failed~%" pass fail)
 (when (> fail 0)
   (exit 1))