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

@@ -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() {

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)
         }

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)
+        }
+    }
 }

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
     }

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.

src/eval/lang.rs

@@ -24,19 +24,18 @@ 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 {
         s.intern_nil()
     }
 }
 
-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)]

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,