IxVM kernel: cut Aiur FFT cost on reduction- and check-heavy shards

Ix/IxVM/Convert.lean

@@ -68,7 +68,9 @@ def convert := ⟦
     match load(idxs) {
       ListNode.Nil => store(ListNode.Nil),
       ListNode.Cons(idx, rest) =>
-        let u = load(list_lookup_u64(univs, idx));
+        -- universe indices are small; walk with a field index (cheap per-step
+        -- field sub) instead of `list_lookup_u64`'s per-step U64 predecessor.
+        let u = load(list_lookup(univs, flatten_u64(idx)));
         store(ListNode.Cons(store(convert_univ(u)), convert_univ_idxs(rest, univs))),
     }
   }
@@ -145,7 +147,9 @@ def convert := ⟦
   ) -> KExpr {
     match load(e) {
       Expr.Srt(univ_idx) =>
-        let u = load(list_lookup_u64(univs, univ_idx));
+        -- field-indexed walk (see `convert_univ_idxs`): avoids the per-step
+        -- U64 predecessor of `list_lookup_u64` on this hot universe lookup.
+        let u = load(list_lookup(univs, flatten_u64(univ_idx)));
         store(KExprNode.Srt(store(convert_univ(u)))),
 
       Expr.Var(idx) =>
@@ -199,7 +203,8 @@ def convert := ⟦
           convert_expr(body, sharing, ref_idxs, recur_idxs, lit_blobs, univs))),
 
       Expr.Share(idx) =>
-        convert_expr(list_lookup_u64(sharing, idx), sharing, ref_idxs, recur_idxs, lit_blobs, univs),
+        let ListNode.Cons(e, _) = load(list_drop(sharing, flatten_u64(idx)));
+        convert_expr(e, sharing, ref_idxs, recur_idxs, lit_blobs, univs),
     }
   }
 

Ix/IxVM/Ingress.lean

@@ -53,17 +53,23 @@ def ingress := ⟦
     bytes
   }
 
