typed-rust: lower for/fold over in-range to a native Rust for loop
ober
09ee436568a18ac65ac9f794436f077535e3bfe3
--- 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. --- 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 --- 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) --- 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)])) --- 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)] new file mode 100644 --- /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)))) --- 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) --- 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>")