[CHIA-3793] extend SIMPLE_GENERATOR flag

crates/chia-consensus/src/additions_and_removals.rs

@@ -36,7 +36,7 @@ where
 
     let program = node_from_bytes_backrefs(&mut a, program)?;
 
-    let args = setup_generator_args(&mut a, block_refs)?;
+    let args = setup_generator_args(&mut a, block_refs, flags)?;
     let dialect = ChiaDialect::new(flags);
 
     let Reduction(clvm_cost, all_spends) = run_program(&mut a, &dialect, program, args, cost_left)?;

crates/chia-consensus/src/get_puzzle_and_solution.rs

@@ -252,7 +252,7 @@ mod test {
                 .collect();
 
             let dialect = &ChiaDialect::new(MEMPOOL_MODE);
-            let args = setup_generator_args(&mut a2, blocks).expect("setup_generator_args");
+            let args = setup_generator_args(&mut a2, blocks, 0).expect("setup_generator_args");
             let Reduction(_, result) =
                 run_program(&mut a2, dialect, generator_node, args, MAX_COST)
                     .expect("run_program (generator)");

crates/chia-consensus/src/run_block_generator.rs

@@ -39,10 +39,19 @@ pub fn subtract_cost(
 pub fn setup_generator_args<GenBuf: AsRef<[u8]>, I: IntoIterator<Item = GenBuf>>(
     a: &mut Allocator,
     block_refs: I,
+    flags: u32,
 ) -> Result<NodePtr, ValidationErr>
 where
     <I as IntoIterator>::IntoIter: DoubleEndedIterator,
 {
+    // once we have soft-forked in requiring simple generators, we no longer
+    // need to pass in the deserialization program
+    if (flags & SIMPLE_GENERATOR) != 0 {
+        if block_refs.into_iter().next().is_some() {
+            return Err(ValidationErr(a.nil(), ErrorCode::TooManyGeneratorRefs));
+        }
+        return Ok(a.nil());
+    }
     let clvm_deserializer = node_from_bytes(a, &CHIALISP_DESERIALISATION)?;
 
     // iterate in reverse order since we're building a linked list from
@@ -219,7 +228,7 @@ where
     let program = node_from_bytes_backrefs(a, program)?;
     check_generator_node(a, program, flags)?;
 
-    let args = setup_generator_args(a, block_refs)?;
+    let args = setup_generator_args(a, block_refs, flags)?;
     let dialect = ChiaDialect::new(flags);
 
     let Reduction(clvm_cost, all_spends) = run_program(a, &dialect, program, args, cost_left)?;
@@ -307,7 +316,7 @@ where
 
     let program = node_from_bytes_backrefs(&mut a, generator)?;
     check_generator_node(&a, program, flags)?;
-    let args = setup_generator_args(&mut a, refs)?;
+    let args = setup_generator_args(&mut a, refs, flags)?;
     let dialect = ChiaDialect::new(flags);
 
     let Reduction(_clvm_cost, res) = run_program(
@@ -382,7 +391,7 @@ where
 
     let program = node_from_bytes_backrefs(&mut a, generator)?;
     check_generator_node(&a, program, flags)?;
-    let args = setup_generator_args(&mut a, refs)?;
+    let args = setup_generator_args(&mut a, refs, flags)?;
     let dialect = ChiaDialect::new(flags);
 
     let Reduction(_clvm_cost, res) = run_program(

crates/chia-consensus/src/spendbundle_conditions.rs

@@ -231,7 +231,7 @@ mod tests {
         let mut a = make_allocator(MEMPOOL_MODE);
 
         let generator = node_from_bytes_backrefs(&mut a, generator).expect("node_from_bytes");
-        let args = setup_generator_args(&mut a, block_refs).expect("setup_generator_args");
+        let args = setup_generator_args(&mut a, block_refs, 0).expect("setup_generator_args");
         let dialect = ChiaDialect::new(MEMPOOL_MODE);
         let Reduction(_, mut all_spends) =
             run_program(&mut a, &dialect, generator, args, 11_000_000_000).expect("run_program");

crates/chia-consensus/src/test_generators.rs

@@ -11,8 +11,9 @@ use crate::validation_error::ErrorCode;
 use chia_bls::Signature;
 use chia_protocol::Program;
 use chia_protocol::{Bytes, Bytes48};
+use chia_puzzles::CHIALISP_DESERIALISATION;
 use clvmr::allocator::NodePtr;
-use clvmr::serde::node_from_bytes_backrefs;
+use clvmr::serde::{node_from_bytes, node_from_bytes_backrefs};
 use clvmr::Allocator;
 use std::iter::zip;
 use text_diff::diff;
@@ -282,7 +283,9 @@ fn run_generator(#[case] name: &str) {
             flags |= COST_CONDITIONS;
         }
 
-        if ![
+        // These are generators that are programs, not a simple list of spends
+        // when the SIMPLE_GENERATOR flag is set, these should fail
+        if [
             "single-coin-only-garbage",
             "many-coins-announcement-cap",
             "puzzle-hash-stress-test",
@@ -303,32 +306,50 @@ fn run_generator(#[case] name: &str) {
         ]
         .contains(&name)
         {
-            // test that we allow quoted generators with the flag
-            flags |= SIMPLE_GENERATOR;
-        } else {
             // lets test that the procedural generators are filtered with the flag
             let mut a = make_allocator(flags);
             let test_conds = run_block_generator2(
                 &mut a,
                 &generator,
-                &block_refs,
+                &block_refs, // we're not allowed to pass in block references when SIMPLE_GENERATOR is set
                 11_000_000_000,
                 flags | DONT_VALIDATE_SIGNATURE | SIMPLE_GENERATOR,
                 &Signature::default(),
                 None,
                 &TEST_CONSTANTS,
             );
-            assert!(test_conds.is_err());
             assert_eq!(
                 test_conds.unwrap_err().1,
                 ErrorCode::ComplexGeneratorReceived
             );
+
             // now lets specifically check the node generator check
-            let program =
-                node_from_bytes_backrefs(&mut a, generator.as_ref()).expect("should be ok");
+            let program = node_from_bytes_backrefs(&mut a, generator.as_ref())
+                .expect("node_from_bytes_backref");
             let res = check_generator_node(&a, program, flags | SIMPLE_GENERATOR);
-            assert!(res.is_err());
             assert_eq!(res.unwrap_err().1, ErrorCode::ComplexGeneratorReceived);
+        } else {
+            flags |= SIMPLE_GENERATOR;
+            // ensure SIMPLE_GENERATOR fails if there are any block references
+            // passed in. We pass in a dummy block reference
+            let mut a = make_allocator(flags);
+            let test_conds = run_block_generator2(
+                &mut a,
+                &generator,
+                &[&[0_u8, 1, 2, 3]],
+                11_000_000_000,
+                flags | DONT_VALIDATE_SIGNATURE,
+                &Signature::default(),
+                None,
+                &TEST_CONSTANTS,
+            );
+            assert_eq!(test_conds.unwrap_err().1, ErrorCode::TooManyGeneratorRefs);
+
+            // now lets specifically check the node generator check
+            let program = node_from_bytes_backrefs(&mut a, generator.as_ref())
+                .expect("node_from_bytes_backref");
+            let res = check_generator_node(&a, program, flags | SIMPLE_GENERATOR);
+            assert!(res.is_ok());
         }
 
         println!("flags: {flags:x}");
@@ -352,6 +373,20 @@ fn run_generator(#[case] name: &str) {
                 // the generator itself has execution cost. At least the cost of
                 // a quote
                 assert!(exe_cost <= conditions.execution_cost);
+
+                if (flags & SIMPLE_GENERATOR) != 0 {
+                    // when running generators with the SIMPLE_GENERATOR flag
+                    // set, we don't pass in the CLVM deserializer program. This
+                    // causes the atoms and pairs counters to be lower than
+                    // before (when the test cases were created). In order to
+                    // match the test cases, we increment those counters in
+                    // order to print compatible output.
+                    // these are the atoms and pairs that would have been
+                    // allocated by the deserializer program
+                    let _ =
+                        node_from_bytes(&mut a2, &CHIALISP_DESERIALISATION).expect("deserializer");
+                    a2.add_ghost_pair(2).expect("add_ghost_pair");
+                }
                 (conditions.cost, print_conditions(&a2, &conditions, &a2))
             }
             Err(code) => (0, format!("FAILED: {}\n", u32::from(code.1))),

wheel/src/api.rs

@@ -215,7 +215,7 @@ pub fn get_puzzle_and_solution_for_coin2<'a>(
 
     let generator = node_from_bytes_backrefs(&mut allocator, generator.as_ref())
         .map_err(|e| map_pyerr_w_ptr(&e, &allocator))?;
-    let args = setup_generator_args(&mut allocator, refs)?;
+    let args = setup_generator_args(&mut allocator, refs, flags)?;
     let dialect = &ChiaDialect::new(flags);
 
     let (puzzle, solution) = py