-  -- Compare two 32-byte addresses for equality
-  fn address_eq(a: Addr, b: Addr) -> G {
-    let [a0, a1, a2, a3, a4, a5, a6, a7,
+  -- Compare two 32-byte addresses for equality.
+  --
+  -- Cold path: limb 0 already matched, compare the remaining 31 limbs.
+  -- Factored into its own function so it forms a separate circuit whose height
+  -- is only the (rare) limb-0-match rows; Aiur charges a function's full width
+  -- on every one of its rows, so a nested match in `address_eq` would not save
+  -- anything — the split must be a function boundary.
+  fn address_eq_tail(a: Addr, b: Addr) -> G {
+    let [_, a1, a2, a3, a4, a5, a6, a7,
          a8, a9, a10, a11, a12, a13, a14, a15,
          a16, a17, a18, a19, a20, a21, a22, a23,
          a24, a25, a26, a27, a28, a29, a30, a31] = load(a);
-    let [b0, b1, b2, b3, b4, b5, b6, b7,
+    let [_, b1, b2, b3, b4, b5, b6, b7,
          b8, b9, b10, b11, b12, b13, b14, b15,
          b16, b17, b18, b19, b20, b21, b22, b23,
          b24, b25, b26, b27, b28, b29, b30, b31] = load(b);
-    match [to_field(a0) - to_field(b0), to_field(a1) - to_field(b1),
+    match [to_field(a1) - to_field(b1),
            to_field(a2) - to_field(b2), to_field(a3) - to_field(b3),
            to_field(a4) - to_field(b4), to_field(a5) - to_field(b5),
            to_field(a6) - to_field(b6), to_field(a7) - to_field(b7),
@@ -80,7 +86,20 @@ def ingress := ⟦
            to_field(a28) - to_field(b28), to_field(a29) - to_field(b29),
            to_field(a30) - to_field(b30), to_field(a31) - to_field(b31)] {
       [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
-       0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0] => 1,
+       0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0] => 1,
+      _ => 0,
+    }
+  }
+
+  -- Limb-0 prefilter: a single differing limb proves inequality, and almost
+  -- every comparison (the primitive-dispatch gauntlet in whnf) mismatches at
+  -- limb 0. Hot rows reject here at narrow width; only limb-0 matches pay the
+  -- wide `address_eq_tail` compare. Identical result to a full 32-limb compare.
+  fn address_eq(a: Addr, b: Addr) -> G {
+    let av = load(a);
+    let bv = load(b);
+    match to_field(av[0]) - to_field(bv[0]) {
+      0 => address_eq_tail(a, b),
       _ => 0,
     }
   }
@@ -792,7 +811,9 @@ def ingress := ⟦
   -- Deref Expr.Share via the constant's sharing list.
   fn deref_share(e: Expr, sharing: List‹&Expr›) -> Expr {
     match e {
-      Expr.Share(idx) => deref_share(load(list_lookup_u64(sharing, idx)), sharing),
+      Expr.Share(idx) =>
+        let ListNode.Cons(e, _) = load(list_drop(sharing, flatten_u64(idx)));
+        deref_share(load(e), sharing),
       _ => e,
     }
   }

Ix/IxVM/IxonDeserialize.lean

@@ -101,15 +101,17 @@ def ixonDeserialize := ⟦
   -- Expression deserialization
   -- ============================================================================
 
-  -- App telescope: read count args, wrapping func in App nodes
-  fn get_app_telescope(func: Expr, stream: ByteStream, count: U64) -> (Expr, ByteStream) {
+  -- App telescope: read count args, wrapping func in App nodes. `func` is passed
+  -- by pointer (loaded only at the base case) so the recursion doesn't carry the
+  -- wide `Expr` union by value on every row.
+  fn get_app_telescope(func: &Expr, stream: ByteStream, count: U64) -> (Expr, ByteStream) {
     let is_zero = u64_is_zero(count);
     match is_zero {
-      1 => (func, stream),
+      1 => (load(func), stream),
       0 =>
         let (arg, s) = get_expr(stream);
-        let app = Expr.App(store(func), store(arg));
-        get_app_telescope(app, s, relaxed_u64_pred(count)),
+        let app = Expr.App(func, store(arg));
+        get_app_telescope(store(app), s, relaxed_u64_pred(count)),
     }
   }
 
@@ -180,7 +182,7 @@ def ixonDeserialize := ⟦
       -- App: Tag4(0x7, count) + func + args...
       0x7 =>
         let (func, s2) = get_expr(s);
-        get_app_telescope(func, s2, size),
+        get_app_telescope(store(func), s2, size),
 
       -- Lam: Tag4(0x8, count) + types... + body
       0x8 => get_lam_telescope(s, size),
@@ -189,17 +191,23 @@ def ixonDeserialize := ⟦
       0x9 => get_all_telescope(s, size),
 
       -- Let: Tag4(0xA, non_dep) + expr(ty) + expr(val) + expr(body)
-      0xA =>
-        let (ty, s2) = get_expr(s);
-        let (val, s3) = get_expr(s2);
-        let (body, s4) = get_expr(s3);
-        (Expr.Let(size, store(ty), store(val), store(body)), s4),
+      0xA => get_expr_let(s, size),
 
       -- Share: Tag4(0xB, idx)
       0xB => (Expr.Share(size), s),
     }
   }
 
+  -- Let arm of get_expr, split out: three recursive `get_expr` calls make it the
+  -- widest (and a rare) arm, so inlined it taxes every get_expr row.
+  fn get_expr_let(s: ByteStream, size: U64) -> (Expr, ByteStream) {
+    let (ty, s2) = get_expr(s);
+    let (val, s3) = get_expr(s2);
+    let (body, s4) = get_expr(s3);
+    (Expr.Let(size, store(ty), store(val), store(body)), s4)
+  }
+
+
   -- ============================================================================
   -- Universe deserialization
   -- ============================================================================

Ix/IxVM/Kernel/Claim.lean

@@ -804,6 +804,43 @@ def claim := ⟦
     }
   }
 
