[wip] Trenchcoat Loam <> Lair integration

src/lair/air.rs

@@ -16,7 +16,7 @@ use super::{
     chipset::Chipset,
     func_chip::{ColumnLayout, FuncChip, LayoutSizes},
     provenance::{DepthLessThan, DEPTH_LESS_THAN_SIZE, DEPTH_W},
-    relations::{CallRelation, MemoryRelation},
+    relations::{CallRelation, LoamRelation, MemoryRelation},
     toplevel::Toplevel,
     trace::ColumnIndex,
 };
@@ -87,6 +87,18 @@ impl<'a, T> ColumnSlice<'a, T> {
         slice
     }
 
+    pub fn next_provide(&self, index: &mut ColumnIndex) -> ProvideRecord<T>
+    where
+        T: Copy,
+    {
+        let last_nonce = self.next_aux(index);
+        let last_count = self.next_aux(index);
+        ProvideRecord {
+            last_nonce,
+            last_count,
+        }
+    }
+
     pub fn next_require(&self, index: &mut ColumnIndex) -> RequireRecord<T>
     where
         T: Copy,
@@ -197,17 +209,12 @@ impl<F: Field> Func<F> {
 
         let toplevel_sel = self.body.return_sel::<AB>(local);
         builder.assert_bool(toplevel_sel.clone());
-        let last_nonce = local.next_aux(index);
-        let last_count = local.next_aux(index);
-        // provenance and range check
-        let record = ProvideRecord {
-            last_nonce,
-            last_count,
-        };
+        let record = local.next_provide(index);
         let mut out = (0..self.output_size)
             .map(|i| local.output[i].into())
             .collect::<Vec<_>>();
         let depth = if self.partial {
+            // provenance and range check
             let depth = (0..DEPTH_W)
                 .map(|_| local.next_aux(index))
                 .collect::<Vec<_>>();
@@ -428,6 +435,27 @@ impl<F: Field> Op<F> {
                     sel.clone(),
                 );
             }
+            Op::Provide(rel_idx, args) => {
+                let args = args.iter().map(|i| map[*i].to_expr());
+                let record = local.next_provide(index);
+
+                builder.provide(
+                    LoamRelation(F::from_canonical_usize(*rel_idx), args),
+                    record,
+                    sel.clone(),
+                );
+            }
+            Op::Require(rel_idx, args) => {
+                let args = args.iter().map(|i| map[*i].to_expr());
+                let record = local.next_require(index);
+
+                builder.require(
+                    LoamRelation(F::from_canonical_usize(*rel_idx), args),
+                    *local.nonce,
+                    record,
+                    sel.clone(),
+                );
+            }
             Op::Store(values) => {
                 let ptr = local.next_aux(index);
                 map.push(Val::Expr(ptr.into()));
@@ -524,6 +552,7 @@ impl<F: Field> Ctrl<F> {
                     process(block)
                 };
             }
+            Ctrl::Exit => todo!(),
             Ctrl::ChooseMany(_, cases) => {
                 let map_len = map.len();
                 let init_state = index.save();
@@ -975,4 +1004,15 @@ mod tests {
             .contains_key(&vec![f(3), f(1), f(100), f(12)]));
         assert!(queries.memoset.contains_key(&vec![f(3), f(1), f(64), f(0)]));
     }
+
+    #[test]
+    fn lair_provide_test() {
+        let func_provide = func!(
+        loam fn provide_test(x: [3]): [0] {
+            provide!(f, x);
+            return ()
+        });
+        println!("{:#?}", func_provide);
+        let toplevel = Toplevel::<F, Nochip>::new_pure(&[func_provide]);
+    }
 }

src/lair/bytecode.rs

