typed/llvmir: lower scalar functions (Phase 2)

ober

305d47b21bebc2fef754e1cee342e5ec65594c5b

diff --git a/lib/jerboa/typed/llvmir.ss b/lib/jerboa/typed/llvmir.ss
index fa5e568..dc7b478 100644
--- a/lib/jerboa/typed/llvmir.ss
+++ b/lib/jerboa/typed/llvmir.ss
@@ -208,6 +208,9 @@
       [(typed-ir-var? ir) (lower-var env ir)]
       [(typed-ir-begin? ir)
        (lower-begin env (typed-ir-begin-exprs ir) (typed-ir-begin-type ir))]
+      [(typed-ir-let? ir) (lower-let env ir)]
+      [(typed-ir-if? ir) (lower-if env ir)]
+      [(typed-ir-call? ir) (lower-call env ir)]
       [else
        (error 'typed-llvmir "unsupported typed IR node for LLVM lowering" ir)]))
 
@@ -238,6 +241,263 @@
              value
              (loop (cdr rest)))))]))
 
+  ;; Immutable let needs no stack slot: bind each name to the SSA value of its
+  ;; initializer. Bindings are sequential (matching the Rust backend's `let`
+  ;; statements); the outer variable scope is restored after the body.
+  (def (lower-let env ir)
+    (let ([saved-vars (llvm-env-vars env)])
+      (for-each
+        (lambda (binding)
+          (let ([value (lower-expr env (typed-ir-binding-expr binding))])
+            (bind-var! env (typed-ir-binding-name binding) value)))
+        (typed-ir-let-bindings ir))
+      (let ([result
+             (lower-begin env (typed-ir-let-body ir) (typed-ir-let-type ir))])
+        (llvm-env-vars-set! env saved-vars)
+        result)))
+
+  ;; Expression-valued conditionals use real basic blocks and a phi at the
+  ;; join. The phi's predecessor labels are the blocks each branch *ends* in
+  ;; (nested control flow moves them past then/else themselves).
+  (def (lower-if env ir)
+    (let* ([test (lower-expr env (typed-ir-if-test ir))]
+           [n (number->string (fresh-label-index! env))]
+           [then-label (string-append "then" n)]
+           [else-label (string-append "else" n)]
+           [join-label (string-append "join" n)]
+           [type (typed-ir-if-type ir)])
+      (finish-block! env
+        (string-append
+          "br i1 " (llvm-value-text test)
+          ", label %" then-label ", label %" else-label))
+      (start-block! env then-label)
+      ;; let* (not let): the branch body must lower before its end label is
+      ;; read, and Chez may evaluate plain let inits right-to-left.
+      (let* ([then-value (lower-expr env (typed-ir-if-then ir))]
+             [then-end (llvm-env-current-label env)])
+        (finish-block! env (string-append "br label %" join-label))
+        (start-block! env else-label)
+        (let* ([else-value (lower-expr env (typed-ir-if-else ir))]
+               [else-end (llvm-env-current-label env)])
+          (finish-block! env (string-append "br label %" join-label))
+          (start-block! env join-label)
+          (if (eq? type 'Unit)
+            unit-value
+            (let ([result (fresh-value! env)]
+                  [ty (llvm-type type)])
+              (emit-instr! env
+                (string-append
+                  result " = phi " ty
+                  " [ " (llvm-value-text then-value) ", %" then-end " ]"
+                  ", [ " (llvm-value-text else-value) ", %" else-end " ]"))
+              (make-llvm-value ty result)))))))
+
+  ;; --- call lowering -----------------------------------------------------------
+
+  ;; Lower an operand and require it to have the expected LLVM type. Mixed
+  ;; integer/float operands would be malformed IR, so they are rejected here
+  ;; instead of reaching the verifier.
+  (def (lower-operand env ir expected-ty context)
+    (let ([value (lower-expr env ir)])
+      (unless (string=? (llvm-value-type value) expected-ty)
+        (error 'typed-llvmir "operand type mismatch in LLVM lowering"
+          (list context 'expected expected-ty 'got (llvm-value-type value))))
+      value))
+
+  (def (arith-instruction operator type)
+    (case type
+      [(Float)
+       (case operator
+         [(+) "fadd"] [(-) "fsub"] [(*) "fmul"] [(/) "fdiv"]
+         [else (error 'typed-llvmir "unknown arithmetic operator" operator)])]
+      [(Nat)
+       (case operator
+         [(+) "add"] [(-) "sub"] [(*) "mul"] [(/) "udiv"]
+         [else (error 'typed-llvmir "unknown arithmetic operator" operator)])]
+      [(Int)
+       (case operator
+         [(+) "add"] [(-) "sub"] [(*) "mul"] [(/) "sdiv"]
+         [else (error 'typed-llvmir "unknown arithmetic operator" operator)])]
+      [else
+       (error 'typed-llvmir "unsupported arithmetic type for LLVM lowering" type)]))
+
+  ;; N-ary arithmetic lowers as a left fold, like the Rust backend's binary
+  ;; chains. A single operand returns its value unchanged (also Rust parity).
+  (def (lower-arith env operator args type)
+    (when (null? args)
+      (error 'typed-llvmir "arithmetic primitive needs at least one operand"
+        operator))
+    (let* ([instr (arith-instruction operator type)]
+           [ty (llvm-type type)])
+      (let loop ([acc (lower-operand env (car args) ty operator)]
+                 [rest (cdr args)])
+        (if (null? rest)
+          acc
+          (let ([rhs (lower-operand env (car rest) ty operator)]
+                [result (fresh-value! env)])
+            (emit-instr! env
+              (string-append
+                result " = " instr " " ty " "
+                (llvm-value-text acc) ", " (llvm-value-text rhs)))
+            (loop (make-llvm-value ty result) (cdr rest)))))))
+
+  ;; Signedness lives in operations, not types: Nat compares unsigned, Int
+  ;; signed, Float ordered (NaN-aware unordered forms are not exposed).
+  (def (cmp-instruction operator operand-type)
+    (case operand-type
+      [(Float)
+       (string-append "fcmp "
+         (case operator
+           [(=) "oeq"] [(<) "olt"] [(<=) "ole"] [(>) "ogt"] [(>=) "oge"]
+           [else (error 'typed-llvmir "unknown comparison operator" operator)]))]
+      [(Nat)
+       (string-append "icmp "
+         (case operator
+           [(=) "eq"] [(<) "ult"] [(<=) "ule"] [(>) "ugt"] [(>=) "uge"]
+           [else (error 'typed-llvmir "unknown comparison operator" operator)]))]
+      [(Int)
+       (string-append "icmp "
+         (case operator
+           [(=) "eq"] [(<) "slt"] [(<=) "sle"] [(>) "sgt"] [(>=) "sge"]
+           [else (error 'typed-llvmir "unknown comparison operator" operator)]))]
+      [else
+       (error 'typed-llvmir "unsupported comparison operand type" operand-type)]))
+
+  (def (merge-cmp-operand-type t1 t2)
+    (cond
+      [(or (eq? t1 'Float) (eq? t2 'Float)) 'Float]
+      [(or (eq? t1 'Int) (eq? t2 'Int)) 'Int]
+      [else 'Nat]))
+
+  (def (lower-cmp env operator args)
+    (unless (= (length args) 2)
+      (error 'typed-llvmir "comparison expects exactly two operands" operator))
+    (let* ([operand-type
+            (merge-cmp-operand-type
+              (typed-ir-node-type (car args))
+              (typed-ir-node-type (cadr args)))]
+           [ty (llvm-type operand-type)]
+           [lhs (lower-operand env (car args) ty operator)]
+           [rhs (lower-operand env (cadr args) ty operator)]
+           [result (fresh-value! env)])
+      (emit-instr! env
+        (string-append
+          result " = " (cmp-instruction operator operand-type) " " ty " "
+          (llvm-value-text lhs) ", " (llvm-value-text rhs)))
+      (make-llvm-value "i1" result)))
+
+  ;; (equal? a b) on scalars only; the checker records the operand type.
+  (def (lower-eq env ir args)
+    (unless (= (length args) 2)
+      (error 'typed-llvmir "equal? expects exactly two operands" args))
+    (let ([entry (assq 'operand-type (typed-ir-call-info ir))])
+      (unless (and entry (scalar-value-type? (cdr entry)))
+        (error 'typed-llvmir "equal? is only supported on scalar types"
+          (and entry (cdr entry))))
+      (let* ([operand-type (cdr entry)]
+             [ty (llvm-type operand-type)]
+             [instr (if (eq? operand-type 'Float) "fcmp oeq" "icmp eq")]
+             [lhs (lower-operand env (car args) ty 'equal?)]
+             [rhs (lower-operand env (cadr args) ty 'equal?)]
+             [result (fresh-value! env)])
+        (emit-instr! env
+          (string-append
+            result " = " instr " " ty " "
+            (llvm-value-text lhs) ", " (llvm-value-text rhs)))
+        (make-llvm-value "i1" result))))
+
+  ;; Primitive boolean operations on already-evaluated i1 values. The typed
+  ;; core IR carries no short-circuit promise for scalar and/or, so they
+  ;; lower to plain `and`/`or` instructions (operands are effect-free here).
+  (def (lower-bool env operator args)
+    (case operator
+      [(not)
+       (unless (= (length args) 1)
+         (error 'typed-llvmir "not expects one operand" args))
+       (let ([value (lower-operand env (car args) "i1" 'not)]
+             [result (fresh-value! env)])
+         (emit-instr! env
+           (string-append result " = xor i1 " (llvm-value-text value) ", true"))
+         (make-llvm-value "i1" result))]
+      [(and or)
+       (let ([instr (symbol->string operator)])
+         (cond
+           [(null? args)
+            (make-llvm-value "i1" (if (eq? operator 'and) "true" "false"))]
+           [else
+            (let loop ([acc (lower-operand env (car args) "i1" operator)]
+                       [rest (cdr args)])
+              (if (null? rest)
+                acc
+                (let ([rhs (lower-operand env (car rest) "i1" operator)]
+                      [result (fresh-value! env)])
+                  (emit-instr! env
+                    (string-append
+                      result " = " instr " i1 "
+                      (llvm-value-text acc) ", " (llvm-value-text rhs)))
+                  (loop (make-llvm-value "i1" result) (cdr rest)))))]))]
+      [else (error 'typed-llvmir "unknown bool operator" operator)]))
+
+  ;; Direct same-module calls. Signatures were resolved before lowering any
+  ;; body, so argument types come from the callee declaration.
+  (def (lower-function-call env name args)
+    (let ([sig (lookup-fn env name)])
+      (unless sig
+        (error 'typed-llvmir "unknown function in LLVM lowering" name))
+      (let ([param-types (llvm-fn-sig-param-types sig)])
+        (unless (= (length param-types) (length args))
+          (error 'typed-llvmir "function call arity mismatch in LLVM lowering"
+            (list name (length param-types) (length args))))
+        (let loop ([rest args] [rest-types param-types] [out '()])
+          (cond
+            [(null? rest)
+             (let* ([arg-values (reverse out)]
+                    [args-text
+                     (join-strings
+                       (map (lambda (value)
+                              (string-append
+                                (llvm-value-type value)
+                                " "
+                                (llvm-value-text value)))
+                            arg-values)
+                       ", ")]
+                    [return-type (llvm-fn-sig-return-type sig)])
+               (if (eq? return-type 'Unit)
+                 (begin
+                   (emit-instr! env
+                     (string-append
+                       "call void @" (llvm-fn-sig-symbol sig)
+                       "(" args-text ")"))
+                   unit-value)
+                 (let ([result (fresh-value! env)]
+                       [ret-ty (llvm-type return-type)])
+                   (emit-instr! env
+                     (string-append
+                       result " = call " ret-ty " @" (llvm-fn-sig-symbol sig)
+                       "(" args-text ")"))
+                   (make-llvm-value ret-ty result))))]
+            [else
+             (loop (cdr rest)
+                   (cdr rest-types)
+                   (cons (lower-operand env (car rest)
+                           (llvm-param-type (car rest-types))
+                           name)
+                         out))])))))
+
+  (def (lower-call env ir)
+    (let ([kind (typed-ir-call-kind ir)]
+          [operator (typed-ir-call-operator ir)]
+          [args (typed-ir-call-args ir)]
+          [type (typed-ir-call-type ir)])
+      (case kind
+        [(prim-arith) (lower-arith env operator args type)]
+        [(prim-cmp) (lower-cmp env operator args)]
+        [(prim-eq) (lower-eq env ir args)]
+        [(prim-bool) (lower-bool env operator args)]
+        [(function) (lower-function-call env operator args)]
+        [else
+         (error 'typed-llvmir "unsupported call kind for LLVM lowering" kind)])))
+
   ;; --- function lowering ---------------------------------------------------------
 
   (def (param-llvm-name index)
diff --git a/tests/fixtures/typed/llvmir-basic.ss b/tests/fixtures/typed/llvmir-basic.ss
index 4dd737c..97ba900 100644
--- a/tests/fixtures/typed/llvmir-basic.ss
+++ b/tests/fixtures/typed/llvmir-basic.ss
@@ -1,5 +1,5 @@
 (typed-library (sample typed llvmir-basic)
-  (export answer truthy echo)
+  (export answer truthy echo add-one double-add positive? div-nat div-int)
 
   (def (answer) : Nat
     42)
@@ -8,4 +8,20 @@
     #t)
 
   (def (echo (x : Nat)) : Nat
-    x))
+    x)
+
+  (def (add-one (x : Nat)) : Nat
+    (+ x 1))
+
+  (def (double-add (x : Nat)) : Nat
+    (let ((y (+ x x)))
+      (+ y 1)))
+
+  (def (positive? (x : Nat)) : Bool
+    (> x 0))
+
+  (def (div-nat (a : Nat) (b : Nat)) : Nat
+    (/ a b))
+
+  (def (div-int (a : Int) (b : Int)) : Int
+    (/ a b)))
diff --git a/tests/fixtures/typed/llvmir-call.ss b/tests/fixtures/typed/llvmir-call.ss
new file mode 100644
index 0000000..207ac10
--- /dev/null
+++ b/tests/fixtures/typed/llvmir-call.ss
@@ -0,0 +1,16 @@
+(typed-library (sample typed llvmir-call)
+  (export add2 scale compute lever)
+
+  (def (add2 (a : Nat) (b : Nat)) : Nat
+    (+ a b))
+
+  (def (scale (x : Nat)) : Nat
+    (* x 3))
+
+  (def (compute (x : Nat) (y : Nat)) : Nat
+    (scale (add2 x y)))
+
+  (def (lever (flag : Bool) (x : Nat)) : Nat
+    (if flag
+      (scale x)
+      (add2 x 1))))
diff --git a/tests/fixtures/typed/llvmir-if.ss b/tests/fixtures/typed/llvmir-if.ss
new file mode 100644
index 0000000..b85dd55
--- /dev/null
+++ b/tests/fixtures/typed/llvmir-if.ss
@@ -0,0 +1,20 @@
+(typed-library (sample typed llvmir-if)
+  (export choose max2 sign toggled)
+
+  (def (choose (flag : Bool)) : Nat
+    (if flag 1 0))
+
+  (def (max2 (a : Nat) (b : Nat)) : Nat
+    (if (> a b) a b))
+
+  ;; branches must agree on one numeric type, so Int 1 and Int 0 are
+  ;; computed from n rather than written as (Nat) literals
+  (def (sign (n : Int)) : Int
+    (if (< n 0)
+      -1
+      (if (> n 0)
+        (/ n n)
+        (- n n))))
+
+  (def (toggled (flag : Bool)) : Bool
+    (if (not flag) #t #f)))
diff --git a/tests/test-typed-llvmir.ss b/tests/test-typed-llvmir.ss
index 1e6b906..6dfd88d 100644
--- a/tests/test-typed-llvmir.ss
+++ b/tests/test-typed-llvmir.ss
@@ -86,6 +86,64 @@
      (def (echo (x : Nat)) : Nat
        x)))
 
+(define scalar-form
+  '(typed-library (sample typed llvmir-scalar)
+     (export add-one add3 double-add positive? div-nat div-int negate
+             both either same-nat half-up)
+     (def (add-one (x : Nat)) : Nat
+       (+ x 1))
+     (def (add3 (a : Nat) (b : Nat) (c : Nat)) : Nat
+       (+ a b c))
+     (def (double-add (x : Nat)) : Nat
+       (let ((y (+ x x)))
+         (+ y 1)))
+     (def (positive? (x : Nat)) : Bool
+       (> x 0))
+     (def (div-nat (a : Nat) (b : Nat)) : Nat
+       (/ a b))
+     (def (div-int (a : Int) (b : Int)) : Int
+       (/ a b))
+     (def (negate (flag : Bool)) : Bool
+       (not flag))
+     (def (both (a : Bool) (b : Bool)) : Bool
+       (and a b))
+     (def (either (a : Bool) (b : Bool)) : Bool
+       (or a b))
+     (def (same-nat (a : Nat) (b : Nat)) : Bool
+       (equal? a b))
+     (def (half-up (n : Int)) : Int
+       (+ (/ n 2) 1))))
+
+(define if-form
+  '(typed-library (sample typed llvmir-if)
+     (export choose max2 sign)
+     (def (choose (flag : Bool)) : Nat
+       (if flag 1 0))
+     (def (max2 (a : Nat) (b : Nat)) : Nat
+       (if (> a b) a b))
+     ;; branches must agree on one numeric type, so Int 1 and Int 0 are
+     ;; computed from n rather than written as (Nat) literals
+     (def (sign (n : Int)) : Int
+       (if (< n 0)
+         -1
+         (if (> n 0)
+           (/ n n)
+           (- n n))))))
+
+(define call-form
+  '(typed-library (sample typed llvmir-call)
+     (export add2 scale compute lever)
+     (def (add2 (a : Nat) (b : Nat)) : Nat
+       (+ a b))
+     (def (scale (x : Nat)) : Nat
+       (* x 3))
+     (def (compute (x : Nat) (y : Nat)) : Nat
+       (scale (add2 x y)))
+     (def (lever (flag : Bool) (x : Nat)) : Nat
+       (if flag
+         (scale x)
+         (add2 x 1)))))
+
 (define basic-ll (typed-library-form->llvmir-string basic-form))
 
 (test "module header names the backend"
@@ -128,12 +186,134 @@
   (substring? basic-ll "entry:")
   #t)
 
+;; --- scalar arithmetic, comparisons, bool, let --------------------------------------
+
+(define scalar-ll (typed-library-form->llvmir-string scalar-form))
+
+(test "nat addition lowers to add"
+  (substring? scalar-ll "%v0 = add i64 %a0, 1")
+  #t)
+
+(test "n-ary addition chains left to right"
+  (and (substring? scalar-ll "%v0 = add i64 %a0, %a1")
+       (substring? scalar-ll "%v1 = add i64 %v0, %a2"))
+  #t)
+
+(test "let binds initializer SSA value"
+  (and (substring? scalar-ll "%v0 = add i64 %a0, %a0")
+       (substring? scalar-ll "%v1 = add i64 %v0, 1"))
+  #t)
+
+(test "nat comparison is unsigned"
+  (substring? scalar-ll "icmp ugt i64 %a0, 0")
+  #t)
+
+(test "nat division is unsigned"
+  (substring? scalar-ll "udiv i64 %a0, %a1")
+  #t)
+
+(test "int division is signed"
+  (substring? scalar-ll "sdiv i64 %a0, %a1")
+  #t)
+
+(test "not lowers to xor true"
+  (substring? scalar-ll "%v0 = xor i1 %a0, true")
+  #t)
+
+(test "and lowers to i1 and"
+  (substring? scalar-ll "%v0 = and i1 %a0, %a1")
+  #t)
+
+(test "or lowers to i1 or"
+  (substring? scalar-ll "%v0 = or i1 %a0, %a1")
+  #t)
+
+(test "equal? on nats lowers to icmp eq"
+  (substring? scalar-ll "%v0 = icmp eq i64 %a0, %a1")
+  #t)
+
+(test "nested arithmetic threads SSA values"
+  (and (substring? scalar-ll "%v0 = sdiv i64 %a0, 2")
+       (substring? scalar-ll "%v1 = add i64 %v0, 1"))
+  #t)
+
+;; --- if lowering --------------------------------------------------------------------
+
+(define if-ll (typed-library-form->llvmir-string if-form))
+
+(test "if branches on the test value"
+  (substring? if-ll "br i1 %a0, label %then0, label %else0")
+  #t)
+
+(test "if branch blocks jump to the join"
+  (and (substring? if-ll "then0:")
+       (substring? if-ll "else0:")
+       (substring? if-ll "join0:")
+       (substring? if-ll "br label %join0"))
+  #t)
+
+(test "if joins with a phi over branch values"
+  (substring? if-ll "phi i64 [ 1, %then0 ], [ 0, %else0 ]")
+  #t)
+
+(test "if test instructions precede the branch"
+  (substring? if-ll "%v0 = icmp ugt i64 %a0, %a1")
+  #t)
+
+(test "int comparison is signed"
+  (substring? if-ll "icmp slt i64 %a0, 0")
+  #t)
+
+(test "nested if keeps phi predecessors at branch end blocks"
+  ;; sign: outer else contains the inner if, so the outer phi's else
+  ;; predecessor is the inner join block, not else0.
+  (substring? if-ll "phi i64 [ -1, %then0 ], [ %v4, %join1 ]")
+  #t)
+
+(test "nested if inner phi uses inner labels"
+  (substring? if-ll "phi i64 [ %v2, %then1 ], [ %v3, %else1 ]")
+  #t)
+
+;; --- direct calls -------------------------------------------------------------------
+
+(define call-ll (typed-library-form->llvmir-string call-form))
+
+(test "direct call passes typed arguments"
+  (substring? call-ll
+    "%v0 = call i64 @jt_llvm_sample_typed_llvmir_call__add2(i64 %a0, i64 %a1)")
+  #t)
+
+(test "nested call threads the result"
+  (substring? call-ll
+    "%v1 = call i64 @jt_llvm_sample_typed_llvmir_call__scale(i64 %v0)")
+  #t)
+
+(test "calls inside if branches stay in their blocks"
+  (and (substring? call-ll
+         "%v0 = call i64 @jt_llvm_sample_typed_llvmir_call__scale(i64 %a1)")
+       (substring? call-ll
+         "%v1 = call i64 @jt_llvm_sample_typed_llvmir_call__add2(i64 %a1, i64 1)")
+       (substring? call-ll "phi i64 [ %v0, %then0 ], [ %v1, %else0 ]"))
+  #t)
+
 ;; --- determinism ------------------------------------------------------------------
 
 (test "emission is deterministic across runs"
   (string=? basic-ll (typed-library-form->llvmir-string basic-form))
   #t)
 
+(test "scalar emission is deterministic across runs"
+  (string=? scalar-ll (typed-library-form->llvmir-string scalar-form))
+  #t)
+
+(test "if emission is deterministic across runs"
+  (string=? if-ll (typed-library-form->llvmir-string if-form))
+  #t)
+
+(test "call emission is deterministic across runs"
+  (string=? call-ll (typed-library-form->llvmir-string call-form))
+  #t)
+
 (test "forms->module emission is deterministic"
   (string=? (typed-library-forms->llvmir-module (list basic-form))
             (typed-library-forms->llvmir-module (list basic-form)))