CollectDialect

src/chia_dialect.rs

@@ -14,11 +14,8 @@ use crate::more_ops::{
     op_logand, op_logior, op_lognot, op_logxor, op_lsh, op_mod, op_modpow, op_multiply, op_not,
     op_point_add, op_pubkey_for_exp, op_sha256, op_strlen, op_substr, op_subtract, op_unknown,
 };
-use crate::reduction::{Reduction, Response};
-use crate::run_program::CollectedOp;
-use crate::secp_ops::{
-    op_secp256k1_verify, op_secp256r1_verify, SECP256K1_VERIFY_COST, SECP256R1_VERIFY_COST,
-};
+use crate::reduction::Response;
+use crate::secp_ops::{op_secp256k1_verify, op_secp256r1_verify};
 
 // unknown operators are disallowed
 // (otherwise they are no-ops with well defined cost)
@@ -72,7 +69,6 @@ impl Dialect for ChiaDialect {
         argument_list: NodePtr,
         max_cost: Cost,
         extension: OperatorSet,
-        collected_ops: Option<&mut Vec<CollectedOp>>,
     ) -> Response {
         let flags = self.flags
             | match extension {
@@ -105,33 +101,15 @@ impl Dialect for ChiaDialect {
             // the secp operators have a fixed cost of 1850000 and 1300000,
             // which makes the multiplier 0x1c3a8f and 0x0cf84f (there is an
             // implied +1) and cost function 0
-
-            let cost;
-
             let f = match opcode {
                 // If we add more operators here, we will have to exclude them from signature collection.
-                0x13d61f00 => {
-                    cost = SECP256K1_VERIFY_COST;
-                    op_secp256k1_verify
-                }
-                0x1c3a8f00 => {
-                    cost = SECP256R1_VERIFY_COST;
-                    op_secp256r1_verify
-                }
+                0x13d61f00 => op_secp256k1_verify,
+                0x1c3a8f00 => op_secp256r1_verify,
                 _ => {
                     return unknown_operator(allocator, o, argument_list, flags, max_cost);
                 }
             };
 
-            if let Some(ops) = collected_ops {
-                ops.push(CollectedOp {
-                    op: o,
-                    args: argument_list,
-                });
-
-                return Ok(Reduction(cost, NodePtr::NIL));
-            }
-
             return f(allocator, argument_list, max_cost);
         }
         if op_len != 1 {

src/collect_dialect.rs

@@ -0,0 +1,90 @@
+use std::cell::RefCell;
+
+use crate::{
+    cost::Cost,
+    dialect::{Dialect, OperatorSet},
+    reduction::{Reduction, Response},
+    secp_ops::{SECP256K1_VERIFY_COST, SECP256R1_VERIFY_COST},
+    Allocator, NodePtr,
+};
+
+#[derive(Debug, Clone, Copy)]
+pub struct CollectedOp {
+    pub op: NodePtr,
+    pub args: NodePtr,
+}
+
+#[derive(Debug, Default, Clone)]
+pub struct CollectDialect<T> {
+    dialect: T,
+    collected_ops: RefCell<Vec<CollectedOp>>,
+}
+
+impl<T> CollectDialect<T> {
+    pub fn new(dialect: T) -> Self {
+        Self {
+            dialect,
+            collected_ops: RefCell::new(Vec::new()),
+        }
+    }
+
+    pub fn collect(self) -> Vec<CollectedOp> {
+        self.collected_ops.into_inner()
+    }
+}
+
+impl<T> Dialect for CollectDialect<T>
+where
+    T: Dialect,
+{
+    fn apply_kw(&self) -> u32 {
+        self.dialect.apply_kw()
+    }
+
+    fn quote_kw(&self) -> u32 {
+        self.dialect.quote_kw()
+    }
+
+    fn softfork_kw(&self) -> u32 {
+        self.dialect.softfork_kw()
+    }
+
+    fn allow_unknown_ops(&self) -> bool {
+        self.dialect.allow_unknown_ops()
+    }
+
+    fn softfork_extension(&self, ext: u32) -> OperatorSet {
+        self.dialect.softfork_extension(ext)
+    }
+
+    fn op(
+        &self,
+        allocator: &mut Allocator,
+        op: NodePtr,
+        args: NodePtr,
+        max_cost: Cost,
+        extensions: OperatorSet,
+    ) -> Response {
+        let response = self.dialect.op(allocator, op, args, max_cost, extensions);
+
+        let op_len = allocator.atom_len(op);
+        if op_len != 4 {
+            return response;
+        }
+
+        let atom = allocator.atom(op);
+        let opcode = u32::from_be_bytes(atom.as_ref().try_into().unwrap());
+
+        let cost = match opcode {
+            0x13d61f00 => SECP256K1_VERIFY_COST,
+            0x1c3a8f00 => SECP256R1_VERIFY_COST,
+            _ => return response,
+        };
+
+        self.collected_ops
+            .borrow_mut()
+            .push(CollectedOp { op, args });
+
+        Ok(Reduction(cost, NodePtr::NIL))
+    }
+}

src/dialect.rs

@@ -1,7 +1,6 @@
 use crate::allocator::{Allocator, NodePtr};
 use crate::cost::Cost;
 use crate::reduction::Response;
-use crate::run_program::CollectedOp;
 
 /// The set of operators that are available in the dialect.
 #[repr(u32)]
@@ -31,7 +30,6 @@ pub trait Dialect {
         args: NodePtr,
         max_cost: Cost,
         extensions: OperatorSet,
-        collected_ops: Option<&mut Vec<CollectedOp>>,
     ) -> Response;
     fn allow_unknown_ops(&self) -> bool;
 }

src/lib.rs

@@ -1,6 +1,7 @@
 pub mod allocator;
 pub mod bls_ops;
 pub mod chia_dialect;
+pub mod collect_dialect;
 pub mod core_ops;
 pub mod cost;
 pub mod dialect;
@@ -19,6 +20,7 @@ pub mod traverse_path;
 
 pub use allocator::{Allocator, Atom, NodePtr, SExp};
 pub use chia_dialect::ChiaDialect;
+pub use collect_dialect::{CollectDialect, CollectedOp};
 pub use run_program::run_program;
 
 pub use chia_dialect::{LIMIT_HEAP, MEMPOOL_MODE, NO_UNKNOWN_OPS};

src/run_program.rs

@@ -64,12 +64,6 @@ impl Counters {
     }
 }
 
-#[derive(Debug, Clone, Copy)]
-pub struct CollectedOp {
-    pub op: NodePtr,
-    pub args: NodePtr,
-}
-
 // this represents the state we were in before entering a soft-fork guard. We
 // may need this to long-jump out of the guard, and also to validate the cost
 // when exiting the guard
@@ -109,7 +103,6 @@ struct RunProgramContext<'a, D> {
     pre_eval: Option<PreEval>,
     #[cfg(feature = "pre-eval")]
     posteval_stack: Vec<Box<PostEval>>,
-    collected_ops: Option<Vec<CollectedOp>>,
 }
 
 fn augment_cost_errors(r: Result<Cost, EvalErr>, max_cost: NodePtr) -> Result<Cost, EvalErr> {
@@ -201,11 +194,10 @@ impl<'a, D: Dialect> RunProgramContext<'a, D> {
             counters: Counters::new(),
             pre_eval,
             posteval_stack: Vec::new(),
-            collected_ops: None,
         }
     }
 
