typed-rust: lower bitwise/shift and bytevector-u8-ref (constant-time-equal)

ober

d64dded253fa4d9b09bd69712c172b7de77adf71

diff --git a/docs/jerboa-to-rust.md b/docs/jerboa-to-rust.md
index 912b5d9..d1c914f 100644
--- a/docs/jerboa-to-rust.md
+++ b/docs/jerboa-to-rust.md
@@ -72,6 +72,21 @@ Landed:
 - Safe Rust text for primitive values, `let`, `if`, arithmetic, numeric
   comparisons, boolean operators, same-type `equal?`, `debug-string`,
   `string-length`, `string-append`, and `bytevector-length`.
+- Bitwise and shift primitives: `bitwise-and`/`bitwise-ior`/`bitwise-xor`
+  lower to `&`/`|`/`^`, `bitwise-not` to `!`, and
+  `bitwise-arithmetic-shift-left`/`-right` to `<<`/`>>`. Operands and result
+  are the same merged numeric type (shifts follow the value being shifted).
+  Lowering is fixed-width (Nat -> u64), so `bitwise-not` and masking behave as
+  fixed-width complements -- the intended semantics for byte/word twiddling in
+  TLV parsing, hex codecs, and constant-time comparison.
+- `bytevector-u8-ref` is checked as `(Bytes Nat -> Nat)` and lowers to indexed
+  byte access `buf[(i) as usize] as u64`. With the bitwise primitives and
+  same-module recursion this is enough to express a constant-time `equal?`
+  (XOR-accumulate fold) entirely in Typed Jerboa; see
+  `tests/fixtures/typed/rust-ct-equal.ss`. Bytevector construction
+  (`make-bytevector`) and in-place mutation (`bytevector-u8-set!`,
+  `bytevector-copy!`) are not yet lowered -- they need the mutable-ownership
+  model that the current Clone-heavy lowering does not yet provide.
 - 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 d96f2bb..80c073c 100644
--- a/docs/typed-jerboa.md
+++ b/docs/typed-jerboa.md
@@ -222,13 +222,19 @@ Current landing:
   diagnostics expose them through `typed-check-error-source`.
 - This is still a front-end milestone. Function-body checking currently covers
   literals, variables, `begin`, simple `let`, `if`, arithmetic primitives,
-  numeric comparisons, same-type `equal?`, boolean primitives, calls to typed
-  functions defined in the same module, generated record/variant operations,
+  bitwise primitives (`bitwise-and`/`-ior`/`-xor` variadic, `bitwise-not`
+  unary) and the two-operand shift primitives, numeric comparisons, same-type
+  `equal?`, boolean primitives, calls to typed functions defined in the same
+  module, generated record/variant operations,
   and exhaustive `match` over same-module variants. The first builtin string primitive,
   `string-length`, is checked as `(String -> Nat)` and lowers to Rust
   `.len()`; `string-append` is checked as a two-argument `(String String ->
   String)` builtin and lowers to Rust `format!`; `bytevector-length` is checked