+  -- Pack the first 4 address bytes (LE) into a u32 key for the skip-set rbtree.
+  --
+  -- Capped at 4 bytes because `RBTreeMap` orders keys with `u32_less_than`, a
+  -- 32-bit comparator gadget — a wider key (a single `G` could hold 7 bytes in
+  -- Goldilocks) would overflow it and corrupt the tree ordering. A 32-bit key
+  -- space means key collisions are possible (~N²/2^33 over N leaves), but they
+  -- are harmless: a collision makes `addr_in_set`'s confirming `address_eq`
+  -- fail, yielding a false negative (a frontier const gets rechecked) never a
+  -- false positive (a wrong skip). See `build_skip_set`.
+  fn addr_key(a: Addr) -> G {
+    let arr = load(a);
+    to_field(arr[0])
+      + to_field(arr[1]) * 256
+      + to_field(arr[2]) * 65536
+      + to_field(arr[3]) * 16777216
+  }
+
+  -- Build an O(log N) membership set from the assumption-leaf list, keyed on
+  -- `addr_key`. Key collisions overwrite — sound because the only consequence is
+  -- a missed skip (a frontier const gets rechecked, never wrongly trusted); the
+  -- confirming `address_eq` in `addr_in_set` rules out false positives.
+  fn build_skip_set(leaves: List‹Addr›, acc: RBTreeMap‹Addr›) -> RBTreeMap‹Addr› {
+    match load(leaves) {
+      ListNode.Nil => acc,
+      ListNode.Cons(a, rest) =>
+        build_skip_set(rest, rbtree_map_insert(addr_key(a), a, acc)),
+    }
+  }
+
+  -- Membership via one rbtree lookup (cheap u32 key compares) + one confirming
+  -- full `address_eq`, replacing the O(N) linear `addr_in_list` scan that
+  -- dominated address_eq cost on sharded checks.
+  fn addr_in_set(target: Addr, skip_set: RBTreeMap‹Addr›) -> G {
+    let found = rbtree_map_lookup_or_default(addr_key(target), skip_set, store([0u8; 32]));
+    address_eq(found, target)
+  }
+
   -- ============================================================================
   -- check_all variant that skips positions whose addr is in the
   -- supplied assumption-leaf list.
