Read bytecodes from Lean in Rust

Author: arthurpaulino

Ix/Aiur2/Bytecode.lean

@@ -2,4 +2,42 @@ namespace Aiur
 
 abbrev G := Fin 0xFFFFFFFF00000001
 
+namespace Bytecode
+
+abbrev FunIdx := Nat
+abbrev ValIdx := Nat
+abbrev SelIdx := Nat
+
+inductive Op
+  | const : G → Op
+  | add : ValIdx → ValIdx → Op
+  | sub : ValIdx → ValIdx → Op
+  | mul : ValIdx → ValIdx → Op
+  | call : FunIdx → Array ValIdx → Op
+  | store : Array ValIdx → Op
+  | load : (width : Nat) → ValIdx → Op
+
+mutual
+  inductive Ctrl where
+    | match : ValIdx → Array (G × Block) → Option Block → Ctrl
+    | return : SelIdx → Array ValIdx → Ctrl
+
+  structure Block where
+    ops : Array Op
+    ctrl : Ctrl
+    minSelIncluded: SelIdx
+    maxSelExcluded: SelIdx
+end
+
+structure Function where
+  inputSize : Nat
+  outputSize : Nat
+  body : Block
+
+structure Toplevel where
+  functions : Array Function
+  memoryWidths : Array Nat
+
+end Bytecode
+
 end Aiur

src/aiur2/bytecode.rs