-    fn new(allocator: &'a mut Allocator, dialect: &'a D, collect_signatures: bool) -> Self {
+    fn new(allocator: &'a mut Allocator, dialect: &'a D) -> Self {
         RunProgramContext {
             allocator,
             dialect,
@@ -219,11 +211,6 @@ impl<'a, D: Dialect> RunProgramContext<'a, D> {
             pre_eval: None,
             #[cfg(feature = "pre-eval")]
             posteval_stack: Vec::new(),
-            collected_ops: if collect_signatures {
-                Some(Vec::new())
-            } else {
-                None
-            },
         }
     }
 
@@ -422,7 +409,6 @@ impl<'a, D: Dialect> RunProgramContext<'a, D> {
                 operand_list,
                 max_cost,
                 current_extensions,
-                self.collected_ops.as_mut(),
             )?;
             self.push(r.1)?;
             Ok(r.0)
@@ -524,25 +510,10 @@ pub fn run_program<'a, D: Dialect>(
     env: NodePtr,
     max_cost: Cost,
 ) -> Response {
-    let mut rpc = RunProgramContext::new(allocator, dialect, false);
+    let mut rpc = RunProgramContext::new(allocator, dialect);
     rpc.run_program(program, env, max_cost)
 }
 
-pub fn run_program_with_signatures<'a, D: Dialect>(
-    allocator: &'a mut Allocator,
-    dialect: &'a D,
-    program: NodePtr,
-    env: NodePtr,
-    max_cost: Cost,
-) -> Result<(Reduction, Vec<CollectedOp>), EvalErr> {
-    let mut rpc = RunProgramContext::new(allocator, dialect, true);
-    let reduction = rpc.run_program(program, env, max_cost)?;
-    Ok((
-        reduction,
-        rpc.collected_ops.expect("missing collected signature ops"),
-    ))
-}
-
 #[cfg(feature = "pre-eval")]
 pub fn run_program_with_pre_eval<'a, D: Dialect>(
     allocator: &'a mut Allocator,
