typed-rust: lower for/fold over in-range to a native Rust for loop

ober

09ee436568a18ac65ac9f794436f077535e3bfe3

diff --git a/docs/jerboa-to-rust.md b/docs/jerboa-to-rust.md
index 4f2efad..683808a 100644
--- a/docs/jerboa-to-rust.md
+++ b/docs/jerboa-to-rust.md
@@ -92,6 +92,15 @@ Landed:
   is `(Bytes -> String)` and lowers to lossy decode
   `String::from_utf8_lossy(&(buf)).into_owned()` (matches the String input ABI,
   never panics). Both borrow rather than move their source.
+- `for/fold` over a single `in-range` clause with one accumulator lowers to a
+  native Rust `for` loop: `{ let mut acc = init; for i in (start)..(end) { acc =
+  body; } acc }`. The accumulator type is pinned by its initial value and the
+  body must reproduce that type; the index runs `[start, end)` (start defaults
+  to `0`). This is the idiomatic typed replacement for a counting named-let --
+  it avoids recursion-based inference (the checker has no type variables) and
+  the per-call clone storms tail recursion would incur. It is enough to write a
+  constant-time difference fold and a rolling checksum directly; see
+  `tests/fixtures/typed/rust-for-fold.ss`.
 - Rust `struct`/`enum` generation for Typed Jerboa records and variants,
   including `Clone`, `Debug`, and `PartialEq` derives.
 - Recursive variant field boxing and match rebinding for owned child values.
diff --git a/docs/typed-jerboa.md b/docs/typed-jerboa.md
index 35106c9..6c1cb8a 100644
--- a/docs/typed-jerboa.md
+++ b/docs/typed-jerboa.md
@@ -238,7 +238,13 @@ Current landing:
   (`String -> Bytes`) and `utf8->string` (`Bytes -> String`) cross the
   text/binary boundary, lowering to `as_bytes().to_vec()` and a lossy
   `from_utf8_lossy` decode, so a constant-time auth-token comparison ports too
-  (`tests/fixtures/typed/rust-utf8.ss`). `debug-string` checks one typed
+  (`tests/fixtures/typed/rust-utf8.ss`). `for/fold` with a single accumulator
+  over one `in-range` clause is checked by pinning the accumulator type from its
+  initializer and requiring the body to reproduce it; it lowers to a native Rust
+  `for` loop (`{ let mut acc = init; for i in (start)..(end) { acc = body; } acc
+  }`), the idiomatic typed stand-in for a counting named-let that needs neither
+  type variables nor clone-heavy tail recursion (`tests/fixtures/typed/rust-for-fold.ss`).
+  `debug-string` checks one typed
   operand and lowers to Rust `format!("{:?}", ...)`, using the Debug derives on
   generated records and variants. Imported calls and richer forms are reported
   as unsupported. Typed core IR is now produced by the checker and consumed by
diff --git a/lib/jerboa/typed/checker.ss b/lib/jerboa/typed/checker.ss
index fa7dc31..e0a80a6 100644
--- a/lib/jerboa/typed/checker.ss
+++ b/lib/jerboa/typed/checker.ss
@@ -1162,6 +1162,121 @@
                    (expr-source expr) 'prim-shift op arg-irs '()))
             (append errors operand-errors))))))
 