@@ -6,27 +6,29 @@ use p3_goldilocks::Goldilocks as G;
 use rustc_hash::FxBuildHasher;
 
 pub struct Toplevel {
-    pub(super) functions: Vec<Function>,
-    pub(super) memory_widths: Vec<usize>,
+    pub(crate) functions: Vec<Function>,
+    pub(crate) memory_widths: Vec<usize>,
 }
 
 pub struct Function {
-    input_size: usize,
-    output_size: usize,
-    pub(super) body: Block,
+    pub(crate) input_size: usize,
+    pub(crate) output_size: usize,
+    pub(crate) body: Block,
 }
 
 pub type FxIndexMap<K, V> = IndexMap<K, V, FxBuildHasher>;
 
 pub struct Block {
-    pub(super) ops: Vec<Op>,
-    pub(super) ctrl: Ctrl,
-    sel_range: SelRange,
+    pub(crate) ops: Vec<Op>,
+    pub(crate) ctrl: Ctrl,
+    pub(crate) min_sel_included: SelIdx,
+    pub(crate) max_sel_excluded: SelIdx,
 }
 
 pub enum Op {
     Const(G),
     Add(ValIdx, ValIdx),
+    Sub(ValIdx, ValIdx),
     Mul(ValIdx, ValIdx),
     Call(FunIdx, Vec<ValIdx>),
     Store(Vec<ValIdx>),
@@ -38,13 +40,8 @@ pub enum Ctrl {
     Return(SelIdx, Vec<ValIdx>),
 }
 
-pub struct SelRange {
-    min_included: SelIdx,
-    max_excluded: SelIdx,
-}
-
 #[derive(Clone, Copy)]
-pub struct SelIdx(usize);
+pub struct SelIdx(pub(crate) usize);
 
 impl SelIdx {
     #[inline]
@@ -54,7 +51,7 @@ impl SelIdx {
 }
 
 #[derive(Clone, Copy)]
-pub struct ValIdx(usize);
+pub struct ValIdx(pub(crate) usize);
 
 impl ValIdx {
     #[inline]
@@ -64,7 +61,7 @@ impl ValIdx {
 }
 
 #[derive(Clone, Copy)]
-pub struct FunIdx(usize);
+pub struct FunIdx(pub(crate) usize);
 
 impl FunIdx {
     #[inline]

src/aiur2/execute.rs

@@ -5,7 +5,7 @@ use crate::aiur2::bytecode::{Ctrl, FunIdx, Function, FxIndexMap, Op, Toplevel};
 
 pub struct QueryResult {
     output: Vec<G>,
-    multiplicity: usize,
+    multiplicity: G,
 }
 
 pub type QueryMap = FxIndexMap<Vec<G>, QueryResult>;
@@ -78,6 +78,11 @@ impl Function {
                     let b = map[b.to_usize()];
                     map.push(a + b);
                 }
+                ExecEntry::Op(Op::Sub(a, b)) => {
+                    let a = map[a.to_usize()];
+                    let b = map[b.to_usize()];
+                    map.push(a - b);
+                }
                 ExecEntry::Op(Op::Mul(a, b)) => {
                     let a = map[a.to_usize()];
                     let b = map[b.to_usize()];
@@ -88,7 +93,7 @@ impl Function {
                     if let Some(result) =
                         record.function_queries[callee_idx.to_usize()].get_mut(&args)
                     {
-                        result.multiplicity += 1;
+                        result.multiplicity += G::ONE;
                         map.extend(result.output.clone());
                     } else {
                         let saved_map = std::mem::replace(&mut map, args);
@@ -111,13 +116,13 @@ impl Function {
                         .get_mut(&width)
                         .expect("Invalid memory width");
                     if let Some(result) = memory_queries.get_mut(&values) {
-                        result.multiplicity += 1;
+                        result.multiplicity += G::ONE;
                         map.extend(&result.output);
                     } else {
                         let ptr = G::ZERO;
                         let result = QueryResult {
                             output: vec![ptr],
-                            multiplicity: 1,
+                            multiplicity: G::ONE,
                         };
                         memory_queries.insert(values, result);
                         map.push(ptr);
@@ -134,7 +139,7 @@ impl Function {
                     let (args, result) = memory_queries
                         .get_index_mut(ptr_usize)
                         .expect("Unbound pointer");
-                    result.multiplicity += 1;
+                    result.multiplicity += G::ONE;
                     map.extend(args);
                 }
                 ExecEntry::Ctrl(Ctrl::Match(val_idx, cases, default)) => {

src/lean/ffi/aiur2/mod.rs

@@ -0,0 +1 @@
+pub mod toplevel;

src/lean/ffi/aiur2/toplevel.rs

@@ -0,0 +1,161 @@
+// TODO: remove
+#![allow(dead_code)]
+
+use p3_field::integers::QuotientMap;
+use p3_goldilocks::Goldilocks as G;
+use std::ffi::c_void;
+
+use crate::{
+    aiur2::bytecode::{Block, Ctrl, FunIdx, Function, FxIndexMap, Op, SelIdx, Toplevel, ValIdx},
+    lean::{
+        array::LeanArrayObject,
+        ctor::LeanCtorObject,
+        ffi::{as_ref_unsafe, lean_is_scalar},
+    },
+    lean_unbox,
+};
+
+fn lean_unbox_nat_as_usize(ptr: *const c_void) -> usize {
+    assert!(lean_is_scalar(ptr));
+    lean_unbox!(usize, ptr)
+}
+
+fn lean_unbox_nat_as_g(ptr: *const c_void) -> G {
+    assert!(lean_is_scalar(ptr));
+    unsafe { G::from_canonical_unchecked(lean_unbox!(u64, ptr)) }
+}
+
+fn lean_ptr_to_vec_val_idx(ptr: *const c_void) -> Vec<ValIdx> {
+    let array: &LeanArrayObject = as_ref_unsafe(ptr.cast());
+    array.to_vec(|ptr| ValIdx(lean_unbox_nat_as_usize(ptr)))
+}
+
+fn lean_ptr_to_op(ptr: *const c_void) -> Op {
+    let ctor: &LeanCtorObject = as_ref_unsafe(ptr.cast());
+    match ctor.tag() {
+        0 => {
+            let [const_val_ptr] = ctor.objs();
+            Op::Const(lean_unbox_nat_as_g(const_val_ptr))
+        }
+        1 => {
+            let [a_ptr, b_ptr] = ctor.objs();
+            Op::Add(
+                ValIdx(lean_unbox_nat_as_usize(a_ptr)),
+                ValIdx(lean_unbox_nat_as_usize(b_ptr)),
+            )
+        }
+        2 => {
+            let [a_ptr, b_ptr] = ctor.objs();
+            Op::Sub(
+                ValIdx(lean_unbox_nat_as_usize(a_ptr)),
+                ValIdx(lean_unbox_nat_as_usize(b_ptr)),
+            )
+        }
+        3 => {
+            let [a_ptr, b_ptr] = ctor.objs();
+            Op::Mul(
+                ValIdx(lean_unbox_nat_as_usize(a_ptr)),
+                ValIdx(lean_unbox_nat_as_usize(b_ptr)),
+            )
+        }
+        4 => {
+            let [fun_idx_ptr, val_idxs_ptr] = ctor.objs();
+            let fun_idx = FunIdx(lean_unbox_nat_as_usize(fun_idx_ptr));
+            let val_idxs = lean_ptr_to_vec_val_idx(val_idxs_ptr);
+            Op::Call(fun_idx, val_idxs)
+        }
+        5 => {
+            let [val_idxs_ptr] = ctor.objs();
+            Op::Store(lean_ptr_to_vec_val_idx(val_idxs_ptr))
+        }
+        6 => {
+            let [width_ptr, val_idx_ptr] = ctor.objs();
+            Op::Load(
+                lean_unbox_nat_as_usize(width_ptr),
+                ValIdx(lean_unbox_nat_as_usize(val_idx_ptr)),
+            )
+        }
+        _ => unreachable!(),
+    }
+}
+
+fn lean_ptr_to_g_block_pair(ptr: *const c_void) -> (G, Block) {
+    let ctor: &LeanCtorObject = as_ref_unsafe(ptr.cast());
+    let [g_ptr, block_ptr] = ctor.objs();
+    let g = lean_unbox_nat_as_g(g_ptr);
+    let block = lean_ctor_to_block(as_ref_unsafe(block_ptr.cast()));
+    (g, block)
+}
+
+fn lean_ctor_to_ctrl(ctor: &LeanCtorObject) -> Ctrl {
+    match ctor.tag() {
+        0 => {
+            let [val_idx_ptr, cases_ptr, default_ptr] = ctor.objs();
+            let val_idx = ValIdx(lean_unbox_nat_as_usize(val_idx_ptr));
+            let cases_array: &LeanArrayObject = as_ref_unsafe(cases_ptr.cast());
+            let vec_cases = cases_array.to_vec(lean_ptr_to_g_block_pair);
+            let cases = FxIndexMap::from_iter(vec_cases);
+            let default = if lean_is_scalar(default_ptr) {
+                None
+            } else {
+                let default_ctor: &LeanCtorObject = as_ref_unsafe(default_ptr.cast());
+                let [block_ptr] = default_ctor.objs();
+                let block = lean_ctor_to_block(as_ref_unsafe(block_ptr.cast()));
+                Some(Box::new(block))
+            };
+            Ctrl::Match(val_idx, cases, default)
+        }
+        1 => {
+            let [sel_idx_ptr, val_idxs_ptr] = ctor.objs();
+            let sel_idx = SelIdx(lean_unbox_nat_as_usize(sel_idx_ptr));
+            let val_idxs = lean_ptr_to_vec_val_idx(val_idxs_ptr);
+            Ctrl::Return(sel_idx, val_idxs)
+        }
+        _ => unreachable!(),
+    }
+}
+
+fn lean_ctor_to_block(ctor: &LeanCtorObject) -> Block {
+    let [
+        ops_ptr,
+        ctrl_ptr,
+        min_sel_included_ptr,
+        max_sel_excluded_ptr,
+    ] = ctor.objs();
+    let ops_array: &LeanArrayObject = as_ref_unsafe(ops_ptr.cast());
+    let ops = ops_array.to_vec(lean_ptr_to_op);
+    let ctrl = lean_ctor_to_ctrl(as_ref_unsafe(ctrl_ptr.cast()));
+    let min_sel_included = SelIdx(lean_unbox_nat_as_usize(min_sel_included_ptr));
+    let max_sel_excluded = SelIdx(lean_unbox_nat_as_usize(max_sel_excluded_ptr));
+    Block {
+        ops,
+        ctrl,
+        min_sel_included,
+        max_sel_excluded,
+    }
+}
+
+fn lean_ptr_to_function(ptr: *const c_void) -> Function {
+    let ctor: &LeanCtorObject = as_ref_unsafe(ptr.cast());
+    let [input_size_ptr, output_size_ptr, body_ptr] = ctor.objs();
+    let input_size = lean_unbox_nat_as_usize(input_size_ptr);
+    let output_size = lean_unbox_nat_as_usize(output_size_ptr);
+    let body = lean_ctor_to_block(as_ref_unsafe(body_ptr.cast()));
+    Function {
+        input_size,
+        output_size,
+        body,
+    }
+}
+
+fn lean_ctor_to_toplevel(ctor: &LeanCtorObject) -> Toplevel {
+    let [functions_ptr, memory_widths_ptr] = ctor.objs();
+    let functions_array: &LeanArrayObject = as_ref_unsafe(functions_ptr.cast());
+    let functions = functions_array.to_vec(lean_ptr_to_function);
+    let memory_widths_array: &LeanArrayObject = as_ref_unsafe(memory_widths_ptr.cast());
+    let memory_widths = memory_widths_array.to_vec(lean_unbox_nat_as_usize);
+    Toplevel {
+        functions,
+        memory_widths,
+    }
+}

src/lean/ffi/mod.rs

@@ -1,3 +1,4 @@
+pub mod aiur2;
 pub mod archon;
 pub mod binius;
 pub mod byte_array;