Signature operator collection

Cargo.toml

@@ -80,6 +80,7 @@ sha3 = "0.10.8"
 rstest = { workspace = true }
 criterion = { workspace = true }
 hex = { workspace = true }
+anyhow = "1.0.93"
 
 [[bench]]
 name = "run-program"

src/chia_dialect.rs

@@ -14,8 +14,11 @@ 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::Response;
-use crate::secp_ops::{op_secp256k1_verify, op_secp256r1_verify};
+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,
+};
 
 // unknown operators are disallowed
 // (otherwise they are no-ops with well defined cost)
@@ -69,6 +72,7 @@ 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 {
@@ -101,13 +105,33 @@ 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 {
-                0x13d61f00 => op_secp256k1_verify,
-                0x1c3a8f00 => op_secp256r1_verify,
+                // 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
+                }
                 _ => {
                     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/dialect.rs

@@ -1,6 +1,7 @@
 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)]
@@ -30,6 +31,7 @@ pub trait Dialect {
         args: NodePtr,
         max_cost: Cost,
         extensions: OperatorSet,
+        collected_ops: Option<&mut Vec<CollectedOp>>,
     ) -> Response;
     fn allow_unknown_ops(&self) -> bool;
 }

src/run_program.rs

@@ -64,6 +64,12 @@ 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
@@ -99,11 +105,11 @@ struct RunProgramContext<'a, D> {
     softfork_stack: Vec<SoftforkGuard>,
     #[cfg(feature = "counters")]
     pub counters: Counters,
-
     #[cfg(feature = "pre-eval")]
     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> {
@@ -195,10 +201,11 @@ 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) -> Self {
+    fn new(allocator: &'a mut Allocator, dialect: &'a D, collect_signatures: bool) -> Self {
         RunProgramContext {
             allocator,
             dialect,
@@ -212,6 +219,11 @@ 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
+            },
         }
     }
 
@@ -410,6 +422,7 @@ 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)
@@ -511,10 +524,25 @@ pub fn run_program<'a, D: Dialect>(
     env: NodePtr,
     max_cost: Cost,
 ) -> Response {
-    let mut rpc = RunProgramContext::new(allocator, dialect);
+    let mut rpc = RunProgramContext::new(allocator, dialect, false);
     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,
@@ -536,7 +564,7 @@ pub fn run_program_with_counters<'a, D: Dialect>(
     env: NodePtr,
     max_cost: Cost,
 ) -> (Counters, Response) {
-    let mut rpc = RunProgramContext::new(allocator, dialect);
+    let mut rpc = RunProgramContext::new(allocator, dialect, false);
     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;
@@ -551,6 +579,7 @@ mod tests {
 
     use crate::chia_dialect::{ENABLE_KECCAK, ENABLE_KECCAK_OPS_OUTSIDE_GUARD, NO_UNKNOWN_OPS};
     use crate::test_ops::parse_exp;
+    use crate::ChiaDialect;
 
     use rstest::rstest;
 
@@ -1580,4 +1609,39 @@ mod tests {
 
         assert_eq!(result.unwrap().0, cost);
     }
+
+    #[test]
+    fn test_signature_collection() -> anyhow::Result<()> {
+        let mut a = Allocator::new();
+
+        let op = a.new_atom(&[0x13, 0xd6, 0x1f, 0x00])?;
+        let fake_arg = a.new_atom(&[1, 2, 3])?;
+        let op_q = a.one();
+        let quoted_fake_arg = a.new_pair(op_q, fake_arg)?;
+        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();
+
+        assert!(a.atom(reduction.1).is_empty());
+        assert_eq!(collected.len(), 1);
+
+        let collected = collected[0];
+        assert_eq!(collected.op, op);
+
+        let SExp::Pair(f, r) = a.sexp(collected.args) else {
+            unreachable!();
+        };
+        assert!(a.atom(r).is_empty());
+        assert_eq!(f, fake_arg);
+
+        Ok(())
+    }
 }

src/runtime_dialect.rs

@@ -6,6 +6,7 @@ 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 {
@@ -41,6 +42,7 @@ 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();

src/secp_ops.rs

@@ -7,8 +7,8 @@ use k256::ecdsa::{Signature as K1Signature, VerifyingKey as K1VerifyingKey};
 use p256::ecdsa::signature::hazmat::PrehashVerifier;
 use p256::ecdsa::{Signature as P1Signature, VerifyingKey as P1VerifyingKey};
 
-const SECP256R1_VERIFY_COST: Cost = 1850000;
-const SECP256K1_VERIFY_COST: Cost = 1300000;
+pub const SECP256R1_VERIFY_COST: Cost = 1850000;
+pub const SECP256K1_VERIFY_COST: Cost = 1300000;
 
 // expects: pubkey msg sig
 pub fn op_secp256r1_verify(a: &mut Allocator, input: NodePtr, max_cost: Cost) -> Response {