From 867f0c2a1a24980bd924d01d2f81d6ce01fe72b3 Mon Sep 17 00:00:00 2001
From: Arthur Paulino <arthurleonardo.ap@gmail.com>
Date: Thu, 29 Feb 2024 13:25:29 -0300
Subject: [PATCH] chore: refine Coprocessors' API (#1180)

* Allow coprocessor to alloc their own global constants
* Remove Coprocessor::eval_arity since the implementer already defines CoCircuit::arity
---
 lurk-macros/src/lib.rs        | 19 -------
 src/coprocessor/circom.rs     | 14 +++--
 src/coprocessor/mod.rs        | 99 ++++++++++++++++++++++++++++++-----
 src/coprocessor/sha256.rs     |  4 --
 src/coprocessor/trie/mod.rs   | 30 ++++++-----
 src/eval/lang.rs              | 11 ++--
 src/lem/circuit.rs            | 30 +++++++----
 src/lem/multiframe.rs         | 13 ++---
 src/proof/tests/nova_tests.rs | 31 +++++++++--
 9 files changed, 171 insertions(+), 80 deletions(-)

diff --git a/lurk-macros/src/lib.rs b/lurk-macros/src/lib.rs
index 1bf1cadb22..6fe41f70a2 100644
--- a/lurk-macros/src/lib.rs
+++ b/lurk-macros/src/lib.rs
@@ -31,7 +31,6 @@ pub fn derive_enum_coproc(input: TokenStream) -> TokenStream {
 }
 
 fn impl_enum_coproc(name: &Ident, variants: &DataEnum) -> TokenStream {
-    let eval_arity_arms = eval_arity_match_arms(name, variants);
     let evaluate_internal_arms = evaluate_internal_match_arms(name, variants);
     let evaluate_arms = evaluate_match_arms(name, variants);
     let evaluate_simple_arms = evaluate_simple_match_arms(name, variants);
@@ -46,12 +45,6 @@ fn impl_enum_coproc(name: &Ident, variants: &DataEnum) -> TokenStream {
 
     let res = quote! {
         impl <F: lurk::field::LurkField> lurk::coprocessor::Coprocessor<F> for #name<F> {
-            fn eval_arity(&self) -> usize {
-                match self {
-                    #eval_arity_arms
-                }
-            }
-
             fn evaluate_internal(&self, s: &lurk::lem::store::Store<F>, ptrs: &[lurk::lem::pointers::Ptr]) -> Vec<lurk::lem::pointers::Ptr> {
                 match self {
                     #evaluate_internal_arms
@@ -131,18 +124,6 @@ fn impl_enum_coproc(name: &Ident, variants: &DataEnum) -> TokenStream {
     res.into()
 }
 
-fn eval_arity_match_arms(name: &Ident, variants: &DataEnum) -> proc_macro2::TokenStream {
-    let mut match_arms = quote! {};
-    for variant in variants.variants.iter() {
-        let variant_ident = &variant.ident;
-
-        match_arms.extend(quote! {
-            #name::#variant_ident(coprocessor) => coprocessor.eval_arity(),
-        });
-    }
-    match_arms
-}
-
 fn evaluate_internal_match_arms(name: &Ident, variants: &DataEnum) -> proc_macro2::TokenStream {
     let mut match_arms = quote! {};
     for variant in variants.variants.iter() {
diff --git a/src/coprocessor/circom.rs b/src/coprocessor/circom.rs
index 218647d734..e6277d6d34 100644
--- a/src/coprocessor/circom.rs
+++ b/src/coprocessor/circom.rs
@@ -123,14 +123,18 @@ Then run `lurk coprocessor --name {name} <{}_FOLDER>` to instantiate a new gadge
 
             Ok(res)
         }
-    }
 
-    impl<F: LurkField, C: CircomGadget<F> + Debug> Coprocessor<F> for CircomCoprocessor<F, C> {
-        /// TODO: Generalize
-        fn eval_arity(&self) -> usize {
-            0
+        fn alloc_globals<CS: ConstraintSystem<F>>(
+            &self,
+            cs: &mut CS,
+            g: &crate::lem::circuit::GlobalAllocator<F>,
+            _s: &Store<F>,
+        ) {
+            g.alloc_tag(cs, &crate::tag::ExprTag::Num);
         }
+    }
 
+    impl<F: LurkField, C: CircomGadget<F> + Debug> Coprocessor<F> for CircomCoprocessor<F, C> {
         fn evaluate_simple(&self, s: &Store<F>, args: &[Ptr]) -> Ptr {
             self.gadget.evaluate_simple(s, args)
         }
diff --git a/src/coprocessor/mod.rs b/src/coprocessor/mod.rs
index 2e2ab7607c..504ee8d85f 100644
--- a/src/coprocessor/mod.rs
+++ b/src/coprocessor/mod.rs
@@ -27,8 +27,6 @@ pub mod trie;
 /// - An enum such as [`crate::eval::lang::Coproc`], which "closes" the hierarchy of possible coprocessor
 ///   implementations we want to instantiate at a particular point in the code.
 pub trait Coprocessor<F: LurkField>: Clone + Debug + Sync + Send + CoCircuit<F> {
-    fn eval_arity(&self) -> usize;
-
     /// Returns true if this Coprocessor actually implements a circuit.
     fn has_circuit(&self) -> bool {
         false
@@ -57,9 +55,7 @@ pub trait Coprocessor<F: LurkField>: Clone + Debug + Sync + Send + CoCircuit<F>
 ///
 /// The trait is implemented by concrete coprocessor types, such as `DumbCoprocessor`.
 pub trait CoCircuit<F: LurkField>: Send + Sync + Clone {
-    fn arity(&self) -> usize {
-        todo!()
-    }
+    fn arity(&self) -> usize;
 
     /// Function for internal plumbing. Reimplementing is not recommended
     fn synthesize_internal<CS: ConstraintSystem<F>>(
@@ -104,6 +100,15 @@ pub trait CoCircuit<F: LurkField>: Send + Sync + Clone {
     ) -> Result<AllocatedPtr<F>, SynthesisError> {
         unimplemented!()
     }
+
+    /// Perform global allocations needed for synthesis
+    fn alloc_globals<CS: ConstraintSystem<F>>(
+        &self,
+        _cs: &mut CS,
+        _g: &GlobalAllocator<F>,
+        _s: &Store<F>,
+    ) {
+    }
 }
 
 #[cfg(test)]
@@ -165,6 +170,7 @@ pub(crate) mod test {
     }
 
     impl<F: LurkField> CoCircuit<F> for DumbCoprocessor<F> {
+        /// Dumb Coprocessor takes two arguments.
         fn arity(&self) -> usize {
             2
         }
@@ -188,14 +194,19 @@ pub(crate) mod test {
 
             Ok(vec![expr, env, cont])
         }
-    }
 
-    impl<F: LurkField> Coprocessor<F> for DumbCoprocessor<F> {
-        /// Dumb Coprocessor takes two arguments.
-        fn eval_arity(&self) -> usize {
-            2
+        fn alloc_globals<CS: ConstraintSystem<F>>(
+            &self,
+            cs: &mut CS,
+            g: &GlobalAllocator<F>,
+            s: &Store<F>,
+        ) {
+            g.alloc_tag(cs, &ExprTag::Num);
+            g.alloc_ptr(cs, &s.cont_error(), s);
         }
+    }
 
+    impl<F: LurkField> Coprocessor<F> for DumbCoprocessor<F> {
         fn has_circuit(&self) -> bool {
             true
         }
@@ -250,13 +261,19 @@ pub(crate) mod test {
             let term = g.alloc_ptr(cs, &s.cont_terminal(), s);
             Ok(vec![nil, env.clone(), term])
         }
-    }
 
-    impl<F: LurkField> Coprocessor<F> for Terminator<F> {
-        fn eval_arity(&self) -> usize {
-            0
+        fn alloc_globals<CS: ConstraintSystem<F>>(
+            &self,
+            cs: &mut CS,
+            g: &GlobalAllocator<F>,
+            s: &Store<F>,
+        ) {
+            g.alloc_ptr(cs, &s.intern_nil(), s);
+            g.alloc_ptr(cs, &s.cont_terminal(), s);
         }
+    }
 
+    impl<F: LurkField> Coprocessor<F> for Terminator<F> {
         fn evaluate(&self, s: &Store<F>, _args: &[Ptr], env: &Ptr, _cont: &Ptr) -> Vec<Ptr> {
             vec![s.intern_nil(), *env, s.cont_terminal()]
         }
@@ -277,4 +294,58 @@ pub(crate) mod test {
             }
         }
     }
+
+    #[derive(Clone, Debug, Serialize, Deserialize)]
+    pub(crate) struct HelloWorld<F> {
+        _p: PhantomData<F>,
+    }
+
+    impl<F: LurkField> HelloWorld<F> {
+        pub(crate) fn new() -> Self {
+            Self {
+                _p: Default::default(),
+            }
+        }
+
+        #[inline]
+        pub(crate) fn intern_hello_world(s: &Store<F>) -> Ptr {
+            s.intern_string("Hello, world!")
+        }
+    }
+
+    impl<F: LurkField> CoCircuit<F> for HelloWorld<F> {
+        fn arity(&self) -> usize {
+            0
+        }
+
+        fn synthesize_simple<CS: ConstraintSystem<F>>(
+            &self,
+            cs: &mut CS,
+            g: &GlobalAllocator<F>,
+            s: &Store<F>,
+            _not_dummy: &Boolean,
+            _args: &[AllocatedPtr<F>],
+        ) -> Result<AllocatedPtr<F>, SynthesisError> {
+            Ok(g.alloc_ptr(cs, &Self::intern_hello_world(s), s))
+        }
+
+        fn alloc_globals<CS: ConstraintSystem<F>>(
+            &self,
+            cs: &mut CS,
+            g: &GlobalAllocator<F>,
+            s: &Store<F>,
+        ) {
+            g.alloc_ptr(cs, &Self::intern_hello_world(s), s);
+        }
+    }
+
+    impl<F: LurkField> Coprocessor<F> for HelloWorld<F> {
+        fn has_circuit(&self) -> bool {
+            true
+        }
+
+        fn evaluate_simple(&self, s: &Store<F>, _args: &[Ptr]) -> Ptr {
+            Self::intern_hello_world(s)
+        }
+    }
 }
diff --git a/src/coprocessor/sha256.rs b/src/coprocessor/sha256.rs
index d70a0b7a87..bfcea4440b 100644
--- a/src/coprocessor/sha256.rs
+++ b/src/coprocessor/sha256.rs
@@ -106,10 +106,6 @@ impl<F: LurkField> CoCircuit<F> for Sha256Coprocessor<F> {
 }
 
 impl<F: LurkField> Coprocessor<F> for Sha256Coprocessor<F> {
-    fn eval_arity(&self) -> usize {
-        self.n
-    }
-
     fn has_circuit(&self) -> bool {
         true
     }
diff --git a/src/coprocessor/trie/mod.rs b/src/coprocessor/trie/mod.rs
index 04f03391ae..d321b292a0 100644
--- a/src/coprocessor/trie/mod.rs
+++ b/src/coprocessor/trie/mod.rs
@@ -63,10 +63,6 @@ pub struct NewCoprocessor<F> {
 }
 
 impl<F: LurkField> Coprocessor<F> for NewCoprocessor<F> {
-    fn eval_arity(&self) -> usize {
-        0
-    }
-
     fn evaluate_simple(&self, s: &Store<F>, _args: &[Ptr]) -> Ptr {
         let trie: StandardTrie<'_, F> = Trie::new(&s.poseidon_cache, &s.inverse_poseidon_cache);
         // TODO: Use a custom type.
@@ -109,10 +105,6 @@ pub struct LookupCoprocessor<F> {
 }
 
 impl<F: LurkField> Coprocessor<F> for LookupCoprocessor<F> {
-    fn eval_arity(&self) -> usize {
-        2
-    }
-
     fn evaluate_simple(&self, s: &Store<F>, args: &[Ptr]) -> Ptr {
         let root_ptr = &args[0];
         let key_ptr = &args[1];
@@ -203,6 +195,15 @@ impl<F: LurkField> CoCircuit<F> for LookupCoprocessor<F> {
             result_commitment_val,
         ))
     }
+
+    fn alloc_globals<CS: ConstraintSystem<F>>(
+        &self,
+        cs: &mut CS,
+        g: &lurk::lem::circuit::GlobalAllocator<F>,
+        _s: &Store<F>,
+    ) {
+        g.alloc_tag(cs, &ExprTag::Comm);
+    }
 }
 
 #[derive(Clone, Debug, Serialize, Default, Deserialize)]
@@ -211,10 +212,6 @@ pub struct InsertCoprocessor<F> {
 }
 
 impl<F: LurkField> Coprocessor<F> for InsertCoprocessor<F> {
-    fn eval_arity(&self) -> usize {
-        3
-    }
-
     fn evaluate_simple(&self, s: &Store<F>, args: &[Ptr]) -> Ptr {
         let root_ptr = &args[0];
         let key_ptr = &args[1];
@@ -308,6 +305,15 @@ impl<F: LurkField> CoCircuit<F> for InsertCoprocessor<F> {
         let num_tag = g.alloc_tag(cs, &ExprTag::Num);
         Ok(AllocatedPtr::from_parts(num_tag.clone(), new_root_val))
     }
+
+    fn alloc_globals<CS: ConstraintSystem<F>>(
+        &self,
+        cs: &mut CS,
+        g: &lurk::lem::circuit::GlobalAllocator<F>,
+        _s: &Store<F>,
+    ) {
+        g.alloc_tag(cs, &ExprTag::Num);
+    }
 }
 
 /// Add the `Trie`-associated functions to a `Lang` with standard bindings.
diff --git a/src/eval/lang.rs b/src/eval/lang.rs
index 0ad5f6ca65..b28a0ea996 100644
--- a/src/eval/lang.rs
+++ b/src/eval/lang.rs
@@ -24,11 +24,6 @@ pub struct DummyCoprocessor<F> {
 }
 
 impl<F: LurkField> Coprocessor<F> for DummyCoprocessor<F> {
-    /// Dummy Coprocessor takes no arguments.
-    fn eval_arity(&self) -> usize {
-        0
-    }
-
     /// And does nothing but return nil. It should probably never be used and can perhaps be eliminated,
     /// but for now it exists as an exemplar demonstrating the intended shape of enums like the default, `Coproc`.
     fn evaluate_simple(&self, s: &Store<F>, _args: &[Ptr]) -> Ptr {
@@ -36,7 +31,11 @@ impl<F: LurkField> Coprocessor<F> for DummyCoprocessor<F> {
     }
 }
 
-impl<F: LurkField> CoCircuit<F> for DummyCoprocessor<F> {}
+impl<F: LurkField> CoCircuit<F> for DummyCoprocessor<F> {
+    fn arity(&self) -> usize {
+        0
+    }
+}
 
 impl<F> DummyCoprocessor<F> {
     #[allow(dead_code)]
diff --git a/src/lem/circuit.rs b/src/lem/circuit.rs
index eca6856125..cd35b7de0b 100644
--- a/src/lem/circuit.rs
+++ b/src/lem/circuit.rs
@@ -379,15 +379,22 @@ pub fn build_slots_allocations<F: LurkField, CS: ConstraintSystem<F>>(
 }
 
 impl Block {
-    fn alloc_consts<F: LurkField, CS: ConstraintSystem<F>>(
+    fn alloc_consts<F: LurkField, CS: ConstraintSystem<F>, C: Coprocessor<F>>(
         &self,
         cs: &mut CS,
         store: &Store<F>,
         g: &GlobalAllocator<F>,
+        lang: &Lang<F, C>,
     ) {
         for op in &self.ops {
             match op {
-                Op::Call(_, func, _) => func.body.alloc_consts(cs, store, g),
+                Op::Cproc(_, sym, _) => {
+                    let Some(c) = lang.lookup_by_sym(sym) else {
+                        panic!("No coprocessor is bound to {sym}")
+                    };
+                    c.alloc_globals(cs, g, store);
+                }
+                Op::Call(_, func, _) => func.body.alloc_consts(cs, store, g, lang),
                 Op::Cons2(_, tag, _)
                 | Op::Cons3(_, tag, _)
                 | Op::Cons4(_, tag, _)
@@ -446,24 +453,24 @@ impl Block {
         }
         match &self.ctrl {
             Ctrl::If(.., a, b) => {
-                a.alloc_consts(cs, store, g);
-                b.alloc_consts(cs, store, g);
+                a.alloc_consts(cs, store, g, lang);
+                b.alloc_consts(cs, store, g, lang);
             }
             Ctrl::MatchTag(_, cases, def) => {
                 for block in cases.values() {
-                    block.alloc_consts(cs, store, g);
+                    block.alloc_consts(cs, store, g, lang);
                 }
                 if let Some(def) = def {
-                    def.alloc_consts(cs, store, g);
+                    def.alloc_consts(cs, store, g, lang);
                 }
             }
             Ctrl::MatchSymbol(_, cases, def) => {
                 g.alloc_tag(cs, &Sym);
                 for block in cases.values() {
-                    block.alloc_consts(cs, store, g);
+                    block.alloc_consts(cs, store, g, lang);
                 }
                 if let Some(def) = def {
-                    def.alloc_consts(cs, store, g);
+                    def.alloc_consts(cs, store, g, lang);
                 }
             }
             Ctrl::Return(..) => (),
@@ -1373,13 +1380,14 @@ impl Func {
         }
     }
 
-    pub fn alloc_consts<F: LurkField, CS: ConstraintSystem<F>>(
+    pub fn alloc_consts<F: LurkField, CS: ConstraintSystem<F>, C: Coprocessor<F>>(
         &self,
         cs: &mut CS,
         store: &Store<F>,
+        lang: &Lang<F, C>,
     ) -> GlobalAllocator<F> {
         let g = GlobalAllocator::default();
-        self.body.alloc_consts(cs, store, &g);
+        self.body.alloc_consts(cs, store, &g, lang);
         g
     }
 
@@ -1483,7 +1491,7 @@ impl Func {
         lang: &Lang<F, C>,
     ) -> Result<()> {
         let bound_allocations = &mut BoundAllocations::new();
-        let global_allocator = self.alloc_consts(cs, store);
+        let global_allocator = self.alloc_consts(cs, store, lang);
         self.allocate_input(cs, store, frame, bound_allocations);
         self.synthesize_frame(
             cs,
diff --git a/src/lem/multiframe.rs b/src/lem/multiframe.rs
index 95d15e4c51..3a5a3fba51 100644
--- a/src/lem/multiframe.rs
+++ b/src/lem/multiframe.rs
@@ -770,7 +770,7 @@ impl<'a, F: LurkField, C: Coprocessor<F>> Circuit<F> for MultiFrame<'a, F, C> {
                 allocated_output.push(AllocatedPtr::from_parts(allocated_tag, allocated_hash));
             }
 
-            let g = self.lurk_step.alloc_consts(cs, store);
+            let g = self.lurk_step.alloc_consts(cs, store, self.get_lang());
 
             let allocated_output_result =
                 self.synthesize_frames(cs, store, allocated_input, frames, &g)?;
@@ -891,13 +891,13 @@ impl<'a, F: LurkField, C: Coprocessor<F>> nova::traits::circuit::StepCircuit<F>
             if self.pc != 0 {
                 assert_eq!(frames.len(), 1);
             }
-            let g = self.lurk_step.alloc_consts(cs, store);
+            let g = self.lurk_step.alloc_consts(cs, store, self.get_lang());
             self.synthesize_frames(cs, store, input, frames, &g)?
         } else {
             let store = Store::default();
             let blank_frame = Frame::blank(self.get_func(), self.pc, &store);
             let frames = vec![blank_frame; self.num_frames];
-            let g = self.lurk_step.alloc_consts(cs, &store);
+            let g = self.lurk_step.alloc_consts(cs, &store, self.get_lang());
             self.synthesize_frames(cs, &store, input, &frames, &g)?
         };
 
@@ -1066,7 +1066,10 @@ mod tests {
                 }
             }
         }
-        let g = lurk_step.alloc_consts(&mut cs, &store);
+
+        let lang = Lang::<Bn>::new();
+
+        let g = lurk_step.alloc_consts(&mut cs, &store, &lang);
 
         // asserting equality of frames witnesses
         let frame = frames.first().unwrap();
@@ -1082,8 +1085,6 @@ mod tests {
 
         let mut cs_clone = cs.clone();
 
-        let lang = Lang::<Bn>::new();
-
         let output_sequential = synthesize_frames_sequential(
             &mut cs,
             &g,
diff --git a/src/proof/tests/nova_tests.rs b/src/proof/tests/nova_tests.rs
index 2ca29e510d..6b44b08ae6 100644
--- a/src/proof/tests/nova_tests.rs
+++ b/src/proof/tests/nova_tests.rs
@@ -9,8 +9,7 @@ use crate::{
         tag::Tag,
     },
     num::Num,
-    state::user_sym,
-    state::State,
+    state::{user_sym, State},
     tag::{ExprTag, Op, Op1, Op2},
 };
 
@@ -4047,7 +4046,7 @@ fn test_dumb_lang() {
 
 #[test]
 fn test_terminator_lang() {
-    use crate::{coprocessor::test::Terminator, state::user_sym};
+    use crate::coprocessor::test::Terminator;
 
     let mut lang = Lang::<Fr, Terminator<Fr>>::new();
     let dumb = Terminator::new();
@@ -4073,6 +4072,32 @@ fn test_terminator_lang() {
     );
 }
 
+#[test]
+fn test_hello_world_lang() {
+    use crate::coprocessor::test::HelloWorld;
+
+    let mut lang = Lang::<Fr, HelloWorld<Fr>>::new();
+    let hello_world = HelloWorld::new();
+    let name = user_sym("hello-world");
+
+    let s = &Store::default();
+    lang.add_coprocessor(name, hello_world);
+
+    let res = HelloWorld::intern_hello_world(s);
+    let terminal = s.cont_terminal();
+
+    test_aux::<_, HelloWorld<_>>(
+        s,
+        "(hello-world)",
+        Some(res),
+        None,
+        Some(terminal),
+        None,
+        &expect!["1"],
+        &Some(lang.into()),
+    );
+}
+
 #[test]
 fn test_trie_lang() {
     use crate::coprocessor::trie::{install, TrieCoproc};