-  as `(Bytes -> Nat)` and lowers the same way. `debug-string` checks one typed
+  as `(Bytes -> Nat)` and lowers the same way. `bytevector-u8-ref` is checked as
+  `(Bytes Nat -> Nat)` and lowers to indexed access `buf[(i) as usize] as u64`;
+  combined with the bitwise primitives and same-module recursion, this is enough
+  to express a constant-time `equal?` (XOR-accumulate fold) entirely in Typed
+  Jerboa — see `tests/fixtures/typed/rust-ct-equal.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 93881c3..9589095 100644
--- a/lib/jerboa/typed/checker.ss
+++ b/lib/jerboa/typed/checker.ss
@@ -362,7 +362,10 @@
               'string-append '()))
       (cons 'bytevector-length
             (make-typed-call-sig (list 'Bytes) 'Nat '()
-              'bytevector-length '()))))
+              'bytevector-length '()))
+      (cons 'bytevector-u8-ref
+            (make-typed-call-sig (list 'Bytes 'Nat) 'Nat '()
+              'bytevector-u8-ref '()))))
 
   (def (field-types fields)
     (map typed-field-type fields))
@@ -1103,6 +1106,56 @@
                     'prim-bool op arg-irs '()))
              (append errors operand-errors))))]))
 
+  (def (infer-bitwise op args env type-names expr)
+    ;; bitwise-and/ior/xor are variadic (>= 1 operand, matching Chez where a
+    ;; single operand returns itself); bitwise-not is unary. All operands and
+    ;; the result are the same merged numeric type. Note the Rust lowering is
+    ;; fixed-width (Nat -> u64), so bitwise-not is a fixed-width complement,
+    ;; unlike Chez's unbounded integers -- the intended semantics for masking.
+    (cond
+      [(and (eq? op 'bitwise-not) (not (= (length args) 1)))
+       (values #f
+         (list (error-at expr 'bad-primitive-arity
+                 "bitwise-not needs exactly one operand"
+                 expr)))]
+      [(null? args)
+       (values #f
+         (list (error-at expr 'bad-primitive-arity
+                 "bitwise primitive needs at least one operand"
+                 expr)))]
+      [else
+       (let-values ([(arg-irs errors) (infer-args args env type-names)])
+         (let* ([types (ir-list-types arg-irs)]
+                [operand-errors (operand-type-errors 'numeric types args expr)]
+                [ok? (and (null? errors) (null? operand-errors)
+                          (all-irs-valid? arg-irs))])
+           (values
+             (and ok?
+                  (make-typed-ir-call (merge-numeric-types types)
+                    (expr-source expr) 'prim-bitwise op arg-irs '()))
+             (append errors operand-errors))))]))
+
+  (def (infer-shift op args env type-names expr)
+    ;; (bitwise-arithmetic-shift-left value amount) and the right variant take
+    ;; exactly two operands. Both must be numeric; the result follows the type
+    ;; of the value being shifted (the first operand). Rust accepts a shift
+    ;; amount of any integer width, so mixed Nat/Int operands are fine here.
+    (if (not (= (length args) 2))
+      (values #f
+        (list (error-at expr 'bad-primitive-arity
+                "shift primitive needs exactly two operands"
+                expr)))
+      (let-values ([(arg-irs errors) (infer-args args env type-names)])
+        (let* ([types (ir-list-types arg-irs)]
+               [operand-errors (operand-type-errors 'numeric types args expr)]
+               [ok? (and (null? errors) (null? operand-errors)
+                         (all-irs-valid? arg-irs))])
+          (values
+            (and ok?
+                 (make-typed-ir-call (car types)
+                   (expr-source expr) 'prim-shift op arg-irs '()))
+            (append errors operand-errors))))))
+
   (def (bad-constructor-arity expr name expected args)
     (list (error-at expr 'bad-call-arity
             "typed constructor arity does not match"
@@ -1427,6 +1480,10 @@
               (infer-debug-string args env type-names expr)]
              [(not and or)
               (infer-boolean head args env type-names expr)]
+             [(bitwise-and bitwise-ior bitwise-xor bitwise-not)
+              (infer-bitwise head args env type-names expr)]
+             [(bitwise-arithmetic-shift-left bitwise-arithmetic-shift-right)
+              (infer-shift head 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 1e5aa61..539682d 100644
--- a/lib/jerboa/typed/core.ss
+++ b/lib/jerboa/typed/core.ss
@@ -94,6 +94,8 @@
       prim-cmp
       prim-eq
       prim-bool
+      prim-bitwise
+      prim-shift
       string-length
       string-append
       bytevector-length
diff --git a/lib/jerboa/typed/rust.ss b/lib/jerboa/typed/rust.ss
index e7e68b5..892914a 100644
--- a/lib/jerboa/typed/rust.ss
+++ b/lib/jerboa/typed/rust.ss
@@ -1292,6 +1292,16 @@
       (emit-expression (car args))
       ").len() as u64"))
 
+  (def (emit-bytevector-u8-ref args)
+    (unless (= (length args) 2)
+      (error 'typed-rust "bytevector-u8-ref expects two operands" args))
+    (string-append
+      "("
+      (emit-expression (car args))
+      "[("
+      (emit-expression (cadr args))
+      ") as usize] as u64)"))
+
   (def (emit-equality args)
     (unless (= (length args) 2)
       (error 'typed-rust "equal? expects two operands" args))
@@ -1474,9 +1484,27 @@
            [(and) (emit-bool-chain "&&" args)]
            [(or) (emit-bool-chain "||" args)]
            [else (error 'typed-rust "unknown bool operator" operator)])]
+        [(prim-bitwise)
+         (case operator
+           [(bitwise-not)
+            (unless (= (length args) 1)
+              (error 'typed-rust "bitwise-not expects one operand" args))
+            (string-append "(!" (emit-expression (car args)) ")")]
+           [(bitwise-and) (emit-binary-chain "&" args)]
+           [(bitwise-ior) (emit-binary-chain "|" args)]
+           [(bitwise-xor) (emit-binary-chain "^" args)]
+           [else (error 'typed-rust "unknown bitwise operator" operator)])]
+        [(prim-shift)
+         (emit-binary-chain
+           (case operator
+             [(bitwise-arithmetic-shift-left) "<<"]
+             [(bitwise-arithmetic-shift-right) ">>"]
+             [else (error 'typed-rust "unknown shift operator" operator)])
+           args)]
         [(string-length) (emit-string-length args)]
         [(string-append) (emit-string-append args)]
         [(bytevector-length) (emit-bytevector-length args)]
+        [(bytevector-u8-ref) (emit-bytevector-u8-ref args)]
         [(debug-string) (emit-debug-string args)]
         [(record-ctor)
          (let ([record (lookup-name (ir-call-info-ref ir 'record)
diff --git a/tests/fixtures/typed/rust-bitwise.ss b/tests/fixtures/typed/rust-bitwise.ss
new file mode 100644
index 0000000..8a364f0
--- /dev/null
+++ b/tests/fixtures/typed/rust-bitwise.ss
@@ -0,0 +1,45 @@
+(typed-library (sample typed bitwise-kernel)
+  (export band bior bxor bnot shl shr
+          low-byte high-nibble combine-bytes rotl8 xor-accumulate)
+
+  (def (band (x : Nat) (m : Nat)) : Nat
+    (bitwise-and x m))
+
+  (def (bior (x : Nat) (y : Nat)) : Nat
+    (bitwise-ior x y))
+
+  (def (bxor (x : Nat) (y : Nat)) : Nat
+    (bitwise-xor x y))
+
+  (def (bnot (x : Nat)) : Nat
+    (bitwise-not x))
+
+  (def (shl (x : Nat) (n : Nat)) : Nat
+    (bitwise-arithmetic-shift-left x n))
+
+  (def (shr (x : Nat) (n : Nat)) : Nat
+    (bitwise-arithmetic-shift-right x n))
+
+  ;; mask off the low 8 bits
+  (def (low-byte (x : Nat)) : Nat
+    (bitwise-and x 255))
+
+  ;; extract the high nibble of a byte
+  (def (high-nibble (b : Nat)) : Nat
+    (bitwise-and (bitwise-arithmetic-shift-right b 4) 15))
+
+  ;; pack two bytes big-endian into a 16-bit word
+  (def (combine-bytes (hi : Nat) (lo : Nat)) : Nat
+    (bitwise-ior (bitwise-arithmetic-shift-left (bitwise-and hi 255) 8)
+                 (bitwise-and lo 255)))
+
+  ;; rotate an 8-bit value left by n (fixed width via masks)
+  (def (rotl8 (x : Nat) (n : Nat)) : Nat
+    (bitwise-and
+      (bitwise-ior (bitwise-arithmetic-shift-left x n)
+                   (bitwise-arithmetic-shift-right x (- 8 n)))
+      255))
+
+  ;; xor accumulator, keep the low byte (constant-time-equal building block)
+  (def (xor-accumulate (acc : Nat) (b : Nat)) : Nat
+    (bitwise-and (bitwise-xor acc b) 255)))
diff --git a/tests/fixtures/typed/rust-ct-equal.ss b/tests/fixtures/typed/rust-ct-equal.ss
new file mode 100644
index 0000000..b147e1c
--- /dev/null
+++ b/tests/fixtures/typed/rust-ct-equal.ss
@@ -0,0 +1,34 @@
+(typed-library (sample typed ct-equal)
+  (export byte-at fold-diff ct-equal differ?
+          first-byte last-byte xor-at)
+
+  ;; read one byte as a Nat
+  (def (byte-at (data : Bytes) (i : Nat)) : Nat
+    (bytevector-u8-ref data i))
+
+  ;; xor the i-th bytes of two buffers
+  (def (xor-at (a : Bytes) (b : Bytes) (i : Nat)) : Nat
+    (bitwise-xor (bytevector-u8-ref a i)
+                 (bytevector-u8-ref b i)))
+
+  ;; constant-time difference accumulator: OR every byte-xor together so
+  ;; the running time depends only on n, never on where a mismatch occurs.
+  (def (fold-diff (a : Bytes) (b : Bytes) (i : Nat) (n : Nat) (acc : Nat)) : Nat
+    (if (>= i n)
+      acc
+      (fold-diff a b (+ i 1) n
+                 (bitwise-ior acc (xor-at a b i)))))
+
+  ;; 0 iff the first n bytes of a and b are equal (constant time)
+  (def (ct-equal (a : Bytes) (b : Bytes) (n : Nat)) : Nat
+    (fold-diff a b 0 n 0))
+
+  ;; non-zero (truthy) iff the buffers differ over n bytes
+  (def (differ? (a : Bytes) (b : Bytes) (n : Nat)) : Bool
+    (> (ct-equal a b n) 0))
+
+  (def (first-byte (data : Bytes)) : Nat
+    (bytevector-u8-ref data 0))
+
+  (def (last-byte (data : Bytes) (n : Nat)) : Nat
+    (bytevector-u8-ref data (- n 1))))
diff --git a/tests/test-typed-checker.ss b/tests/test-typed-checker.ss
index 6be7cbb..2e4e164 100644
--- a/tests/test-typed-checker.ss
+++ b/tests/test-typed-checker.ss
@@ -597,6 +597,30 @@
          (bytevector-length x))))
   '(argument-type-mismatch))
 
+(test "builtin bytevector-u8-ref returns Nat"
+  (error-kinds
+    '(typed-library (body bv-ref-ok)
+       (export f)
+       (def (f (data : Bytes) (i : Nat)) : Nat
+         (bytevector-u8-ref data i))))
+  '())
+
+(test "builtin bytevector-u8-ref rejects non-Bytes buffer"
+  (error-kinds
+    '(typed-library (body bv-ref-bad-buf)
+       (export f)
+       (def (f (data : String) (i : Nat)) : Nat
+         (bytevector-u8-ref data i))))
+  '(argument-type-mismatch))
+
+(test "builtin bytevector-u8-ref rejects bad arity"
+  (error-kinds
+    '(typed-library (body bv-ref-bad-arity)
+       (export f)
+       (def (f (data : Bytes)) : Nat
+         (bytevector-u8-ref data))))
+  '(bad-call-arity))
+
 (test "record constructor and accessor calls"
   (error-kinds
     '(typed-library (body record-ok)
@@ -835,6 +859,46 @@
          (+ x 1))))
   '(operand-type-mismatch))
 
+(test "bitwise primitive returns numeric type"
+  (error-kinds
+    '(typed-library (body bitwise)
+       (export f)
+       (def (f (x : Nat) (m : Nat)) : Nat
+         (bitwise-and (bitwise-ior x m) (bitwise-xor x (bitwise-not m))))))
+  '())
+
+(test "bitwise primitive rejects nonnumeric operand"
+  (error-kinds
+    '(typed-library (body bitwise-bad)
+       (export f)
+       (def (f (x : String)) : Nat
+         (bitwise-and x 1))))
+  '(operand-type-mismatch))
+
+(test "bitwise-not rejects bad arity"
+  (error-kinds
+    '(typed-library (body bitnot-bad)
+       (export f)
+       (def (f (x : Nat)) : Nat
+         (bitwise-not x x))))
+  '(bad-primitive-arity))
+
+(test "shift primitive returns numeric type"
+  (error-kinds
+    '(typed-library (body shift)
+       (export f)
+       (def (f (x : Nat) (n : Nat)) : Nat
+         (bitwise-arithmetic-shift-left x n))))
+  '())
+
+(test "shift primitive rejects bad arity"
+  (error-kinds
+    '(typed-library (body shift-bad)
+       (export f)
+       (def (f (x : Nat)) : Nat
+         (bitwise-arithmetic-shift-left x))))
+  '(bad-primitive-arity))
+
 (test "comparison primitive feeds if condition"
   (error-kinds
     '(typed-library (body compare)
diff --git a/tests/test-typed-rust.ss b/tests/test-typed-rust.ss
index 3cd1ec5..b653474 100644
--- a/tests/test-typed-rust.ss
+++ b/tests/test-typed-rust.ss
@@ -143,6 +143,47 @@
   (string-append rust-file-header safe-prelude
     "pub fn bytes_length(data: Vec<u8>) -> u64 {\n    (data).len() as u64\n}\n\n#[unsafe(no_mangle)]\npub extern \"C\" fn jt_sample_typed_bytes_bytes_length(data_ptr: *const u8, data_len: usize) -> u64 {\n    match std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {\n        let data = {\n            let bytes: &[u8] = if data_ptr.is_null() {\n                &[]\n            } else {\n                // unsafe: pointer and length are produced by the generated Jerboa wrapper.\n                unsafe { std::slice::from_raw_parts(data_ptr, data_len) }\n            };\n            bytes.to_vec()\n        };\n        bytes_length(data)\n    })) {\n        Ok(value) => value,\n        Err(payload) => { jt_capture_panic(payload); 0u64 },\n    }\n}\n\n"))
 
+(define bitwise-form
+  '(typed-library (sample typed bitwise)
+     (export bit-and bit-ior bit-xor bit-not shl shr mask-byte)
+     (def (bit-and (x : Nat) (m : Nat)) : Nat
+       (bitwise-and x m))
+     (def (bit-ior (x : Nat) (y : Nat)) : Nat
+       (bitwise-ior x y))
+     (def (bit-xor (x : Nat) (y : Nat)) : Nat
+       (bitwise-xor x y))
+     (def (bit-not (x : Nat)) : Nat
+       (bitwise-not x))
+     (def (shl (x : Nat) (n : Nat)) : Nat
+       (bitwise-arithmetic-shift-left x n))
+     (def (shr (x : Nat) (n : Nat)) : Nat
+       (bitwise-arithmetic-shift-right x n))
+     (def (mask-byte (x : Nat) (n : Nat)) : Nat
+       (bitwise-and (bitwise-arithmetic-shift-right x n) 255))))
+
+(define bitwise-rust
+  (typed-library-form->rust-string bitwise-form))
+
+;; bytevector-u8-ref + recursion + bitwise = constant-time-equal?, the
+;; headline crypto primitive: XOR-accumulate over two buffers, 0 iff equal.
+(define bvref-form
+  '(typed-library (sample typed bvref)
+     (export byte-at ct-acc ct-equal)
+     (def (byte-at (data : Bytes) (i : Nat)) : Nat
+       (bytevector-u8-ref data i))
+     (def (ct-acc (a : Bytes) (b : Bytes) (i : Nat) (n : Nat) (acc : Nat)) : Nat
+       (if (>= i n)
+         acc
+         (ct-acc a b (+ i 1) n
+                 (bitwise-ior acc
+                              (bitwise-xor (bytevector-u8-ref a i)
+                                           (bytevector-u8-ref b i))))))
+     (def (ct-equal (a : Bytes) (b : Bytes) (n : Nat)) : Nat
+       (ct-acc a b 0 n 0))))
+
+(define bvref-rust
+  (typed-library-form->rust-string bvref-form))
+
 (define return-form
   '(typed-library (sample typed return-values)
      (export greeting echo-text echo-bytes)
@@ -383,6 +424,33 @@
   (typed-library-form->rust-string ops-form)
   ops-rust)
 
+(test "rust lowers bitwise and shift ops"
+  (and (substring? bitwise-rust "pub fn bit_and(x: u64, m: u64) -> u64")
+       (substring? bitwise-rust "(x & m)")
+       (substring? bitwise-rust "pub fn bit_ior(x: u64, y: u64) -> u64")
+       (substring? bitwise-rust "(x | y)")
+       (substring? bitwise-rust "pub fn bit_xor(x: u64, y: u64) -> u64")
+       (substring? bitwise-rust "(x ^ y)")
+       (substring? bitwise-rust "pub fn bit_not(x: u64) -> u64")
+       (substring? bitwise-rust "(!x)")
+       (substring? bitwise-rust "pub fn shl(x: u64, n: u64) -> u64")
+       (substring? bitwise-rust "(x << n)")
+       (substring? bitwise-rust "pub fn shr(x: u64, n: u64) -> u64")
+       (substring? bitwise-rust "(x >> n)")
+       ;; nested: mask off the low byte after a shift
+       (substring? bitwise-rust "((x >> n) & 255u64)"))
+  #t)
+
+(test "rust lowers bytevector-u8-ref and constant-time-equal"
+  (and (substring? bvref-rust "pub fn byte_at(data: Vec<u8>, i: u64) -> u64")
+       (substring? bvref-rust "(data[(i) as usize] as u64)")
+       ;; the constant-time core: XOR two indexed bytes, OR into accumulator
+       (substring? bvref-rust
+         "((a[(i) as usize] as u64) ^ (b[(i) as usize] as u64))")
+       (substring? bvref-rust
+         "(acc | ((a[(i) as usize] as u64) ^ (b[(i) as usize] as u64)))"))
+  #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>")