typed/llvmir: Option lowering (P6)

ober

ba64305a8fe55da9563db464b455a2276b21c711

diff --git a/Makefile b/Makefile
index 9a0ec1a..6dbd6b1 100644
--- a/Makefile
+++ b/Makefile
@@ -40,7 +40,7 @@ PURE_AUDIT_ARGS ?= --summary --discover $(PURE_AUDIT_ROOT)
 TYPED_SOURCES ?= tests/fixtures/typed/valid-split-tree.ss
 TYPED_RUST_SOURCES ?= $(TYPED_SOURCES)
 TYPED_RUST_DIR ?= build/typed/rust
-TYPED_LLVMIR_SOURCES ?= tests/fixtures/typed/llvmir-basic.ss tests/fixtures/typed/llvmir-if.ss tests/fixtures/typed/llvmir-call.ss tests/fixtures/typed/llvmir-float.ss tests/fixtures/typed/llvmir-bitwise.ss tests/fixtures/typed/llvmir-for-fold.ss tests/fixtures/typed/llvmir-bytes.ss tests/fixtures/typed/llvmir-record.ss tests/fixtures/typed/llvmir-smoke.ss
+TYPED_LLVMIR_SOURCES ?= tests/fixtures/typed/llvmir-basic.ss tests/fixtures/typed/llvmir-if.ss tests/fixtures/typed/llvmir-call.ss tests/fixtures/typed/llvmir-float.ss tests/fixtures/typed/llvmir-bitwise.ss tests/fixtures/typed/llvmir-for-fold.ss tests/fixtures/typed/llvmir-bytes.ss tests/fixtures/typed/llvmir-record.ss tests/fixtures/typed/llvmir-option.ss tests/fixtures/typed/llvmir-smoke.ss
 TYPED_LLVMIR_DIR ?= build/typed/llvmir
 TYPED_LLVMIR_SMOKE_MODULE ?= sample_typed_llvmir_smoke
 TYPED_LLVMIR_SMOKE_EXPECT ?= 42