@@ -564,7 +535,7 @@ pub fn run_program_with_counters<'a, D: Dialect>(
     env: NodePtr,
     max_cost: Cost,
 ) -> (Counters, Response) {
-    let mut rpc = RunProgramContext::new(allocator, dialect, false);
+    let mut rpc = RunProgramContext::new(allocator, dialect);
     let ret = rpc.run_program(program, env, max_cost);
     rpc.counters.atom_count = rpc.allocator.atom_count() as u32;
     rpc.counters.small_atom_count = rpc.allocator.small_atom_count() as u32;
@@ -578,6 +549,7 @@ mod tests {
     use super::*;
 
     use crate::chia_dialect::{ENABLE_KECCAK, ENABLE_KECCAK_OPS_OUTSIDE_GUARD, NO_UNKNOWN_OPS};
+    use crate::collect_dialect::CollectDialect;
     use crate::test_ops::parse_exp;
     use crate::ChiaDialect;
 
@@ -1621,14 +1593,10 @@ mod tests {
         let args = a.new_pair(quoted_fake_arg, NodePtr::NIL)?;
         let program = a.new_pair(op, args)?;
 
-        let (reduction, collected) = run_program_with_signatures(
-            &mut a,
-            &ChiaDialect::new(0),
-            program,
-            NodePtr::NIL,
-            u64::MAX,
-        )
-        .unwrap();
+        let dialect = CollectDialect::new(ChiaDialect::new(0));
+
+        let reduction = run_program(&mut a, &dialect, program, NodePtr::NIL, u64::MAX).unwrap();
+        let collected = dialect.collect();
 
         assert!(a.atom(reduction.1).is_empty());
         assert_eq!(collected.len(), 1);

src/runtime_dialect.rs

@@ -6,7 +6,6 @@ use crate::err_utils::err;
 use crate::f_table::{f_lookup_for_hashmap, FLookup};
 use crate::more_ops::op_unknown;
 use crate::reduction::Response;
-use crate::run_program::CollectedOp;
 use std::collections::HashMap;
 
 pub struct RuntimeDialect {
@@ -42,7 +41,6 @@ impl Dialect for RuntimeDialect {
         argument_list: NodePtr,
         max_cost: Cost,
         _extensions: OperatorSet,
-        _collected_ops: Option<&mut Vec<CollectedOp>>,
     ) -> Response {
         let atom = allocator.atom(o);
         let b = atom.as_ref();