perf(aiur): Canonical levels, memoized dispatch, flat query record

Cargo.toml

@@ -11,6 +11,7 @@ anyhow = "1"
 blake3 = "1.8.4"
 itertools = "0.14.0"
 indexmap = { version = "2", features = ["rayon"] }
+hashbrown = "0.15"
 lean-ffi = { git = "https://github.com/argumentcomputer/lean-ffi", rev = "cc98ebf67bf453ac3827cb767f78b13ea674dd6a" }
 mimalloc = { version = "0.1", default-features = false }
 multi-stark = { git = "https://github.com/argumentcomputer/multi-stark.git", rev = "9ecab51d553445c0cc7b571af00a76b8a83a6f8c" }

Ix/Cli/RefsOfCmd.lean

@@ -0,0 +1,70 @@
+/-
+  `ix refs-of <Lean.Name> --ixe <path>`: print the direct constant
+  references of a constant — its `Constant.refs` entries that resolve to
+  constants in the env (literal blobs are skipped: they are pinned by
+  content-addressing and are not well-typedness obligations).
+
+  Prints a comma-separated list of 64-char hex addresses on stdout, the
+  exact shape `ix tree canonical` expects, so a frontier assumption tree
+  for a single constant is:
+
+      ix tree canonical $(ix refs-of <name> --ixe env.ixe)
+-/
+module
+public import Cli
+public import Ix.Address
+public import Ix.Common
+public import Ix.Ixon
+
+public section
+
+namespace Ix.Cli.RefsOfCmd
+
+def parseName (arg : String) : Lean.Name :=
+  arg.splitOn "." |>.foldl (init := .anonymous)
+    fun acc s => match s.toNat? with
+      | some n => Lean.Name.mkNum acc n
+      | none   => Lean.Name.mkStr acc s
+
+def runRefsOfCmd (p : Cli.Parsed) : IO UInt32 := do
+  let some nameArg := p.positionalArg? "name"
+    | p.printError "error: must specify <name>"; return 1
+  let name := parseName (nameArg.as! String)
+  let some ixeFlag := p.flag? "ixe"
+    | p.printError "error: --ixe is required"; return 1
+  let path := ixeFlag.as! String
+  let bytes ← IO.FS.readBinFile path
+  let ixonEnv ← match Ixon.rsDeEnv bytes with
+    | .error e =>
+      IO.eprintln s!"error: failed to deserialize {path}: {e}"; return 1
+    | .ok env => pure env
+  let some addr := ixonEnv.getAddr? (Ix.Name.fromLeanName name)
+    | IO.eprintln s!"error: {name} not found in {path}"; return 1
+  let some c := ixonEnv.getConst? addr
+    | IO.eprintln s!"error: constant for {name} missing from env"; return 1
+  let mut seen : Std.HashSet String := {}
+  let mut out : Array String := #[]
+  for r in c.refs do
+    if (ixonEnv.getConst? r).isSome then
+      let s := toString r
+      if !seen.contains s then
+        seen := seen.insert s
+        out := out.push s
+  IO.println (String.intercalate "," out.toList)
+  return 0
+
+end Ix.Cli.RefsOfCmd
+
+open Ix.Cli.RefsOfCmd in
+def refsOfCmd : Cli.Cmd := `[Cli|
+  "refs-of" VIA runRefsOfCmd;
+  "Print a constant's direct constant references (comma-separated hex), the leaf list for a frontier assumption tree"
+
+  FLAGS:
+    "ixe" : String; "Path to a serialized `.ixe` env (required)."
+
+  ARGS:
+    name : String; "Fully-qualified Lean.Name (e.g. `Nat.add_comm`)."
+]
+
+end

Ix/IxVM/Kernel/Claim.lean

@@ -804,43 +804,22 @@ 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› {
+  -- Membership map over the assumption leaves, keyed by `ptr_val`. Sound
+  -- by the same invariant as `build_addr_pos_map` (Ingress.lean): one
+  -- pointer maps to one content, so a ptr hit implies the address IS a
+  -- leaf (no false skip); a de-interned pointer reads as absent, which
+  -- here only means the constant gets CHECKED instead of skipped —
+  -- conservative. One tree lookup per constant, no confirming
+  -- address_eq, no per-lookup address loads (ported from jcb/fixes
+  -- H-14; replaces the 4-byte-content-keyed set + confirm compare).
+  fn build_asm_leaf_map(leaves: List‹Addr›) -> RBTreeMap‹G› {
     match load(leaves) {
-      ListNode.Nil => acc,
+      ListNode.Nil => RBTreeMap.Nil,
       ListNode.Cons(a, rest) =>
-        build_skip_set(rest, rbtree_map_insert(addr_key(a), a, acc)),
+        rbtree_map_insert(ptr_val(a), 1, build_asm_leaf_map(rest)),
     }
   }
 
-  -- 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.
@@ -850,27 +829,32 @@ def claim := ⟦
                         addrs: List‹Addr›,
                         asm_leaves: List‹Addr›) {
     let _ = check_canonical_block_sort(top);
-    -- 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)
+    let asm_map = store(build_asm_leaf_map(asm_leaves));
+    check_all_skipping_iter(consts, top, addrs, addrs, asm_map, 0)
   }
 
+  -- `cur_addrs` is the suffix of `addrs` aligned with `consts`, walked in
+  -- lockstep instead of `list_lookup(addrs, pos)` — which re-walks the
+  -- list prefix every iteration, a standalone quadratic in closure size
+  -- (ported from jcb/fixes H-14).
   fn check_all_skipping_iter(consts: List‹&KConstantInfo›,
                              top: List‹&KConstantInfo›,
                              addrs: List‹Addr›,
-                             skip_set: RBTreeMap‹Addr›,
+                             cur_addrs: List‹Addr›,
+                             asm_map: &RBTreeMap‹G›,
                              pos: G) {
     match load(consts) {
       ListNode.Nil => (),
       ListNode.Cons(&ci, rest) =>
-        let addr = list_lookup(addrs, pos);
-        match addr_in_set(addr, skip_set) {
-          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, skip_set, pos + 1),
+        match load(cur_addrs) {
+          ListNode.Cons(addr, rest_addrs) =>
+            match rbtree_map_lookup_or_default(ptr_val(addr), load(asm_map), 0) {
+              0 =>
+                let _ = check_const(ci, pos, top, addrs);
+                check_all_skipping_iter(rest, top, addrs, rest_addrs, asm_map, pos + 1),
+              _ =>
+                check_all_skipping_iter(rest, top, addrs, rest_addrs, asm_map, pos + 1),
+            },
         },
     }
   }