Ixon.Constant deserialization

Benchmarks/IxVM.lean

@@ -0,0 +1,45 @@
+import Ix.Meta
+import Ix.IxVM
+import Ix.Aiur.Simple
+import Ix.Aiur.Compile
+import Ix.Aiur.Protocol
+import Ix.Benchmark.Bench
+
+def commitmentParameters : Aiur.CommitmentParameters := {
+  logBlowup := 1
+}
+
+def friParameters : Aiur.FriParameters := {
+  logFinalPolyLen := 0
+  numQueries := 100
+  commitProofOfWorkBits := 20
+  queryProofOfWorkBits := 0
+}
+
+def main : IO Unit := do
+  let .ok toplevel := IxVM.ixVM
+    | throw (IO.userError "Merging failed")
+  let some funIdx := toplevel.getFuncIdx `ixon_serde_blake3_bench
+    | throw (IO.userError "Aiur function not found")
+  let .ok decls := toplevel.checkAndSimplify
+    | throw (IO.userError "Simplification failed")
+  let .ok bytecode := decls.compile
+    | throw (IO.userError "Compilation failed")
+  let aiurSystem := Aiur.AiurSystem.build bytecode commitmentParameters
+
+  let env ← get_env!
+  let natAddCommName := ``Nat.add_comm
+  let constList := Lean.collectDependencies natAddCommName env.constants
+  let rawEnv ← Ix.CompileM.rsCompileEnvFFI constList
+  let ixonEnv := rawEnv.toEnv
+  let ixonConsts := ixonEnv.consts.valuesIter
+  let (ioBuffer, n) := ixonConsts.fold (init := (default, 0)) fun (ioBuffer, i) c =>
+    let (_, bytes) := Ixon.Serialize.put c |>.run default
+    (ioBuffer.extend #[.ofNat i] (bytes.data.map .ofUInt8), i + 1)
+
+  let _report ← oneShotBench "IxVM benchmarks"
+    (bench "serde/blake3 Nat.add_comm"
+      (aiurSystem.prove friParameters funIdx #[.ofNat n])
+      ioBuffer)
+    { oneShot := true }
+  return

Ix/Aiur/Meta.lean

@@ -97,6 +97,7 @@ declare_syntax_cat                                              trm
 syntax ("." noWs)? ident                                      : trm
 -- syntax "cast" "(" trm ", " typ ")"                            : trm
 syntax num                                                    : trm
+syntax "(" ")"                                                : trm
 syntax "(" trm (", " trm)* ")"                                : trm
 syntax "[" trm (", " trm)* "]"                                : trm
 syntax "[" trm "; " num "]"                                   : trm
@@ -151,6 +152,7 @@ partial def elabTrm : ElabStxCat `trm
   | `(trm| $n:num) => do
     let data ← mkAppM ``Data.field #[← elabG n]
     mkAppM ``Term.data #[data]
+  | `(trm| ()) => pure $ mkConst ``Term.unit
   | `(trm| ($t:trm $[, $ts:trm]*)) => do
     if ts.isEmpty then elabTrm t
     else

Ix/Aiur/Protocol.lean

@@ -43,8 +43,18 @@ structure IOBuffer where
   map : Std.HashMap (Array G) IOKeyInfo
   deriving Inhabited
 
+def IOBuffer.extend (ioBuffer : IOBuffer) (key data : Array G) : IOBuffer :=
+  let idx := ioBuffer.data.size
+  let len := data.size
+  { ioBuffer with
+    data := ioBuffer.data ++ data
+    map := ioBuffer.map.insert key { idx, len } }
+
 instance : BEq IOBuffer where
-  beq x y := x.data == y.data && x.map.toArray == y.map.toArray
+  beq x y :=
+    x.data == y.data &&
+    x.map.size == y.map.size &&
+    x.map.all fun k v => y.map.get? k == some v
 
 namespace AiurSystem
 

Ix/Common.lean

@@ -307,6 +307,62 @@ def runFrontend (input : String) (filePath : FilePath) : IO Environment := do
       (← msgs.toList.mapM (·.toString)).map (String.trimAscii · |>.toString)
   else return s.commandState.env
 
+abbrev ConstList := List (Lean.Name × Lean.ConstantInfo)
+private abbrev CollectM := StateM Lean.NameHashSet
+
+private partial def collectDependenciesAux (const : Lean.ConstantInfo)
+    (consts : Lean.ConstMap) (acc : ConstList) : CollectM ConstList := do
+  modify (·.insert const.name)
+  match const with
+  | .ctorInfo val =>
+    let acc ← collectNames [val.induct] acc
+    goExpr consts acc val.type
+  | .axiomInfo val | .quotInfo val => goExpr consts acc val.type
+  | .inductInfo val =>
+    let acc ← collectNames val.all acc
+    let acc ← collectNames val.ctors acc
+    goExpr consts acc val.type
+  | .defnInfo val | .thmInfo val | .opaqueInfo val =>
+    let acc ← collectNames val.all acc
+    let acc ← goExpr consts acc val.type
+    goExpr consts acc val.value
+  | .recInfo val =>
+    let acc ← collectNames val.all acc
+    let acc ← goExpr consts acc val.type
+    val.rules.foldlM (init := acc) fun acc rule => goExpr consts acc rule.rhs
+where
+  collectNames all acc := do
+    let visited ← get
+    all.foldlM (init := acc) fun acc name =>
+      if visited.contains name then pure acc
+      else
+        let const := consts.find! name
+        collectDependenciesAux const consts $ (name, const) :: acc
+  goExpr (consts : Lean.ConstMap) (acc : ConstList) : Lean.Expr → CollectM ConstList
+    | .bvar _ | .fvar _ | .mvar _ | .sort _ | .lit _ => pure acc
+    | .const name _ => do
+      let visited ← get
+      if visited.contains name then pure acc
+      else
+        let const := consts.find! name
+        collectDependenciesAux const consts $ (name, const) :: acc
+    | .app f a => do
+      let acc ← goExpr consts acc f
+      goExpr consts acc a
+    | .lam _ t b _ | .forallE _ t b _ => do
+      let acc ← goExpr consts acc t
+      goExpr consts acc b
+    | .letE _ t v b _ => do
+      let acc ← goExpr consts acc t
+      let acc ← goExpr consts acc v
+      goExpr consts acc b
+    | .mdata _ e | .proj _ _ e => goExpr consts acc e
+
+def collectDependencies (name : Lean.Name) (consts : Lean.ConstMap) : ConstList :=
+  let const := consts.find! name
+  let (constList, _) := collectDependenciesAux const consts [(name, const)] default
+  constList
+
 end Lean
 
 /-- Format a duration in milliseconds with appropriate unit suffix.

Ix/Environment.lean

@@ -110,6 +110,11 @@ partial def toStringAux : Name → String
 instance : ToString Name where
   toString := toStringAux
 
+def fromLeanName : Lean.Name → Name
+  | .anonymous => .mkAnon
+  | .str pre s => .mkStr (.fromLeanName pre) s
+  | .num pre n => .mkNat (.fromLeanName pre) n
+
 end Name
 
 /-- Compare Ix.Name by hash for ordered collections. -/

Ix/IxVM.lean

@@ -13,26 +13,37 @@ namespace IxVM
 def entrypoints := ⟦
   /- # Test entrypoints -/
 
-  -- fn ixon_blake3_test(h: [[G; 4]; 8]) {
-  --   let key = [
-  --     h[0][0], h[0][1], h[0][2], h[0][3],
-  --     h[1][0], h[1][1], h[1][2], h[1][3],
-  --     h[2][0], h[2][1], h[2][2], h[2][3],
-  --     h[3][0], h[3][1], h[3][2], h[3][3],
-  --     h[4][0], h[4][1], h[4][2], h[4][3],
-  --     h[5][0], h[5][1], h[5][2], h[5][3],
-  --     h[6][0], h[6][1], h[6][2], h[6][3],
-  --     h[7][0], h[7][1], h[7][2], h[7][3]
-  --   ];
-  --   let (idx, len) = io_get_info(key);
-  --   let bytes_unconstrained = read_byte_stream(idx, len);
-  --   let ixon_unconstrained = deserialize(bytes_unconstrained);
-  --   let bytes = serialize(ixon_unconstrained);
-  --   let bytes_hash = blake3(bytes);
-  --   assert_eq!(h, bytes_hash);
-  -- }
+  fn ixon_serde_test(n: G) {
+    match n {
+      0 => (),
+      _ =>
+        let n_minus_1 = n - 1;
+        let (idx, len) = io_get_info([n_minus_1]);
+        let bytes = read_byte_stream(idx, len);
+        let (const, rest) = get_constant(bytes);
+        assert_eq!(rest, ByteStream.Nil);
+        let bytes2 = put_constant(const, ByteStream.Nil);
+        assert_eq!(bytes, bytes2);
+        ixon_serde_test(n_minus_1),
+    }
+  }
 
   /- # Benchmark entrypoints -/
+
+  fn ixon_serde_blake3_bench(n: G) {
+    match n {
+      0 => (),
+      _ =>
+        let n_minus_1 = n - 1;
+        let (idx, len) = io_get_info([n_minus_1]);
+        let bytes = read_byte_stream(idx, len);
+        let (const, rest) = get_constant(bytes);
+        assert_eq!(rest, ByteStream.Nil);
+        let bytes2 = put_constant(const, ByteStream.Nil);
+        assert_eq!(blake3(bytes), blake3(bytes2));
+        ixon_serde_blake3_bench(n_minus_1),
+    }
+  }
 ⟧
 
 def ixVM : Except Aiur.Global Aiur.Toplevel := do

Ix/IxVM/ByteStream.lean

@@ -45,7 +45,9 @@ def byteStream := ⟦
     }
   }
 
-  -- Count bytes needed to represent a u64 (0-8)
+  -- Count bytes needed to represent a u64.
+  -- Important: this implementation differs from the Lean and Rust ones, returning
+  -- 1 for [0; 8] instead of 0.
   fn u64_byte_count(x: [G; 8]) -> G {
     match x {
       [_, 0, 0, 0, 0, 0, 0, 0] => 1,
@@ -189,6 +191,37 @@ def byteStream := ⟦
     Nil
   }
 
+  -- Computes the predecessor of an `u64` assumed to be properly represented in
+  -- little-endian bytes. If that's not the case, this implementation has UB.
+  fn relaxed_u64_pred(bytes: [G; 8]) -> [G; 8] {
+    let [b0, b1, b2, b3, b4, b5, b6, b7] = bytes;
+    match b0 {
+      0 => match b1 {
+        0 => match b2 {
+          0 => match b3 {
+            0 => match b4 {
+              0 => match b5 {
+                0 => match b6 {
+                  0 => match b7 {
+                    0 => [0, 0, 0, 0, 0, 0, 0, 0],
+                    _ => [255, 255, 255, 255, 255, 255, 255, b7 - 1],
+                  },
+                  _ => [255, 255, 255, 255, 255, 255, b6 - 1, b7],
+                },
+                _ => [255, 255, 255, 255, 255, b5 - 1, b6, b7],
+              },
+              _ => [255, 255, 255, 255, b4 - 1, b5, b6, b7],
+            },
+            _ => [255, 255, 255, b3 - 1, b4, b5, b6, b7],
+          },
+          _ => [255, 255, b2 - 1, b3, b4, b5, b6, b7],
+        },
+        _ => [255, b1 - 1, b2, b3, b4, b5, b6, b7],
+      },
+      _ => [b0 - 1, b1, b2, b3, b4, b5, b6, b7],
+    }
+  }
+
   fn u64_list_length(xs: U64List) -> [G; 8] {
     match xs {
       U64List.Nil => [0; 8],