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

Ix/Cli/CheckCmd.lean

@@ -45,6 +45,16 @@ def parseName (arg : String) : Lean.Name :=
       | some n => Lean.Name.mkNum acc n
       | none   => Lean.Name.mkStr acc s
 
+/-- Resolve a CLI name argument against the env. `parseName` can't rebuild
+    private names (`_private.M.0.foo`) — the marker/scope-index components
+    don't round-trip through naive dot-splitting. So if the parsed name
+    isn't present, fall back to matching the arg against each constant's
+    `toString` (the displayed form the user copied). -/
+def resolveName (env : Lean.Environment) (arg : String) : Option Lean.Name :=
+  let parsed := parseName arg
+  if env.constants.contains parsed then some parsed
+  else (env.constants.toList.find? (fun (k, _) => toString k == arg)).map (·.1)
+
 def addrOfHex! (label : String) (s : String) : IO Address := do
   match Address.fromString s with
   | some a => pure a
@@ -88,6 +98,16 @@ partial def ixNameToLeanName : Ix.Name → Lean.Name
   | .str p s _ => .str (ixNameToLeanName p) s
   | .num p n _ => .num (ixNameToLeanName p) n
 
+/-- Resolve a CLI name argument against a `.ixe` env's named map. Like
+    `resolveName` for the compiled-in env: `parseName` can't rebuild private
+    names, so when the direct lookup misses, fall back to matching the arg
+    against each named constant's displayed `toString`. -/
+def resolveIxeAddr (ixonEnv : Ixon.Env) (arg : String) : Option Address :=
+  match ixonEnv.getAddr? (Ix.Name.fromLeanName (parseName arg)) with
+  | some a => some a
+  | none =>
+    (ixonEnv.named.toList.find? (fun (k, _) => toString (ixNameToLeanName k) == arg)).map (·.2.addr)
+
 /-- Build a `ClaimWitness` for the `verify_claim` entrypoint against
     `Ix.Claim.check addr none` (full transitive typecheck of the
     target's closure). -/
@@ -189,14 +209,13 @@ def forEachClaim
           if !keepGoing then break
     else
       for arg in names do
-        let name := parseName arg
-        match ixonEnv.getAddr? (Ix.Name.fromLeanName name) with
+        match resolveIxeAddr ixonEnv arg with
         | none =>
-          IO.eprintln s!"{name} not found in {path}"
-          failures := failures.push (toString name)
+          IO.eprintln s!"{arg} not found in {path}"
+          failures := failures.push arg
           if !keepGoing then break
         | some addr =>
-          let label := toString name
+          let label := arg
           let claim := Ix.Claim.check addr none
           let witness ← mkWitness addr ixonEnv
           if (← runOne claim witness label) ≠ 0 then
@@ -224,17 +243,18 @@ def forEachClaim
           if !keepGoing then break
     else
       for arg in names do
-        let name := parseName arg
-        if !env.constants.contains name then
-          IO.eprintln s!"{name} not found"
-          failures := failures.push (toString name)
+        match resolveName env arg with
+        | none =>
+          IO.eprintln s!"{arg} not found"
+          failures := failures.push arg
           if !keepGoing then break
           else continue
-        let label := toString name
-        let (claim, witness) ← buildOne name
-        if (← runOne claim witness label) ≠ 0 then
-          failures := failures.push label
-          if !keepGoing then break
+        | some name =>
+          let label := toString name
+          let (claim, witness) ← buildOne name
+          if (← runOne claim witness label) ≠ 0 then
+            failures := failures.push label
+            if !keepGoing then break
 
   if failures.isEmpty then pure 0
   else

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),
         },
     }
   }