feat: Implement u32_less_than comparison primitive for Aiur

.github/workflows/ignored.yml

@@ -44,4 +44,4 @@ jobs:
             --errors-for-leak-kinds=definite \
             --track-origins=yes \
             --error-exitcode=1 \
-            .lake/build/bin/IxTests -- --include-ignored aiur aiur-hashes ixvm
+            .lake/build/bin/IxTests -- --include-ignored aiur aiur-hashes ixvm rbtree-map

Ix/Aiur/Bytecode.lean

@@ -34,6 +34,7 @@ inductive Op
   | u8And : ValIdx → ValIdx → Op
   | u8Or : ValIdx → ValIdx → Op
   | u8LessThan : ValIdx → ValIdx → Op
+  | u32LessThan : ValIdx → ValIdx → Op
   | debug : String → Option (Array ValIdx) → Op
   deriving Repr
 

Ix/Aiur/Check.lean

@@ -116,6 +116,7 @@ partial def inferTerm (t : Term) : CheckM TypedTerm := match t with
   | .u8Add a b => inferBinop a b .u8Add (.tuple #[.field, .field])
   | .u8Sub a b => inferBinop a b .u8Sub (.tuple #[.field, .field])
   | .u8LessThan a b => inferBinop a b .u8LessThan .field
+  | .u32LessThan a b => inferBinop a b .u32LessThan .field
   | .ioGetInfo key => inferIoGetInfo key
   | .ioSetInfo key idx len ret => inferIoSetInfo key idx len ret
   | .ioRead idx len => inferIoRead idx len

Ix/Aiur/Compile.lean

@@ -135,6 +135,11 @@ def opLayout : Bytecode.Op → LayoutM Unit
     pushDegree 1
     bumpAuxiliaries 1
     bumpLookups
+  | .u32LessThan .. => do
+    pushDegree 1
+    bumpAuxiliaries 24
+    modify fun stt => { stt with
+      functionLayout := { stt.functionLayout with lookups := stt.functionLayout.lookups + 8 } }
   | .debug .. => pure ()
 
 partial def blockLayout (block : Bytecode.Block) : LayoutM Unit := do
@@ -467,6 +472,10 @@ partial def toIndex
     let i ← expectIdx i
     let j ← expectIdx j
     pushOp (.u8LessThan i j)
+  | .u32LessThan i j => do
+    let i ← expectIdx i
+    let j ← expectIdx j
+    pushOp (.u32LessThan i j)
   | .debug label term ret => do
     let term ← term.mapM (toIndex layoutMap bindings)
     modify fun stt => { stt with ops := stt.ops.push (.debug label term) }

Ix/Aiur/Meta.lean

@@ -132,6 +132,7 @@ syntax "u8_sub" "(" trm ", " trm ")"                          : trm
 syntax "u8_and" "(" trm ", " trm ")"                          : trm
 syntax "u8_or" "(" trm ", " trm ")"                           : trm
 syntax "u8_less_than" "(" trm ", " trm ")"                    : trm
+syntax "u32_less_than" "(" trm ", " trm ")"                   : trm
 syntax "dbg!" "(" str (", " trm)? ")" ";" (trm)?              : trm
 
 syntax trm "[" "@" noWs ident "]"                                                        : trm
@@ -237,6 +238,8 @@ partial def elabTrm : ElabStxCat `trm
     mkAppM ``Term.u8Or #[← elabTrm i, ← elabTrm j]
   | `(trm| u8_less_than($i:trm, $j:trm)) => do
     mkAppM ``Term.u8LessThan #[← elabTrm i, ← elabTrm j]
+  | `(trm| u32_less_than($i:trm, $j:trm)) => do
+    mkAppM ``Term.u32LessThan #[← elabTrm i, ← elabTrm j]
   | `(trm| dbg!($label:str $[, $t:trm]?); $[$ret:trm]?) => do
     let t ← match t with
       | none => mkAppOptM ``Option.none #[some (mkConst ``Term)]
@@ -394,6 +397,10 @@ where
       let i ← replaceToken old new i
       let j ← replaceToken old new j
       `(trm| u8_less_than($i, $j))
+    | `(trm| u32_less_than($i:trm, $j:trm)) => do
+      let i ← replaceToken old new i
+      let j ← replaceToken old new j
+      `(trm| u32_less_than($i, $j))
     | `(trm| dbg!($label:str $[, $t:trm]?); $[$ret:trm]?) => do
       let t' ← t.mapM $ replaceToken old new
       let ret' ← ret.mapM $ replaceToken old new

Ix/Aiur/Term.lean

@@ -100,6 +100,7 @@ inductive Term
   | u8And : Term → Term → Term
   | u8Or : Term → Term → Term
   | u8LessThan : Term → Term → Term
+  | u32LessThan : Term → Term → Term
   | debug : String → Option Term → Term → Term
   deriving Repr, BEq, Hashable, Inhabited
 
@@ -147,6 +148,7 @@ inductive TypedTermInner
   | u8And : TypedTerm → TypedTerm → TypedTermInner
   | u8Or : TypedTerm → TypedTerm → TypedTermInner
   | u8LessThan : TypedTerm → TypedTerm → TypedTermInner
+  | u32LessThan : TypedTerm → TypedTerm → TypedTermInner
   | debug : String → Option TypedTerm → TypedTerm → TypedTermInner
   deriving Repr, Inhabited
 

Ix/IxVM/RBTreeMap.lean

@@ -0,0 +1,159 @@
+module
+public import Ix.Aiur.Meta
+
+public section
+
+namespace IxVM
+
+def rbTreeMap := ⟦
+  enum RBTreeMap {
+    -- Color {0: red, 1: black}, key, value, left, right
+    Node(G, G, G, &RBTreeMap, &RBTreeMap),
+    Nil
+  }
+
+  -- Okasaki-style balance: fixes red-red violations after insertion
+  fn rbtree_map_balance(color: G, key: G, value: G, left: RBTreeMap, right: RBTreeMap) -> RBTreeMap {
+    match (color, left, right) {
+      -- Case 1: left child red, left-left grandchild red
+      (1, RBTreeMap.Node(0, yk, yv, &RBTreeMap.Node(0, xk, xv, a, b), c), d) =>
+        RBTreeMap.Node(0, yk, yv,
+          store(RBTreeMap.Node(1, xk, xv, a, b)),
+          store(RBTreeMap.Node(1, key, value, c, store(d)))),
+      -- Case 2: left child red, left-right grandchild red
+      (1, RBTreeMap.Node(0, xk, xv, a, &RBTreeMap.Node(0, yk, yv, b, c)), d) =>
+        RBTreeMap.Node(0, yk, yv,
+          store(RBTreeMap.Node(1, xk, xv, a, b)),
+          store(RBTreeMap.Node(1, key, value, c, store(d)))),
+      -- Case 3: right child red, right-left grandchild red
+      (1, a, RBTreeMap.Node(0, zk, zv, &RBTreeMap.Node(0, yk, yv, b, c), d)) =>
+        RBTreeMap.Node(0, yk, yv,
+          store(RBTreeMap.Node(1, key, value, store(a), b)),
+          store(RBTreeMap.Node(1, zk, zv, c, d))),
+      -- Case 4: right child red, right-right grandchild red
+      (1, a, RBTreeMap.Node(0, yk, yv, b, &RBTreeMap.Node(0, zk, zv, c, d))) =>
+        RBTreeMap.Node(0, yk, yv,
+          store(RBTreeMap.Node(1, key, value, store(a), b)),
+          store(RBTreeMap.Node(1, zk, zv, c, d))),
+      -- No violation
+      _ => RBTreeMap.Node(color, key, value, store(left), store(right)),
+    }
+  }
+
+  -- Internal insert (new nodes are red)
+  fn rbtree_map_ins(key: G, value: G, tree: RBTreeMap) -> RBTreeMap {
+    match tree {
+      RBTreeMap.Nil =>
+        RBTreeMap.Node(0, key, value, store(RBTreeMap.Nil), store(RBTreeMap.Nil)),
+      RBTreeMap.Node(color, k, v, &left, &right) =>
+        let lt = u32_less_than(key, k);
+        match lt {
+          1 => rbtree_map_balance(color, k, v, rbtree_map_ins(key, value, left), right),
+          _ =>
+            let gt = u32_less_than(k, key);
+            match gt {
+              1 => rbtree_map_balance(color, k, v, left, rbtree_map_ins(key, value, right)),
+              _ => RBTreeMap.Node(color, key, value, store(left), store(right)),
+            },
+        },
+    }
+  }
+
+  -- Public insert: inserts key-value pair, enforces black root
+  fn rbtree_map_insert(key: G, value: G, tree: RBTreeMap) -> RBTreeMap {
+    let result = rbtree_map_ins(key, value, tree);
+    match result {
+      RBTreeMap.Node(_, k, v, left, right) => RBTreeMap.Node(1, k, v, left, right),
+    }
+  }
+
+  -- Lookup: returns the value associated with the key.
+  -- Fails if the key is not found (no Nil case).
+  fn rbtree_map_lookup(key: G, tree: RBTreeMap) -> G {
+    match tree {
+      RBTreeMap.Node(_, k, v, &left, &right) =>
+        let lt = u32_less_than(key, k);
+        match lt {
+          1 => rbtree_map_lookup(key, left),
+          _ =>
+            let gt = u32_less_than(k, key);
+            match gt {
+              1 => rbtree_map_lookup(key, right),
+              _ => v,
+            },
+        },
+    }
+  }
+
+  /- # Test entrypoints -/
+
+  fn rbtree_map_test() -> [G; 22] {
+    -- Single insert and lookup
+    let tree = rbtree_map_insert(5, 42, RBTreeMap.Nil);
+    let r0 = rbtree_map_lookup(5, tree);
+
+    -- Three inserts
+    let tree = rbtree_map_insert(10, 100, RBTreeMap.Nil);
+    let tree = rbtree_map_insert(20, 200, tree);
+    let tree = rbtree_map_insert(5, 50, tree);
+    let r1 = rbtree_map_lookup(5, tree);
+    let r2 = rbtree_map_lookup(10, tree);
+    let r3 = rbtree_map_lookup(20, tree);
+
+    -- Overwrite
+    let tree = rbtree_map_insert(10, 100, RBTreeMap.Nil);
+    let tree = rbtree_map_insert(10, 999, tree);
+    let r4 = rbtree_map_lookup(10, tree);
+
+    -- Ascending insertion order
+    let tree = rbtree_map_insert(1, 10, RBTreeMap.Nil);
+    let tree = rbtree_map_insert(2, 20, tree);
+    let tree = rbtree_map_insert(3, 30, tree);
+    let tree = rbtree_map_insert(4, 40, tree);
+    let tree = rbtree_map_insert(5, 50, tree);
+    let r5 = rbtree_map_lookup(1, tree);
+    let r6 = rbtree_map_lookup(2, tree);
+    let r7 = rbtree_map_lookup(3, tree);
+    let r8 = rbtree_map_lookup(4, tree);
+    let r9 = rbtree_map_lookup(5, tree);
+
+    -- Descending insertion order
+    let tree = rbtree_map_insert(5, 50, RBTreeMap.Nil);
+    let tree = rbtree_map_insert(4, 40, tree);
+    let tree = rbtree_map_insert(3, 30, tree);
+    let tree = rbtree_map_insert(2, 20, tree);
+    let tree = rbtree_map_insert(1, 10, tree);
+    let r10 = rbtree_map_lookup(1, tree);
+    let r11 = rbtree_map_lookup(2, tree);
+    let r12 = rbtree_map_lookup(3, tree);
+    let r13 = rbtree_map_lookup(4, tree);
+    let r14 = rbtree_map_lookup(5, tree);
+
+    -- Mixed insertion order
+    let tree = rbtree_map_insert(50, 500, RBTreeMap.Nil);
+    let tree = rbtree_map_insert(30, 300, tree);
+    let tree = rbtree_map_insert(70, 700, tree);
+    let tree = rbtree_map_insert(20, 200, tree);
+    let tree = rbtree_map_insert(40, 400, tree);
+    let tree = rbtree_map_insert(60, 600, tree);
+    let tree = rbtree_map_insert(80, 800, tree);
+    let r15 = rbtree_map_lookup(20, tree);
+    let r16 = rbtree_map_lookup(30, tree);
+    let r17 = rbtree_map_lookup(40, tree);
+    let r18 = rbtree_map_lookup(50, tree);
+    let r19 = rbtree_map_lookup(60, tree);
+    let r20 = rbtree_map_lookup(70, tree);
+    let r21 = rbtree_map_lookup(80, tree);
+
+    [r0,
+     r1, r2, r3,
+     r4,
+     r5, r6, r7, r8, r9,
+     r10, r11, r12, r13, r14,
+     r15, r16, r17, r18, r19, r20, r21]
+  }
+⟧
+
+end IxVM
+
+end

Tests/Aiur.lean

@@ -3,3 +3,4 @@ module
 public import Tests.Aiur.Common
 public import Tests.Aiur.Aiur
 public import Tests.Aiur.Hashes
+public import Tests.Aiur.RBTreeMap

Tests/Aiur/Aiur.lean

@@ -219,6 +219,10 @@ def toplevel := ⟦
     u8_or(i, j)
   }
 
+  fn u32_less_than_function(x: G, y: G) -> G {
+    u32_less_than(x, y)
+  }
+
   fn fold_matrix_sum(m: [[G; 2]; 2]) -> G {
     fold(0 .. 2, 0, |acc_outer, @i|
       fold(0 .. 2, acc_outer, |acc_inner, @j|
@@ -270,6 +274,11 @@ def aiurTestCases : List AiurTestCase := [
     .noIO `u8_and_function #[45, 131] #[1],
     .noIO `u8_or_function #[45, 131] #[175],
     .noIO `fold_matrix_sum #[1, 2, 3, 4] #[10],
+    { AiurTestCase.noIO `u32_less_than_function #[300, 500] #[1] with label := "u32_less_than(300,500)" },
+    { AiurTestCase.noIO `u32_less_than_function #[500, 300] #[0] with label := "u32_less_than(500,300)" },
+    { AiurTestCase.noIO `u32_less_than_function #[500, 500] #[0] with label := "u32_less_than(500,500)" },
+    { AiurTestCase.noIO `u32_less_than_function #[0, 1] #[1] with label := "u32_less_than(0,1)" },
+    { AiurTestCase.noIO `u32_less_than_function #[0, 0] #[0] with label := "u32_less_than(0,0)" },
   ]
 
 end

Tests/Aiur/RBTreeMap.lean

@@ -0,0 +1,18 @@
+module
+
+public import Tests.Aiur.Common
+public import Ix.IxVM.RBTreeMap
+
+public section
+
+public def rbTreeMapTestCases : List AiurTestCase := [
+  .noIO `rbtree_map_test #[] #[
+    42,
+    50, 100, 200,
+    999,
+    10, 20, 30, 40, 50,
+    10, 20, 30, 40, 50,
+    200, 300, 400, 500, 600, 700, 800],
+]
+
+end

Tests/Main.lean

@@ -70,6 +70,11 @@ def ignoredRunners : List (String × IO UInt32) := [
     return if r1 == 0 && r2 == 0 then 0 else 1),
   ("ixvm", do
     LSpec.lspecIO (.ofList [("ixvm", [mkAiurTests IxVM.ixVM [← serdeNatAddComm]])]) []),
+  ("rbtree-map", do
+    IO.println "rbtree-map"
+    match AiurTestEnv.build (pure IxVM.rbTreeMap) with
+    | .error e => IO.eprintln s!"RBTreeMap setup failed: {e}"; return 1
+    | .ok env => LSpec.lspecEachIO rbTreeMapTestCases fun tc => pure (env.runTestCase tc)),
 ]
 
 def main (args : List String) : IO UInt32 := do

src/aiur/bytecode.rs

@@ -57,6 +57,7 @@ pub enum Op {
   U8And(ValIdx, ValIdx),
   U8Or(ValIdx, ValIdx),
   U8LessThan(ValIdx, ValIdx),
+  U32LessThan(ValIdx, ValIdx),
   Debug(String, Option<Vec<ValIdx>>),
 }