Merge pull request #651 from Chia-Network/get-puzzle-and-solution

add higher level python binding for `get_puzzle_and_solution_for_coin()`

crates/chia-consensus/fuzz/fuzz_targets/puzzle-coin-solution.rs

@@ -2,8 +2,10 @@
 use libfuzzer_sys::fuzz_target;
 
 use chia_consensus::gen::get_puzzle_and_solution::get_puzzle_and_solution_for_coin;
+use chia_protocol::Coin;
 use clvmr::allocator::Allocator;
 use fuzzing_utils::{make_tree, BitCursor};
+use std::collections::HashSet;
 
 const HASH: [u8; 32] = [
     0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
@@ -13,5 +15,10 @@ fuzz_target!(|data: &[u8]| {
     let mut a = Allocator::new();
     let input = make_tree(&mut a, &mut BitCursor::new(data), false);
 
-    let _ret = get_puzzle_and_solution_for_coin(&a, input, HASH.into(), 1337, HASH.into());
+    let _ret = get_puzzle_and_solution_for_coin(
+        &a,
+        input,
+        &HashSet::new(),
+        &Coin::new(HASH.into(), HASH.into(), 1337),
+    );
 });

crates/chia-consensus/src/gen/get_puzzle_and_solution.rs

@@ -1,8 +1,9 @@
 use crate::gen::validation_error::{atom, check_nil, first, next, rest, ErrorCode, ValidationErr};
-use chia_protocol::Bytes32;
-use clvm_utils::tree_hash;
+use chia_protocol::Coin;
+use clvm_utils::{tree_hash_cached, TreeHash};
 use clvmr::allocator::{Allocator, Atom, NodePtr};
 use clvmr::op_utils::u64_from_bytes;
+use std::collections::{HashMap, HashSet};
 
 // returns parent-coin ID, amount, puzzle-reveal and solution
 pub fn parse_coin_spend(
@@ -24,15 +25,15 @@ pub fn parse_coin_spend(
 pub fn get_puzzle_and_solution_for_coin(
     a: &Allocator,
     generator_result: NodePtr,
-    find_parent: Bytes32,
-    find_amount: u64,
-    find_ph: Bytes32,
+    backrefs: &HashSet<NodePtr>,
+    find_coin: &Coin,
 ) -> Result<(NodePtr, NodePtr), ValidationErr> {
     // the output from the block generator is a list of CoinSpends
     // with (parent-coin-id puzzle-reveal amount solution)
     // this function is given the generator output and a parent_coin_id, amount
     // and puzzle_hash and it will return the puzzle and solution for that given
     // coin spend, or fail if it cannot be found
+    let mut cache = HashMap::<NodePtr, TreeHash>::new();
     let mut iter = first(a, generator_result)?;
     while let Some((coin_spend, next)) = next(a, iter)? {
         iter = next;
@@ -41,12 +42,12 @@ pub fn get_puzzle_and_solution_for_coin(
 
         // we want to avoid having to compute the puzzle hash if we don't have to
         // so check parent and amount first
-        if parent.as_ref() != find_parent.as_ref() || amount != find_amount {
+        if parent.as_ref() != find_coin.parent_coin_info.as_ref() || amount != find_coin.amount {
             continue;
         }
 
-        let puzzle_hash = tree_hash(a, puzzle);
-        if puzzle_hash != find_ph.into() {
+        let puzzle_hash = tree_hash_cached(a, puzzle, backrefs, &mut cache);
+        if puzzle_hash != find_coin.puzzle_hash.into() {
             continue;
         }
 
@@ -63,7 +64,9 @@ mod test {
     use crate::gen::conditions::{u64_to_bytes, MempoolVisitor};
     use crate::gen::flags::{ALLOW_BACKREFS, MEMPOOL_MODE};
     use crate::gen::run_block_generator::{run_block_generator2, setup_generator_args};
+    use chia_protocol::Bytes32;
     use clvm_traits::FromClvm;
+    use clvm_utils::tree_hash;
     use clvmr::reduction::Reduction;
     use clvmr::serde::node_from_bytes_backrefs;
     use clvmr::sha2::Sha256;
@@ -131,9 +134,8 @@ mod test {
             get_puzzle_and_solution_for_coin(
                 &a,
                 generator_output,
-                parent,
-                1337,
-                tree_hash(&a, puzzle1).into()
+                &HashSet::new(),
+                &Coin::new(parent, tree_hash(&a, puzzle1).into(), 1337),
             )
             .unwrap(),
             (puzzle1, solution1)
@@ -144,9 +146,8 @@ mod test {
             get_puzzle_and_solution_for_coin(
                 &a,
                 generator_output,
-                make_dummy_id(2),
-                1337,
-                tree_hash(&a, puzzle1).into()
+                &HashSet::new(),
+                &Coin::new(make_dummy_id(2), tree_hash(&a, puzzle1).into(), 1337),
             )
             .unwrap_err()
             .1,
@@ -158,9 +159,8 @@ mod test {
             get_puzzle_and_solution_for_coin(
                 &a,
                 generator_output,
-                parent,
-                42,
-                tree_hash(&a, puzzle1).into()
+                &HashSet::new(),
+                &Coin::new(parent, tree_hash(&a, puzzle1).into(), 42),
             )
             .unwrap_err()
             .1,
@@ -169,9 +169,14 @@ mod test {
 
         // wrong puzzle hash
         assert_eq!(
-            get_puzzle_and_solution_for_coin(&a, generator_output, parent, 1337, make_dummy_id(4))
-                .unwrap_err()
-                .1,
+            get_puzzle_and_solution_for_coin(
+                &a,
+                generator_output,
+                &HashSet::new(),
+                &Coin::new(parent, make_dummy_id(4), 1337),
+            )
+            .unwrap_err()
+            .1,
             ErrorCode::InvalidCondition
         );
     }
@@ -245,8 +250,11 @@ mod test {
         let checkpoint = a2.checkpoint();
         for s in &conds.spends {
             a2.restore_checkpoint(&checkpoint);
-            let mut expected_additions: HashSet<(Bytes32, u64)> =
-                HashSet::from_iter(s.create_coin.iter().map(|c| (c.puzzle_hash, c.amount)));
+            let mut expected_additions: HashSet<(Bytes32, u64)> = s
+                .create_coin
+                .iter()
+                .map(|c| (c.puzzle_hash, c.amount))
+                .collect();
 
             let dialect = &ChiaDialect::new(MEMPOOL_MODE);
             let args = setup_generator_args(&mut a2, blocks).expect("setup_generator_args");
@@ -257,9 +265,12 @@ mod test {
             let (puzzle, solution) = get_puzzle_and_solution_for_coin(
                 &a2,
                 result,
-                a.atom(s.parent_id).as_ref().try_into().unwrap(),
-                s.coin_amount,
-                a.atom(s.puzzle_hash).as_ref().try_into().unwrap(),
+                &HashSet::new(),
+                &Coin::new(
+                    a.atom(s.parent_id).as_ref().try_into().unwrap(),
+                    a.atom(s.puzzle_hash).as_ref().try_into().unwrap(),
+                    s.coin_amount,
+                ),
             )
             .expect("get_puzzle_and_solution_for_coin");
 

tests/test_get_puzzle_and_solution.py

@@ -1,11 +1,15 @@
 from chia_rs import (
     get_puzzle_and_solution_for_coin,
+    get_puzzle_and_solution_for_coin2,
     run_block_generator2,
     ALLOW_BACKREFS,
     run_chia_program,
+    Program,
+    Coin,
 )
 from run_gen import DEFAULT_CONSTANTS
 from chia_rs.sized_bytes import bytes32
+from chia_rs.sized_ints import uint64
 import pytest
 
 DESERIALIZE_MOD = bytes.fromhex(
@@ -61,6 +65,16 @@ def test_get_puzzle_and_solution_for_coin(input_file: str) -> None:
             bytes32(s.puzzle_hash),
             ALLOW_BACKREFS,
         )
+        puzzle2, solution2 = get_puzzle_and_solution_for_coin2(
+            Program.from_bytes(block),
+            [],
+            11000000000,
+            Coin(bytes32(s.parent_id), bytes32(s.puzzle_hash), uint64(s.coin_amount)),
+            ALLOW_BACKREFS,
+        )
+        assert puzzle == bytes(puzzle2)
+        assert solution == bytes(solution2)
+
         assert len(puzzle) > 0
         assert len(solution) > 0
 

wheel/generate_type_stubs.py

@@ -353,6 +353,7 @@ class LazyNode:
 def serialized_length(program: ReadableBuffer) -> int: ...
 def tree_hash(program: ReadableBuffer) -> bytes32: ...
 def get_puzzle_and_solution_for_coin(program: ReadableBuffer, args: ReadableBuffer, max_cost: int, find_parent: bytes32, find_amount: int, find_ph: bytes32, flags: int) -> Tuple[bytes, bytes]: ...
+def get_puzzle_and_solution_for_coin2(program: Program, block_refs: List[ReadableBuffer], max_cost: int, find_coin: Coin, flags: int) -> Tuple[Program, Program]: ...
 
 class BLSCache:
     def __init__(self, cache_size: Optional[int] = 50000) -> None: ...

wheel/python/chia_rs/chia_rs.pyi

@@ -97,6 +97,7 @@ class LazyNode:
 def serialized_length(program: ReadableBuffer) -> int: ...
 def tree_hash(program: ReadableBuffer) -> bytes32: ...
 def get_puzzle_and_solution_for_coin(program: ReadableBuffer, args: ReadableBuffer, max_cost: int, find_parent: bytes32, find_amount: int, find_ph: bytes32, flags: int) -> Tuple[bytes, bytes]: ...
+def get_puzzle_and_solution_for_coin2(program: Program, block_refs: List[ReadableBuffer], max_cost: int, find_coin: Coin, flags: int) -> Tuple[Program, Program]: ...
 
 class BLSCache:
     def __init__(self, cache_size: Optional[int] = 50000) -> None: ...

wheel/src/api.rs

@@ -7,6 +7,7 @@ use chia_consensus::gen::flags::{
     NO_UNKNOWN_CONDS, STRICT_ARGS_COUNT,
 };
 use chia_consensus::gen::owned_conditions::{OwnedSpend, OwnedSpendBundleConditions};
+use chia_consensus::gen::run_block_generator::setup_generator_args;
 use chia_consensus::gen::run_puzzle::run_puzzle as native_run_puzzle;
 use chia_consensus::gen::solution_generator::solution_generator as native_solution_generator;
 use chia_consensus::gen::solution_generator::solution_generator_backrefs as native_solution_generator_backrefs;
@@ -52,6 +53,7 @@ use pyo3::types::PyBytes;
 use pyo3::types::PyList;
 use pyo3::types::PyTuple;
 use pyo3::wrap_pyfunction;
+use std::collections::HashSet;
 use std::iter::zip;
 
 use crate::run_program::{run_chia_program, serialized_length};
@@ -65,7 +67,7 @@ use clvmr::reduction::EvalErr;
 use clvmr::reduction::Reduction;
 use clvmr::run_program;
 use clvmr::serde::node_to_bytes;
-use clvmr::serde::{node_from_bytes, node_from_bytes_backrefs};
+use clvmr::serde::{node_from_bytes, node_from_bytes_backrefs, node_from_bytes_backrefs_record};
 use clvmr::ChiaDialect;
 
 use chia_bls::{
@@ -115,6 +117,8 @@ pub fn tree_hash<'a>(py: Python<'a>, blob: PyBuffer<u8>) -> PyResult<Bound<'_, P
     Ok(PyBytes::new_bound(py, &tree_hash_from_bytes(slice)?))
 }
 
+// there is an updated version of this function that doesn't require serializing
+// and deserializing the generator and arguments.
 #[allow(clippy::too_many_arguments)]
 #[pyfunction]
 pub fn get_puzzle_and_solution_for_coin<'a>(
@@ -145,7 +149,12 @@ pub fn get_puzzle_and_solution_for_coin<'a>(
         .allow_threads(|| -> Result<(NodePtr, NodePtr), EvalErr> {
             let Reduction(_cost, result) =
                 run_program(&mut allocator, dialect, program, args, max_cost)?;
-            match parse_puzzle_solution(&allocator, result, find_parent, find_amount, find_ph) {
+            match parse_puzzle_solution(
+                &allocator,
+                result,
+                &HashSet::new(),
+                &Coin::new(find_parent, find_ph, find_amount),
+            ) {
                 Err(ValidationErr(n, _)) => Err(EvalErr(n, "coin not found".to_string())),
                 Ok(pair) => Ok(pair),
             }
@@ -170,6 +179,55 @@ pub fn get_puzzle_and_solution_for_coin<'a>(
     ))
 }
 
+// This is a new version of get_puzzle_and_solution_for_coin() which uses the
+// right types for generator, blocks_refs and the return value.
+// The old version was written when Program was a python type had to be
+// serialized to bytes through rust boundary.
+#[allow(clippy::too_many_arguments)]
+#[pyfunction]
+pub fn get_puzzle_and_solution_for_coin2<'a>(
+    py: Python<'a>,
+    generator: &Program,
+    block_refs: &Bound<'a, PyList>,
+    max_cost: Cost,
+    find_coin: &Coin,
+    flags: u32,
+) -> PyResult<(Program, Program)> {
+    let mut allocator = make_allocator(LIMIT_HEAP);
+
+    let refs = block_refs.into_iter().map(|b| {
+        let buf = b
+            .extract::<PyBuffer<u8>>()
+            .expect("block_refs should be a list of buffers");
+        py_to_slice::<'a>(buf)
+    });
+
+    let (generator, backrefs) =
+        node_from_bytes_backrefs_record(&mut allocator, generator.as_ref())?;
+    let args = setup_generator_args(&mut allocator, refs)?;
+    let dialect = &ChiaDialect::new(flags);
+
+    let (puzzle, solution) = py
+        .allow_threads(|| -> Result<(NodePtr, NodePtr), EvalErr> {
+            let Reduction(_cost, result) =
+                run_program(&mut allocator, dialect, generator, args, max_cost)?;
+            match parse_puzzle_solution(&allocator, result, &backrefs, find_coin) {
+                Err(ValidationErr(n, _)) => Err(EvalErr(n, "coin not found".to_string())),
+                Ok(pair) => Ok(pair),
+            }
+        })
+        .map_err(|e| {
+            let blob = node_to_bytes(&allocator, e.0).ok().map(hex::encode);
+            PyValueError::new_err((e.1, blob))
+        })?;
+
+    // keep serializing normally, until wallets support backrefs
+    Ok((
+        node_to_bytes(&allocator, puzzle)?.into(),
+        node_to_bytes(&allocator, solution)?.into(),
+    ))
+}
+
 #[pyfunction]
 fn run_puzzle(
     puzzle: &[u8],
@@ -571,6 +629,7 @@ pub fn chia_rs(_py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_function(wrap_pyfunction!(compute_merkle_set_root, m)?)?;
     m.add_function(wrap_pyfunction!(tree_hash, m)?)?;
     m.add_function(wrap_pyfunction!(get_puzzle_and_solution_for_coin, m)?)?;
+    m.add_function(wrap_pyfunction!(get_puzzle_and_solution_for_coin2, m)?)?;
 
     // facilities from chia-bls