@@ -813,24 +850,27 @@ def claim := ⟦
                         addrs: List‹Addr›,
                         asm_leaves: List‹Addr›) {
     let _ = check_canonical_block_sort(top);
-    check_all_skipping_iter(consts, top, addrs, asm_leaves, 0)
+    -- Build the skip-set once (O(N log N)) instead of an O(N) linear scan per
+    -- checked const.
+    let skip_set = build_skip_set(asm_leaves, RBTreeMap.Nil);
+    check_all_skipping_iter(consts, top, addrs, skip_set, 0)
   }
 
   fn check_all_skipping_iter(consts: List‹&KConstantInfo›,
                              top: List‹&KConstantInfo›,
                              addrs: List‹Addr›,
-                             asm_leaves: List‹Addr›,
+                             skip_set: RBTreeMap‹Addr›,
                              pos: G) {
     match load(consts) {
       ListNode.Nil => (),
       ListNode.Cons(&ci, rest) =>
         let addr = list_lookup(addrs, pos);
-        match addr_in_list(addr, asm_leaves) {
+        match addr_in_set(addr, skip_set) {
           1 =>
-            check_all_skipping_iter(rest, top, addrs, asm_leaves, pos + 1),
+            check_all_skipping_iter(rest, top, addrs, skip_set, pos + 1),
           _ =>
             let _ = check_const(ci, pos, top, addrs);
-            check_all_skipping_iter(rest, top, addrs, asm_leaves, pos + 1),
+            check_all_skipping_iter(rest, top, addrs, skip_set, pos + 1),
         },
     }
   }

Ix/IxVM/Kernel/DefEq.lean

@@ -146,7 +146,7 @@ def defEq := ⟦
 
   -- 1 iff ty is `Const(I, _) args` for non-rec 1-ctor 0-field inductive.
   fn is_unit_like_type(ty: KExpr, top: List‹&KConstantInfo›) -> G {
-    match collect_spine_simple(ty) {
+    match collect_spine(ty) {
       (head, _) =>
         match load(head) {
           KExprNode.Const(idx, _) =>
@@ -218,36 +218,66 @@ def defEq := ⟦
     }
   }
 
-  -- Mirror: src/ix/kernel/def_eq.rs:930-948 fn nat_succ_of.
-  -- `Lit(n)` n>0 → (1, Lit(n-1)). `App(Const(Nat.succ), arg)` → (1, arg).
-  -- Else (0, _).
-  fn nat_succ_of(e: KExpr, addrs: List‹Addr›) -> (G, KExpr) {
+
+  -- Decompose a WHNF'd Nat into `base + offset` where `base` is the
+  -- non-offset core and `offset` a KLimbs literal. Recognizes:
+  --   Lit n               -> (matched, 0-base, n)
+  --   Nat.succ e          -> base/offset of e, offset+1
+  --   Nat.add e (Lit m)   -> base/offset of e, offset+m
+  -- `matched=1` iff `e` is offset-shaped (succ/add/lit). The few succ layers
+  -- whnf exposes are peeled, but `Nat.add base (Lit m)` is read in O(1) — so a
+  -- `succ^k(x)` chain (which whnf leaves as `succ(Nat.add x (Lit k-1))`)
+  -- decomposes to `(x, k)` in O(1) instead of k unary steps.
+  fn nat_offset_of(e: KExpr, addrs: List‹Addr›) -> (G, KExpr, KLimbs) {
     match load(e) {
       KExprNode.Lit(lit) =>
         match lit {
-          KLiteral.Nat(limbs) =>
-            match klimbs_is_zero(limbs) {
-              1 => (0, store(KExprNode.BVar(0))),
-              0 => (1, mk_nat_lit(klimbs_dec(limbs))),
-            },
-          _ => (0, store(KExprNode.BVar(0))),
+          KLiteral.Nat(n) => (1, mk_nat_lit(store(ListNode.Nil)), n),
+          _ => (0, e, store(ListNode.Nil)),
         },
       KExprNode.App(f, a) =>
         match load(f) {
           KExprNode.Const(idx, _) =>
             match address_eq(list_lookup(addrs, idx), nat_succ_addr()) {
-              1 => (1, a),
-              0 => (0, store(KExprNode.BVar(0))),
+              1 =>
+                match nat_offset_of(a, addrs) {
+                  (_, base, o) => (1, base, klimbs_succ(o)),
+                },
+              0 => (0, e, store(ListNode.Nil)),
             },
-          _ => (0, store(KExprNode.BVar(0))),
+          KExprNode.App(g, x) =>
+            match load(g) {
+              KExprNode.Const(idx, _) =>
+                match address_eq(list_lookup(addrs, idx), nat_add_addr()) {
+                  1 =>
+                    match load(a) {
+                      KExprNode.Lit(alit) =>
+                        match alit {
+                          KLiteral.Nat(m) =>
+                            match nat_offset_of(x, addrs) {
+                              (1, base, o) => (1, base, klimbs_add(o, m)),
+                              (0, _, _) => (1, x, m),
+                            },
+                          _ => (0, e, store(ListNode.Nil)),
+                        },
+                      _ => (0, e, store(ListNode.Nil)),
+                    },
+                  0 => (0, e, store(ListNode.Nil)),
+                },
+              _ => (0, e, store(ListNode.Nil)),
+            },
+          _ => (0, e, store(ListNode.Nil)),
         },
-      _ => (0, store(KExprNode.BVar(0))),
+      _ => (0, e, store(ListNode.Nil)),
     }
   }
 
-  -- Mirror: src/ix/kernel/def_eq.rs:953-995 is_def_eq_nat / try_def_eq_offset.
-  -- Returns (matched, eq). `matched=1` iff both sides are nat-shaped (both
-  -- zero, both succ-headed, or both literals); `eq` is the verdict.
+  -- Mirror: src/ix/kernel/def_eq.rs:953-995 is_def_eq_nat / try_def_eq_offset,
+  -- generalized to offset form. Returns (matched, eq). Conservative: only
+  -- decides when both sides are offset-shaped with EQUAL offsets (then the
+  -- verdict is `base_a ≟ base_b`, sound because `+k` is injective); differing
+  -- offsets or non-offset shapes fall back (matched=0) to the generic path.
+  -- Collapses `succ^k(x) ≟ succ^k(x)` from k unary steps to one klimbs compare.
   fn try_def_eq_nat(a: KExpr, b: KExpr, types: List‹KExpr›,
                      top: List‹&KConstantInfo›,
                      addrs: List‹Addr›) -> (G, G) {
@@ -256,13 +286,19 @@ def defEq := ⟦
     match za * zb {
       1 => (1, 1),
       0 =>
-        match nat_succ_of(a, addrs) {
-          (1, ap) =>
-            match nat_succ_of(b, addrs) {
-              (1, bp) => (1, k_is_def_eq(ap, bp, types, top, addrs)),
-              _ => (0, 0),
+        match nat_offset_of(a, addrs) {
+          (ma, ba, oa) =>
+            match nat_offset_of(b, addrs) {
+              (mb, bb, ob) =>
+                match ma * mb {
+                  0 => (0, 0),
+                  _ =>
+                    match klimbs_eq(oa, ob) {
+                      0 => (0, 0),
+                      1 => (1, k_is_def_eq(ba, bb, types, top, addrs)),
+                    },
+                },
             },
-          _ => (0, 0),
         },
     }
   }