+  ;; Parse a single for/fold range clause: (var (in-range end)) or
+  ;; (var (in-range start end)). Returns (var-symbol start-expr-or-#f end-expr)
+  ;; or #f when the shape is not recognised.
+  (def (parse-range-clause clause)
+    (and (expr-list? clause)
+         (= (expr-length clause) 2)
+         (expr-symbol? (expr-car clause))
+         (let ([range (expr-cadr clause)])
+           (and (expr-pair? range)
+                (eq? (expr-head range) 'in-range)
+                (let ([rargs (cdr (expr->list range))])
+                  (cond
+                    [(= (length rargs) 1)
+                     (list (expr-value (expr-car clause)) #f (car rargs))]
+                    [(= (length rargs) 2)
+                     (list (expr-value (expr-car clause))
+                           (car rargs) (cadr rargs))]
+                    [else #f]))))))
+
+  (def (infer-for-fold args env type-names expr)
+    ;; (for/fold ([acc init]) ([i (in-range end)]) body ...).  Single
+    ;; accumulator over a single in-range index; lowers to a Rust `for` loop
+    ;; that reassigns the accumulator. The accumulator type is fixed by its
+    ;; init and the body must produce that same type (the fold invariant), so
+    ;; no recursion-result inference is needed -- unlike a named let.
+    (cond
+      [(< (length args) 3)
+       (values #f
+         (list (error-at expr 'bad-for-fold
+                 "for/fold expects accumulators, a range clause, and a body"
+                 expr)))]
+      [(not (and (expr-list? (car args)) (expr-list? (cadr args))))
+       (values #f
+         (list (error-at expr 'bad-for-fold
+                 "for/fold accumulators and clauses must each be a list"
+                 expr)))]
+      [else
+       (let ([accs (expr->list (car args))]
+             [clauses (expr->list (cadr args))]
+             [body (cddr args)])
+         (cond
+           [(not (= (length accs) 1))
+            (values #f
+              (list (error-at expr 'unsupported-for-fold
+                      "only a single for/fold accumulator is supported yet"
+                      expr)))]
+           [(not (= (length clauses) 1))
+            (values #f
+              (list (error-at expr 'unsupported-for-fold
+                      "only a single for/fold range clause is supported yet"
+                      expr)))]
+           [else
+            (infer-for-fold-1 (car accs) (car clauses) body
+                              env type-names expr)]))]))
+
+  (def (infer-for-fold-1 acc-binding clause body env type-names expr)
+    (let ([acc-shape (check-let-binding-shape acc-binding)]
+          [clause-info (parse-range-clause clause)])
+      (cond
+        [(not (null? acc-shape)) (values #f acc-shape)]
+        [(not clause-info)
+         (values #f
+           (list (error-at clause 'bad-for-fold
+                   "for/fold clause must be (var (in-range end)) or (var (in-range start end))"
+                   clause)))]
+        [else
+         (let ([acc-name (binding-name acc-binding)]
+               [acc-init-expr (expr-cadr acc-binding)]
+               [var-name (car clause-info)]
+               [start-expr (cadr clause-info)]
+               [end-expr (caddr clause-info)])
+           (let*-values
+             ([(acc-ir acc-errors)
+               (infer-expression acc-init-expr env type-names)]
+              [(start-ir start-errors)
+               (if start-expr
+                 (infer-expression start-expr env type-names)
+                 (values (make-typed-ir-lit 'Nat (expr-source expr) 0) '()))]
+              [(end-ir end-errors)
+               (infer-expression end-expr env type-names)])
+             (let* ([acc-type (ir-type acc-ir)]
+                    [start-type (ir-type start-ir)]
+                    [end-type (ir-type end-ir)]
+                    [range-errors
+                     (operand-type-errors 'numeric
+                       (list start-type end-type)
+                       (list clause clause) clause)]
+                    [var-type (merge-numeric-types (list start-type end-type))]
+                    [body-env
+                     (if acc-type
+                       (extend-env var-name var-type
+                         (extend-env acc-name acc-type env))
+                       env)])
+               (let-values ([(body-ir body-errors)
+                             (infer-body body body-env type-names
+                                         (expr-source expr))])
+                 (let* ([body-type (ir-type body-ir)]
+                        [invariant-errors
+                         (if (and acc-type body-type
+                                  (not (equal? acc-type body-type)))
+                           (list (error-at expr 'for-fold-body-type-mismatch
+                                   "for/fold body must have the accumulator type"
+                                   (list acc-type body-type)))
+                           '())]
+                        [ok? (and acc-ir start-ir end-ir body-ir acc-type
+                                  (null? acc-errors) (null? start-errors)
+                                  (null? end-errors) (null? body-errors)
+                                  (null? range-errors) (null? invariant-errors))])
+                   (values
+                     (and ok?
+                          (make-typed-ir-for-fold acc-type (expr-source expr)
+                            acc-name acc-ir var-name start-ir end-ir body-ir))
+                     (append acc-errors start-errors end-errors body-errors
+                             range-errors invariant-errors)))))))])))
+
   (def (bad-constructor-arity expr name expected args)
     (list (error-at expr 'bad-call-arity
             "typed constructor arity does not match"
@@ -1490,6 +1605,8 @@
               (infer-bitwise head args env type-names expr)]
              [(bitwise-arithmetic-shift-left bitwise-arithmetic-shift-right)
               (infer-shift head args env type-names expr)]
+             [(for/fold)
+              (infer-for-fold args env type-names expr)]
              [(option-some)
               (infer-option-some args env type-names expr)]
              [(option-none)
diff --git a/lib/jerboa/typed/core.ss b/lib/jerboa/typed/core.ss
index 539682d..f96e7c0 100644
--- a/lib/jerboa/typed/core.ss
+++ b/lib/jerboa/typed/core.ss
@@ -39,6 +39,14 @@
     typed-ir-if-type typed-ir-if-source
     typed-ir-if-test typed-ir-if-then typed-ir-if-else
 
+    typed-ir-for-fold?
+    make-typed-ir-for-fold
+    typed-ir-for-fold-type typed-ir-for-fold-source
+    typed-ir-for-fold-acc-name typed-ir-for-fold-acc-init
+    typed-ir-for-fold-var-name
+    typed-ir-for-fold-range-start typed-ir-for-fold-range-end
+    typed-ir-for-fold-body
+
     typed-ir-match?
     make-typed-ir-match
     typed-ir-match-type typed-ir-match-source
@@ -74,6 +82,10 @@
   (defstruct typed-ir-let (type source bindings body))
   (defstruct typed-ir-binding (name expr))
   (defstruct typed-ir-if (type source test then else))
+  ;; single-accumulator fold over an in-range index: the loop variable runs
+  ;; [range-start, range-end); body (of acc's type) becomes acc's next value.
+  (defstruct typed-ir-for-fold
+    (type source acc-name acc-init var-name range-start range-end body))
   (defstruct typed-ir-match
     (type source scrutinee scrutinee-type clauses default))
   (defstruct typed-ir-match-clause (case bindings field-types body))
@@ -119,6 +131,7 @@
         (typed-ir-begin? x)
         (typed-ir-let? x)
         (typed-ir-if? x)
+        (typed-ir-for-fold? x)
         (typed-ir-match? x)
         (typed-ir-call? x)))
 
@@ -129,6 +142,7 @@
       [(typed-ir-begin? node) (typed-ir-begin-type node)]
       [(typed-ir-let? node) (typed-ir-let-type node)]
       [(typed-ir-if? node) (typed-ir-if-type node)]
+      [(typed-ir-for-fold? node) (typed-ir-for-fold-type node)]
       [(typed-ir-match? node) (typed-ir-match-type node)]
       [(typed-ir-call? node) (typed-ir-call-type node)]
       [else (error 'typed-ir-node-type "not a typed IR node" node)]))
@@ -140,6 +154,7 @@
       [(typed-ir-begin? node) (typed-ir-begin-source node)]
       [(typed-ir-let? node) (typed-ir-let-source node)]
       [(typed-ir-if? node) (typed-ir-if-source node)]
+      [(typed-ir-for-fold? node) (typed-ir-for-fold-source node)]
       [(typed-ir-match? node) (typed-ir-match-source node)]
       [(typed-ir-call? node) (typed-ir-call-source node)]
       [else (error 'typed-ir-node-source "not a typed IR node" node)]))
diff --git a/lib/jerboa/typed/rust.ss b/lib/jerboa/typed/rust.ss
index bb14dbe..c535d15 100644
--- a/lib/jerboa/typed/rust.ss
+++ b/lib/jerboa/typed/rust.ss
@@ -1407,6 +1407,7 @@
       [(typed-ir-begin? ir) (emit-begin (typed-ir-begin-exprs ir))]
       [(typed-ir-let? ir) (emit-ir-let ir)]
       [(typed-ir-if? ir) (emit-ir-if ir)]
+      [(typed-ir-for-fold? ir) (emit-ir-for-fold ir)]
       [(typed-ir-match? ir) (emit-ir-match ir)]
       [(typed-ir-call? ir) (emit-ir-call ir)]
       [else (error 'typed-rust "unsupported IR node" ir)]))
@@ -1440,6 +1441,23 @@
       (emit-expression (typed-ir-if-else ir))
       " }"))
 
+  ;; (for/fold ([acc init]) ([i (in-range start end)]) body) lowers to a block
+  ;; that seeds a mutable accumulator, runs a Rust `for` over the half-open
+  ;; range, reassigns the accumulator from the body each step, and yields it.
+  (def (emit-ir-for-fold ir)
+    (let ([acc (rust-symbol-name (typed-ir-for-fold-acc-name ir))]
+          [var (rust-symbol-name (typed-ir-for-fold-var-name ir))])
+      (string-append
+        "{ "
+        (rust-inline-source-comment (typed-ir-node-source ir))
+        "let mut " acc " = "
+        (emit-expression (typed-ir-for-fold-acc-init ir)) "; "
+        "for " var " in ("
+        (emit-expression (typed-ir-for-fold-range-start ir)) ")..("
+        (emit-expression (typed-ir-for-fold-range-end ir)) ") { "
+        acc " = " (emit-expression (typed-ir-for-fold-body ir)) "; } "
+        acc " }")))
+
   (def (emit-ir-match-clause clause)
     (let* ([case-name (typed-ir-match-clause-case clause)]
            [bindings (typed-ir-match-clause-bindings clause)]
diff --git a/tests/fixtures/typed/rust-for-fold.ss b/tests/fixtures/typed/rust-for-fold.ss
new file mode 100644
index 0000000..55441e0
--- /dev/null
+++ b/tests/fixtures/typed/rust-for-fold.ss
@@ -0,0 +1,23 @@
+(typed-library (sample typed for-fold)
+  (export sum-to ct-diff ct-equal? checksum)
+
+  ;; sum of 0 .. n-1
+  (def (sum-to (n : Nat)) : Nat
+    (for/fold ((acc 0)) ((i (in-range n)))
+      (+ acc i)))
+
+  ;; constant-time difference accumulator over two buffers, written as a
+  ;; for/fold instead of a named-let -- the idiomatic typed loop form.
+  (def (ct-diff (a : Bytes) (b : Bytes) (n : Nat)) : Nat
+    (for/fold ((acc 0)) ((i (in-range n)))
+      (bitwise-ior acc (bitwise-xor (bytevector-u8-ref a i)
+                                    (bytevector-u8-ref b i)))))
+
+  ;; 0 iff equal, in constant time
+  (def (ct-equal? (a : Bytes) (b : Bytes) (n : Nat)) : Bool
+    (= (ct-diff a b n) 0))
+
+  ;; rolling additive checksum over a range, using (in-range start end)
+  (def (checksum (data : Bytes) (start : Nat) (end : Nat)) : Nat
+    (for/fold ((acc 0)) ((i (in-range start end)))
+      (bitwise-and (+ acc (bytevector-u8-ref data i)) 255))))
diff --git a/tests/test-typed-checker.ss b/tests/test-typed-checker.ss
index 06c3a3c..70e97f6 100644
--- a/tests/test-typed-checker.ss
+++ b/tests/test-typed-checker.ss
@@ -647,6 +647,42 @@
          (utf8->string s))))
   '(argument-type-mismatch))
 
+(test "for/fold over in-range returns the accumulator type"
+  (error-kinds
+    '(typed-library (body fold-ok)
+       (export f)
+       (def (f (n : Nat)) : Nat
+         (for/fold ((acc 0)) ((i (in-range n)))
+           (+ acc i)))))
+  '())
+
+(test "for/fold body type must match accumulator type"
+  (error-kinds
+    '(typed-library (body fold-mismatch)
+       (export f)
+       (def (f (s : String) (n : Nat)) : String
+         (for/fold ((acc s)) ((i (in-range n)))
+           i))))
+  '(for-fold-body-type-mismatch))
+
+(test "for/fold rejects multiple accumulators (not yet supported)"
+  (error-kinds
+    '(typed-library (body fold-multi)
+       (export f)
+       (def (f (n : Nat)) : Nat
+         (for/fold ((a 0) (b 0)) ((i (in-range n)))
+           (+ a b)))))
+  '(unsupported-for-fold))
+
+(test "for/fold rejects a clause that is not an in-range"
+  (error-kinds
+    '(typed-library (body fold-bad-clause)
+       (export f)
+       (def (f (n : Nat)) : Nat
+         (for/fold ((acc 0)) ((i n))
+           acc))))
+  '(bad-for-fold))
+
 (test "record constructor and accessor calls"
   (error-kinds
     '(typed-library (body record-ok)
diff --git a/tests/test-typed-rust.ss b/tests/test-typed-rust.ss
index 0610df4..c8d4415 100644
--- a/tests/test-typed-rust.ss
+++ b/tests/test-typed-rust.ss
@@ -199,6 +199,25 @@
 (define utf8-rust
   (typed-library-form->rust-string utf8-form))
 
+;; for/fold over in-range: the idiomatic typed loop, lowering to a Rust `for`
+;; that reassigns a mutable accumulator (no recursion, no clone storms).
+(define for-fold-form
+  '(typed-library (sample typed for-fold)
+     (export sum-to ct-diff window-sum)
+     (def (sum-to (n : Nat)) : Nat
+       (for/fold ((acc 0)) ((i (in-range n)))
+         (+ acc i)))
+     (def (ct-diff (a : Bytes) (b : Bytes) (n : Nat)) : Nat
+       (for/fold ((acc 0)) ((i (in-range n)))
+         (bitwise-ior acc (bitwise-xor (bytevector-u8-ref a i)
+                                       (bytevector-u8-ref b i)))))
+     (def (window-sum (data : Bytes) (start : Nat) (end : Nat)) : Nat
+       (for/fold ((acc 0)) ((i (in-range start end)))
+         (+ acc (bytevector-u8-ref data i))))))
+
+(define for-fold-rust
+  (typed-library-form->rust-string for-fold-form))
+
 (define return-form
   '(typed-library (sample typed return-values)
      (export greeting echo-text echo-bytes)
@@ -476,6 +495,16 @@
          "String::from_utf8_lossy(&((s).as_bytes().to_vec())).into_owned()"))
   #t)
 
+(test "rust lowers for/fold over in-range to a for loop"
+  (and (substring? for-fold-rust
+         "let mut acc = 0u64; for i in (0u64)..(n) { acc = (acc + i); } acc")
+       ;; constant-time fold body reassigns the accumulator each step
+       (substring? for-fold-rust
+         "acc = (acc | ((a[(i) as usize] as u64) ^ (b[(i) as usize] as u64)));")
+       ;; two-argument in-range lowers to a (start)..(end) loop
+       (substring? for-fold-rust "for i in (start)..(end)"))
+  #t)
+
 (test "rust emits handle registry for record and variant ABI"
   (and (substring? handle-rust "fn jt_store_handle<T: Any + Send>(value: T) -> u64")
        (substring? handle-rust "fn jt_clone_handle<T: Any + Clone>(id: u64) -> Option<T>")