@@ -40,6 +40,10 @@ pub enum Op<F> {
     /// use it to format an error message and bail at runtime instead of panicking
     /// when the preimage is not available
     PreImg(usize, List<usize>, Option<fn(&[F]) -> String>),
+    /// TODO: Loam-only, document
+    Provide(usize, List<usize>),
+    /// TODO: Loam-only, document
+    Require(usize, List<usize>),
     /// `Store([y, ...])` pushes the pointer to `[y, ...]` to the stack
     Store(List<usize>),
     /// `Load(len, y)` extends the stack with the `len` values that is pointed by `y`
@@ -79,6 +83,8 @@ pub enum Ctrl<F> {
     /// block. The first `usize` is an unique identifier, representing the
     /// selector used for arithmetization
     Return(usize, List<usize>),
+    /// TODO: Loam-only, document
+    Exit,
 }
 
 /// Represents the cases for `Ctrl::Match`, containing the branches for successfull
@@ -139,6 +145,7 @@ pub struct Func<F> {
     pub(crate) name: Name,
     pub(crate) invertible: bool,
     pub(crate) partial: bool,
+    pub(crate) loam: bool,
     pub(crate) index: usize,
     pub(crate) input_size: usize,
     pub(crate) output_size: usize,

src/lair/execute.rs

@@ -26,7 +26,7 @@ pub(crate) type MemMap<F> = FxIndexMap<List<F>, QueryResult<F>>;
 #[derive(Clone, Debug, Eq, PartialEq, Default)]
 pub struct QueryResult<F> {
     pub(crate) output: Option<List<F>>,
-    pub(crate) provide: Record,
+    pub(crate) provides: Vec<Record>,
     pub(crate) requires: Vec<Record>,
     pub(crate) depth: u32,
     pub(crate) depth_requires: Vec<Record>,
@@ -39,7 +39,7 @@ impl<F: PrimeField32> QueryResult<F> {
     }
 
     pub(crate) fn new_lookup(&mut self, nonce: usize, caller_requires: &mut Vec<Record>) {
-        caller_requires.push(self.provide.new_lookup(nonce as u32));
+        caller_requires.push(self.provides[0].new_lookup(nonce as u32));
     }
 }
 
@@ -63,7 +63,7 @@ pub struct DebugData {
     pub(crate) breakpoints: Vec<usize>,
 }
 
-#[derive(Debug, Clone, Eq, PartialEq)]
+#[derive(Default, Debug, Clone, Eq, PartialEq)]
 pub struct QueryRecord<F: PrimeField32> {
     pub(crate) public_values: Option<Vec<F>>,
     pub(crate) func_queries: Vec<QueryMap<F>>,
@@ -155,7 +155,11 @@ impl<'a, F: PrimeField32> MachineRecord for Shard<'a, F> {
             queries
                 .func_queries
                 .iter()
-                .map(|im| im.values().map(|r| r.provide.count as usize).sum::<usize>())
+                .map(|im| {
+                    im.values()
+                        .map(|r| r.provides[0].count as usize)
+                        .sum::<usize>()
+                })
                 .sum(),
         );
 
@@ -169,7 +173,11 @@ impl<'a, F: PrimeField32> MachineRecord for Shard<'a, F> {
             queries
                 .mem_queries
                 .iter()
-                .map(|im| im.values().map(|r| r.provide.count as usize).sum::<usize>())
+                .map(|im| {
+                    im.values()
+                        .map(|r| r.provides[0].count as usize)
+                        .sum::<usize>()
+                })
                 .sum(),
         );
         map.insert(
@@ -442,7 +450,8 @@ impl<F: PrimeField32> Func<F> {
     ) -> Result<(List<F>, u32)> {
         let mut func_index = self.index;
         let mut query_result = QueryResult::default();
-        query_result.provide.count = 1;
+        query_result.provides.push(Record::default());
+        query_result.provides[0].count = 1;
         let (mut nonce, _) =
             queries.func_queries[func_index].insert_full(args.into(), query_result);
         let mut map = args.to_vec();
@@ -628,6 +637,9 @@ impl<F: PrimeField32> Func<F> {
                         push_block_exec_entries!(&func.body);
                     }
                 }
+                ExecEntry::Op(Op::Provide(..)) | ExecEntry::Op(Op::Require(..)) => {
+                    panic!("Provide and require cannot yet be run on Lair")
+                }
                 ExecEntry::Op(Op::Const(c)) => map.push(*c),
                 ExecEntry::Op(Op::Add(a, b)) => map.push(map[*a] + map[*b]),
                 ExecEntry::Op(Op::Sub(a, b)) => map.push(map[*a] - map[*b]),
@@ -766,6 +778,7 @@ impl<F: PrimeField32> Func<F> {
                         break;
                     }
                 }
+                ExecEntry::Ctrl(Ctrl::Exit) => todo!(),
                 ExecEntry::Ctrl(Ctrl::Choose(v, cases, _)) => {
                     let v = map[*v];
                     let block = cases.match_case(&v).expect("No match");

src/lair/expr.rs

@@ -81,6 +81,10 @@ pub enum OpE<F> {
     /// If `fmt` is `Some`, use it to format an error message and bail at runtime
     /// instead of panicking when the preimage is not available
     PreImg(VarList, Name, VarList, Option<fn(&[F]) -> String>),
+    /// TODO: Loam-only, document
+    Provide(Name, VarList),
+    /// TODO: Loam-only, document
+    Require(Name, VarList),
     /// `Store(x, [y, ...])` binds `x` to a pointer to `[y, ...]`
     Store(Var, VarList),
     /// `Load([x, ...], y)` binds `[x, ...]` to the values that is pointed by `y`
@@ -122,6 +126,8 @@ pub enum CtrlE<F> {
     /// Contains the variables whose bindings will construct the output of the
     /// block
     Return(VarList),
+    /// TODO: Loam-only, document
+    Exit,
 }
 
 /// Represents the cases for `CtrlE::Match`, containing the branches for successfull
@@ -140,6 +146,7 @@ pub struct FuncE<F> {
     pub name: Name,
     pub invertible: bool,
     pub partial: bool,
+    pub loam: bool,
     pub input_params: VarList,
     pub output_size: usize,
     pub body: BlockE<F>,

src/lair/func_chip.rs

@@ -141,6 +141,7 @@ impl<F> Ctrl<F> {
                 // exactly one selector per return
                 *sel += 1;
             }
+            Ctrl::Exit => todo!(),
             Ctrl::Choose(_, cases, branches) => {
                 let degrees_len = degrees.len();
                 let mut max_aux = *aux;
@@ -249,6 +250,8 @@ impl<F> Op<F> {
                 }
                 degrees.extend(vec![1; inp_size]);
             }
+            Op::Provide(_rel_idx, _args) => todo!(),
+            Op::Require(_rel_idx, _args) => todo!(),
             Op::Store(..) => {
                 *aux += 4;
                 degrees.push(1);

src/lair/macros.rs

@@ -6,6 +6,19 @@ macro_rules! func {
             name: $crate::lair::Name(stringify!($name)),
             invertible: false,
             partial: false,
+            loam: false,
+            input_params: [$($crate::var!($in $(, $in_size)?)),*].into(),
+            output_size: $size,
+            body: $crate::block_init!($lair),
+        }
+    }};
+    (loam fn $name:ident($( $in:ident $(: [$in_size:expr])? ),*): [$size:expr] $lair:tt) => {{
+        $(let $in = $crate::var!($in $(, $in_size)?);)*
+        $crate::lair::expr::FuncE {
+            name: $crate::lair::Name(stringify!($name)),
+            invertible: false,
+            partial: false,
+            loam: true,
             input_params: [$($crate::var!($in $(, $in_size)?)),*].into(),
             output_size: $size,
             body: $crate::block_init!($lair),
@@ -17,6 +30,7 @@ macro_rules! func {
             name: $crate::lair::Name(stringify!($name)),
             invertible: false,
             partial: true,
+            loam: false,
             input_params: [$($crate::var!($in $(, $in_size)?)),*].into(),
             output_size: $size,
             body: $crate::block_init!($lair),
@@ -28,6 +42,7 @@ macro_rules! func {
             name: $crate::lair::Name(stringify!($name)),
             invertible: true,
             partial: false,
+            loam: false,
             input_params: [$($crate::var!($in $(, $in_size)?)),*].into(),
             output_size: $size,
             body: $crate::block_init!($lair),
@@ -39,6 +54,7 @@ macro_rules! func {
             name: $crate::lair::Name(stringify!($name)),
             invertible: true,
             partial: true,
+            loam: false,
             input_params: [$($crate::var!($in $(, $in_size)?)),*].into(),
             output_size: $size,
             body: $crate::block_init!($lair),
@@ -74,6 +90,18 @@ macro_rules! block_init {
 #[macro_export]
 macro_rules! block {
     // Operations
+    ({ provide!($func:ident, $($arg:ident),*); $($tail:tt)+ }, $ops:expr) => {{
+        let rel = $crate::lair::Name(stringify!($func));
+        let args = [$($arg),*].into();
+        $ops.push($crate::lair::expr::OpE::Provide(rel, args));
+        $crate::block!({ $($tail)* }, $ops)
+    }};
+    ({ require!($func:ident, $($arg:ident),*); $($tail:tt)+ }, $ops:expr) => {{
+        let rel = $crate::lair::Name(stringify!($func));
+        let args = [$($arg),*].into();
+        $ops.push($crate::lair::expr::OpE::Require(rel, args));
+        $crate::block!({ $($tail)* }, $ops)
+    }};
     ({ range_u8!($($a:ident),*); $($tail:tt)+ }, $ops:expr) => {{
         $ops.push($crate::lair::expr::OpE::RangeU8([$($a),*].into()));
         $crate::block!({ $($tail)* }, $ops)

src/lair/map.rs

@@ -69,7 +69,7 @@ impl<K, V> Map<K, V> {
         self.0
     }
 
-    /// Retrieves the key/value pair by intex
+    /// Retrieves the key/value pair by index
     #[inline]
     pub fn get_index(&self, i: usize) -> Option<&(K, V)> {
         self.0.get(i)

src/lair/memory.rs

@@ -50,7 +50,7 @@ impl<F: PrimeField32> MemChip<F> {
             // .skip(range.start)
             // .take(non_dummy_height)
             .for_each(|(i, (row, (args, mem_result)))| {
-                let provide = mem_result.provide.into_provide();
+                let provide = mem_result.provides[0].into_provide();
 
                 // is_real
                 row[0] = F::one();