diff --git a/lib/jerboa/typed/llvmir.ss b/lib/jerboa/typed/llvmir.ss
index 21bb4fb..be38ea2 100644
--- a/lib/jerboa/typed/llvmir.ss
+++ b/lib/jerboa/typed/llvmir.ss
@@ -114,6 +114,8 @@
   ;; typed-record so type/ctor/accessor lowering can resolve fields and order.
   ;; Set during emission; empty outside it.
   (def *llvm-record-env* (make-parameter '()))
+  ;; Variant name -> typed-variant, for tag/field layout during emission.
+  (def *llvm-variant-env* (make-parameter '()))
 
   (def (lookup-record name)
     (let ([entry (assq name (*llvm-record-env*))])
@@ -122,6 +124,28 @@
   (def (record-type? type)
     (and (symbol? type) (lookup-record type) #t))
 
+  (def (lookup-variant name)
+    (let ([entry (assq name (*llvm-variant-env*))])
+      (and entry (cdr entry))))
+
+  (def (variant-type? type)
+    (and (symbol? type) (lookup-variant type) #t))
+
+  ;; (Option T): a value type carrying an optional payload.
+  (def (option-type? type)
+    (and (pair? type) (eq? (car type) 'Option) (= (length type) 2)))
+
+  ;; Variants lower to a tagged boxed value { i32 tag, ptr payload }: the tag
+  ;; is the case's index, the pointer a heap struct of that case's fields (null
+  ;; for a fieldless case). Boxing handles cases with different field sets
+  ;; uniformly and supports recursive variants. The box is never freed.
+  (def variant-llvm-type "{ i32, ptr }")
+
+  ;; (Option T) lowers to { i1 tag, <T> payload }: tag 1 = Some (payload valid),
+  ;; tag 0 = None (payload undef). By value, like a record.
+  (def (option-llvm-type type)
+    (string-append "{ i1, " (llvm-type (cadr type)) " }"))
+
   ;; Structural struct type: { <field0-llvm-type>, <field1-llvm-type>, ... }.
   ;; Structural (not a named %T) so no type definitions need emitting; the
   ;; checker already enforced field correctness. Recursive records are not
@@ -138,6 +162,8 @@
   (def (llvm-type type)
     (cond
       [(record-type? type) (record-llvm-type (lookup-record type))]
+      [(variant-type? type) variant-llvm-type]
+      [(option-type? type) (option-llvm-type type)]
       [else
        (case type
          [(Bool) "i1"]
@@ -151,9 +177,10 @@
     (and (symbol? type) (memq type '(Bool Nat Int Float)) #t))
 
   ;; A type that can be a function parameter or return value: scalars and
-  ;; buffers directly, plus record types (resolved structurally).
+  ;; buffers directly, plus record, variant, and Option types.
   (def (value-type? type)
-    (or (scalar-value-type? type) (buffer-type? type) (record-type? type)))
+    (or (scalar-value-type? type) (buffer-type? type) (record-type? type)
+        (variant-type? type) (option-type? type)))
 
   (def (llvm-return-type type)
     (if (or (eq? type 'Unit) (value-type? type))
@@ -648,6 +675,36 @@
                result))]
           [else (loop (cdr fields) (+ idx 1))]))))
 
+  ;; --- Option primitives -------------------------------------------------------
+
+  ;; (option-some x) : (Option T) -> { i1 true, T x }. The payload type comes
+  ;; from the call's own (Option T) type so Some and None agree on the struct.
+  (def (lower-option-some env ir args)
+    (unless (= (length args) 1)
+      (error 'typed-llvmir "option-some expects one operand" args))
+    (let* ([opt-type (typed-ir-call-type ir)]
+           [payload-ty (llvm-type (cadr opt-type))]
+           [struct-ty (option-llvm-type opt-type)]
+           [payload (lower-operand env (car args) payload-ty 'option-some)]
+           [t0 (fresh-value! env)]
+           [result (fresh-value! env)])
+      (emit-instr! env
+        (string-append t0 " = insertvalue " struct-ty " undef, i1 true, 0"))
+      (emit-instr! env
+        (string-append
+          result " = insertvalue " struct-ty " " t0 ", "
+          payload-ty " " (llvm-value-text payload) ", 1"))
+      (make-llvm-value struct-ty result)))
+
+  ;; (option-none) : (Option T) -> { i1 false, T undef }.
+  (def (lower-option-none env ir)
+    (let* ([opt-type (typed-ir-call-type ir)]
+           [struct-ty (option-llvm-type opt-type)]
+           [result (fresh-value! env)])
+      (emit-instr! env
+        (string-append result " = insertvalue " struct-ty " undef, i1 false, 0"))
+      (make-llvm-value struct-ty result)))
+
   ;; --- call lowering -----------------------------------------------------------
 
   ;; Lower an operand and require it to have the expected LLVM type. Mixed
@@ -933,6 +990,8 @@
         [(record-ctor) (lower-record-ctor env ir args)]
         [(record-accessor) (lower-record-accessor env ir args)]
         [(record-pred) (make-llvm-value "i1" "true")]
+        [(option-some) (lower-option-some env ir args)]
+        [(option-none) (lower-option-none env ir)]
         [(function) (lower-function-call env operator args)]
         [else
          (error 'typed-llvmir "unsupported call kind for LLVM lowering" kind)])))
@@ -1023,10 +1082,11 @@
                                  (typed-def-return-type def)
                                  (map typed-param-type (typed-def-params def))))
                          out)))]
-          [(typed-record? (car rest)) (loop (cdr rest) out)]
+          [(or (typed-record? (car rest)) (typed-variant? (car rest)))
+           (loop (cdr rest) out)]
           [else
            (error 'typed-llvmir
-             "only function and record declarations are supported by the LLVM backend"
+             "only function, record, and variant declarations are supported by the LLVM backend"
              (car rest))]))))
 
   ;; alist of record-name -> typed-record for the modules in scope.
@@ -1043,6 +1103,20 @@
                      (cons (cons (typed-record-name (car rest)) (car rest)) out))]
              [else (inner (cdr rest) out)]))])))
 
+  ;; alist of variant-name -> typed-variant for the modules in scope.
+  (def (modules-variant-env modules)
+    (let loop ([modules modules] [out '()])
+      (cond
+        [(null? modules) (reverse out)]
+        [else
+         (let inner ([rest (typed-module-declarations (car modules))] [out out])
+           (cond
+             [(null? rest) (loop (cdr modules) out)]
+             [(typed-variant? (car rest))
+              (inner (cdr rest)
+                     (cons (cons (typed-variant-name (car rest)) (car rest)) out))]
+             [else (inner (cdr rest) out)]))])))
+
   (def (elaborate-module-or-error module who)
     (let-values ([(errors defs)
                   (check-and-elaborate-typed-module module)])
@@ -1149,7 +1223,8 @@
     (newline port))
 
   (def (typed-module->llvmir-string module)
-    (parameterize ([*llvm-record-env* (modules-record-env (list module))])
+    (parameterize ([*llvm-record-env* (modules-record-env (list module))]
+                   [*llvm-variant-env* (modules-variant-env (list module))])
       (let* ([lowered (module-functions module)]
              [functions (car lowered)]
              [ctx (cdr lowered)])
@@ -1234,7 +1309,8 @@
                       'typed-modules->llvmir-string)]
            [fns (modules-fn-env modules)]
            [ctx (make-empty-mod-ctx)])
-      (parameterize ([*llvm-record-env* (modules-record-env modules)])
+      (parameterize ([*llvm-record-env* (modules-record-env modules)]
+                     [*llvm-variant-env* (modules-variant-env modules)])
         (let ([functions
                (let loop ([rest modules] [envs ir-envs] [acc '()])
                  (cond
diff --git a/tests/fixtures/typed/llvmir-option.ss b/tests/fixtures/typed/llvmir-option.ss
new file mode 100644
index 0000000..6111ce9
--- /dev/null
+++ b/tests/fixtures/typed/llvmir-option.ss
@@ -0,0 +1,22 @@
+(typed-library (sample typed llvmir-option)
+  (export some-n none-n pass-opt safe-div some-bytes)
+
+  (def (some-n (x : Nat)) : (Option Nat)
+    (option-some x))
+
+  (def (none-n) : (Option Nat)
+    (option-none Nat))
+
+  ;; Option is a real value type: pass one through
+  (def (pass-opt (o : (Option Nat))) : (Option Nat)
+    o)
+
+  ;; build Some/None by a condition
+  (def (safe-div (a : Nat) (b : Nat)) : (Option Nat)
+    (if (> b 0)
+      (option-some (/ a b))
+      (option-none Nat)))
+
+  ;; Option over a buffer payload nests { i1, { ptr, i64 } }
+  (def (some-bytes (b : Bytes)) : (Option Bytes)
+    (option-some b)))
diff --git a/tests/test-typed-llvmir.ss b/tests/test-typed-llvmir.ss
index 71432ff..050308b 100644
--- a/tests/test-typed-llvmir.ss
+++ b/tests/test-typed-llvmir.ss
@@ -594,6 +594,58 @@
   (string=? record-ll (typed-library-form->llvmir-string record-form))
   #t)
 
+;; --- Option lowering ----------------------------------------------------------------
+
+(define option-form
+  '(typed-library (sample typed llvmir-option)
+     (export some-n none-n pass-opt safe-div some-bytes)
+     (def (some-n (x : Nat)) : (Option Nat)
+       (option-some x))
+     (def (none-n) : (Option Nat)
+       (option-none Nat))
+     (def (pass-opt (o : (Option Nat))) : (Option Nat)
+       o)
+     (def (safe-div (a : Nat) (b : Nat)) : (Option Nat)
+       (if (> b 0)
+         (option-some (/ a b))
+         (option-none Nat)))
+     (def (some-bytes (b : Bytes)) : (Option Bytes)
+       (option-some b))))
+
+(define option-ll (typed-library-form->llvmir-string option-form))
+
+(test "Option lowers to a tagged { i1, T } struct"
+  (substring? option-ll
+    "define { i1, i64 } @jt_llvm_sample_typed_llvmir_option__some_n(i64 %a0)")
+  #t)
+
+(test "option-some sets the tag to true and stores the payload"
+  (and (substring? option-ll "insertvalue { i1, i64 } undef, i1 true, 0")
+       (substring? option-ll "insertvalue { i1, i64 } %v0, i64 %a0, 1"))
+  #t)
+
+(test "option-none sets the tag to false with undef payload"
+  (substring? option-ll "insertvalue { i1, i64 } undef, i1 false, 0")
+  #t)
+
+(test "Option is a value type usable as a parameter"
+  (substring? option-ll
+    "@jt_llvm_sample_typed_llvmir_option__pass_opt({ i1, i64 } %a0)")
+  #t)
+
+(test "Option branches join with a struct phi"
+  (substring? option-ll "phi { i1, i64 } [ ")
+  #t)
+
+(test "Option over a buffer payload nests the fat pointer"
+  (substring? option-ll
+    "define { i1, { ptr, i64 } } @jt_llvm_sample_typed_llvmir_option__some_bytes")
+  #t)
+
+(test "Option emission is deterministic across runs"
+  (string=? option-ll (typed-library-form->llvmir-string option-form))
+  #t)
+
 ;; --- @main smoke wrapper --------------------------------------------------------
 
 (define smoke-form
@@ -714,13 +766,13 @@
     'accepted)
   'accepted)
 
-(test "option-returning defs are rejected"
+(test "crypto-primitive defs are rejected (stay on the RustCrypto path)"
   (guard (exn [#t 'rejected])
     (typed-library-form->llvmir-string
-      '(typed-library (sample typed llvmir-option)
-         (export maybe)
-         (def (maybe (x : Nat)) : (Option Nat)
-           (option-some x))))
+      '(typed-library (sample typed llvmir-crypto)
+         (export digest)
+         (def (digest (data : Bytes)) : Bytes
+           (sha256 data))))
     'accepted)
   'rejected)