From 0217d83016f13c5069ed0db87457bb7d106941dc Mon Sep 17 00:00:00 2001 From: arvidn Date: Mon, 19 Dec 2022 16:05:08 +0100 Subject: [PATCH] Add generator ROM and higher level function to run generators --- src/gen/conditions.rs | 5 ++- src/gen/mod.rs | 1 + src/gen/run_block_generator.rs | 57 ++++++++++++++++++++++++++++++ src/gen/validation_error.rs | 15 ++++++++ src/generator_rom.rs | 55 +++++++++++++++++++++++++++++ src/lib.rs | 1 + tests/run_gen.py | 63 ++++++---------------------------- tests/test-generators.py | 12 +++---- wheel/chia_rs.pyi | 4 +++ wheel/generate_type_stubs.py | 4 +++ wheel/src/api.rs | 3 +- wheel/src/run_generator.rs | 39 +++++++++++++++++++-- 12 files changed, 196 insertions(+), 63 deletions(-) create mode 100644 src/gen/run_block_generator.rs create mode 100644 src/generator_rom.rs diff --git a/src/gen/conditions.rs b/src/gen/conditions.rs index a35c60ee0..bdaee8881 100644 --- a/src/gen/conditions.rs +++ b/src/gen/conditions.rs @@ -348,7 +348,10 @@ pub struct SpendBundleConditions { // Unsafe Agg Sig conditions (i.e. not tied to the spend generating it) pub agg_sig_unsafe: Vec<(NodePtr, NodePtr)>, - // the total cost of the spend bundle + // the cost of conditions (when returned by parse_spends()) + // run_generator() will include the CLVM cost + // run_block_generator() will include CLVM cost and byte cost (making this + // the total cost) pub cost: u64, } diff --git a/src/gen/mod.rs b/src/gen/mod.rs index 8cb6fb336..30ac7b743 100644 --- a/src/gen/mod.rs +++ b/src/gen/mod.rs @@ -4,5 +4,6 @@ pub mod conditions; pub mod flags; pub mod get_puzzle_and_solution; pub mod opcodes; +pub mod run_block_generator; mod sanitize_int; pub mod validation_error; diff --git a/src/gen/run_block_generator.rs b/src/gen/run_block_generator.rs new file mode 100644 index 000000000..988cd870b --- /dev/null +++ b/src/gen/run_block_generator.rs @@ -0,0 +1,57 @@ +use crate::gen::conditions::{parse_spends, SpendBundleConditions}; +use crate::gen::validation_error::ValidationErr; +use crate::generator_rom::{COST_PER_BYTE, GENERATOR_ROM}; +use clvmr::allocator::Allocator; +use clvmr::chia_dialect::ChiaDialect; +use clvmr::reduction::Reduction; +use clvmr::run_program::run_program; +use clvmr::serde::node_from_bytes; + +// Runs the generator ROM and passes in the program (transactions generator). +// The program is expected to return a list of spends. Each item being: + +// (parent-coin-id puzzle-reveal amount solution) + +// The puzzle-reveals are then executed with the corresponding solution being +// passed as the argument. The output from those puzzles are lists of +// conditions. The conditions are parsed and returned in the +// SpendBundleConditions. Some conditions are validated, and if invalid may +// cause the function to return an error. + +// the only reason we need to pass in the allocator is because the returned +// SpendBundleConditions contains NodePtr fields. If that's changed, we could +// create the allocator inside this functions as well. +pub fn run_block_generator>( + a: &mut Allocator, + program: &[u8], + block_refs: &[GenBuf], + max_cost: u64, + flags: u32, +) -> Result { + let byte_cost = program.len() as u64 * COST_PER_BYTE; + + let generator_rom = node_from_bytes(a, &GENERATOR_ROM)?; + let program = node_from_bytes(a, program)?; + + // iterate in reverse order since we're building a linked list from + // the tail + let mut args = a.null(); + for g in block_refs.iter().rev() { + let ref_gen = a.new_atom(g.as_ref())?; + args = a.new_pair(ref_gen, args)?; + } + + args = a.new_pair(args, a.null())?; + let args = a.new_pair(args, a.null())?; + let args = a.new_pair(program, args)?; + + let dialect = ChiaDialect::new(flags); + let Reduction(clvm_cost, generator_output) = + run_program(a, &dialect, generator_rom, args, max_cost - byte_cost)?; + + // we pass in what's left of max_cost here, to fail early in case the + // cost of a condition brings us over the cost limit + let mut result = parse_spends(a, generator_output, max_cost - clvm_cost - byte_cost, flags)?; + result.cost += clvm_cost + byte_cost; + Ok(result) +} diff --git a/src/gen/validation_error.rs b/src/gen/validation_error.rs index 8861d9676..464c1e80e 100644 --- a/src/gen/validation_error.rs +++ b/src/gen/validation_error.rs @@ -1,7 +1,9 @@ use clvmr::allocator::{Allocator, NodePtr, SExp}; +use clvmr::reduction::EvalErr; #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum ErrorCode { + GeneratorRuntimeError, NegativeAmount, AmountExceedsMaximum, InvalidConditionOpcode, @@ -33,6 +35,18 @@ pub enum ErrorCode { #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub struct ValidationErr(pub NodePtr, pub ErrorCode); +impl From for ValidationErr { + fn from(v: EvalErr) -> Self { + ValidationErr(v.0, ErrorCode::GeneratorRuntimeError) + } +} + +impl From for ValidationErr { + fn from(_: std::io::Error) -> Self { + ValidationErr(-1, ErrorCode::GeneratorRuntimeError) + } +} + // helper functions that fail with ValidationErr pub fn first(a: &Allocator, n: NodePtr) -> Result { match a.sexp(n) { @@ -45,6 +59,7 @@ pub fn first(a: &Allocator, n: NodePtr) -> Result { impl From for u32 { fn from(err: ErrorCode) -> u32 { match err { + ErrorCode::GeneratorRuntimeError => 117, ErrorCode::NegativeAmount => 124, ErrorCode::AmountExceedsMaximum => 16, ErrorCode::InvalidPuzzleHash => 10, diff --git a/src/generator_rom.rs b/src/generator_rom.rs new file mode 100644 index 000000000..cee39f24e --- /dev/null +++ b/src/generator_rom.rs @@ -0,0 +1,55 @@ +// the generator ROM from: +// https://github.com/Chia-Network/chia-blockchain/blob/main/chia/wallet/puzzles/rom_bootstrap_generator.clvm.hex +pub const GENERATOR_ROM: [u8; 737] = [ + 0xff, 0x02, 0xff, 0xff, 0x01, 0xff, 0x02, 0xff, 0x0c, 0xff, 0xff, 0x04, 0xff, 0x02, 0xff, 0xff, + 0x04, 0xff, 0xff, 0x02, 0xff, 0x05, 0xff, 0xff, 0x04, 0xff, 0x08, 0xff, 0xff, 0x04, 0xff, 0x13, + 0xff, 0x80, 0x80, 0x80, 0x80, 0xff, 0x80, 0x80, 0x80, 0x80, 0xff, 0xff, 0x04, 0xff, 0xff, 0x01, + 0xff, 0xff, 0xff, 0x02, 0xff, 0xff, 0x01, 0xff, 0x05, 0xff, 0xff, 0x02, 0xff, 0x3e, 0xff, 0xff, + 0x04, 0xff, 0x02, 0xff, 0xff, 0x04, 0xff, 0x05, 0xff, 0x80, 0x80, 0x80, 0x80, 0x80, 0xff, 0xff, + 0x04, 0xff, 0xff, 0x01, 0xff, 0xff, 0xff, 0x81, 0xff, 0x7f, 0xff, 0x81, 0xdf, 0x81, 0xbf, 0xff, + 0xff, 0xff, 0x02, 0xff, 0xff, 0x03, 0xff, 0xff, 0x09, 0xff, 0x0b, 0xff, 0xff, 0x01, 0x81, 0x80, + 0x80, 0xff, 0xff, 0x01, 0xff, 0x04, 0xff, 0x80, 0xff, 0xff, 0x04, 0xff, 0x05, 0xff, 0x80, 0x80, + 0x80, 0xff, 0xff, 0x01, 0xff, 0x02, 0xff, 0xff, 0x03, 0xff, 0xff, 0x0a, 0xff, 0x0b, 0xff, 0x18, + 0x80, 0xff, 0xff, 0x01, 0xff, 0x02, 0xff, 0x1a, 0xff, 0xff, 0x04, 0xff, 0x02, 0xff, 0xff, 0x04, + 0xff, 0xff, 0x02, 0xff, 0xff, 0x03, 0xff, 0xff, 0x0a, 0xff, 0x0b, 0xff, 0x1c, 0x80, 0xff, 0xff, + 0x01, 0xff, 0x02, 0xff, 0xff, 0x03, 0xff, 0xff, 0x0a, 0xff, 0x0b, 0xff, 0x14, 0x80, 0xff, 0xff, + 0x01, 0xff, 0x08, 0x80, 0xff, 0xff, 0x01, 0xff, 0x04, 0xff, 0xff, 0x0e, 0xff, 0xff, 0x18, 0xff, + 0xff, 0x01, 0x1f, 0xff, 0x0b, 0x80, 0xff, 0xff, 0x0c, 0xff, 0x05, 0xff, 0x80, 0xff, 0xff, 0x01, + 0x01, 0x80, 0x80, 0xff, 0xff, 0x04, 0xff, 0xff, 0x0c, 0xff, 0x05, 0xff, 0xff, 0x01, 0x01, 0x80, + 0xff, 0x80, 0x80, 0x80, 0x80, 0xff, 0x01, 0x80, 0xff, 0xff, 0x01, 0xff, 0x04, 0xff, 0xff, 0x18, + 0xff, 0xff, 0x01, 0x3f, 0xff, 0x0b, 0x80, 0xff, 0xff, 0x04, 0xff, 0x05, 0xff, 0x80, 0x80, 0x80, + 0x80, 0xff, 0x01, 0x80, 0xff, 0x80, 0x80, 0x80, 0x80, 0xff, 0xff, 0x01, 0xff, 0x04, 0xff, 0x0b, + 0xff, 0xff, 0x04, 0xff, 0x05, 0xff, 0x80, 0x80, 0x80, 0x80, 0xff, 0x01, 0x80, 0x80, 0xff, 0x01, + 0x80, 0xff, 0x04, 0xff, 0xff, 0x0c, 0xff, 0x15, 0xff, 0x80, 0xff, 0x09, 0x80, 0xff, 0xff, 0x04, + 0xff, 0xff, 0x0c, 0xff, 0x15, 0xff, 0x09, 0x80, 0xff, 0x80, 0x80, 0x80, 0xff, 0xff, 0x04, 0xff, + 0xff, 0x04, 0xff, 0x05, 0xff, 0x13, 0x80, 0xff, 0xff, 0x04, 0xff, 0x2b, 0xff, 0x80, 0x80, 0x80, + 0xff, 0xff, 0x02, 0xff, 0x16, 0xff, 0xff, 0x04, 0xff, 0x02, 0xff, 0xff, 0x04, 0xff, 0x09, 0xff, + 0xff, 0x04, 0xff, 0xff, 0x02, 0xff, 0x3e, 0xff, 0xff, 0x04, 0xff, 0x02, 0xff, 0xff, 0x04, 0xff, + 0x15, 0xff, 0x80, 0x80, 0x80, 0x80, 0xff, 0x80, 0x80, 0x80, 0x80, 0x80, 0xff, 0x02, 0xff, 0xff, + 0x03, 0xff, 0xff, 0x09, 0xff, 0xff, 0x0c, 0xff, 0x05, 0xff, 0x80, 0xff, 0xff, 0x01, 0x01, 0x80, + 0xff, 0x10, 0x80, 0xff, 0xff, 0x01, 0xff, 0x02, 0xff, 0x2e, 0xff, 0xff, 0x04, 0xff, 0x02, 0xff, + 0xff, 0x04, 0xff, 0xff, 0x02, 0xff, 0x3e, 0xff, 0xff, 0x04, 0xff, 0x02, 0xff, 0xff, 0x04, 0xff, + 0xff, 0x0c, 0xff, 0x05, 0xff, 0xff, 0x01, 0x01, 0x80, 0xff, 0x80, 0x80, 0x80, 0x80, 0xff, 0x80, + 0x80, 0x80, 0x80, 0xff, 0xff, 0x01, 0xff, 0x02, 0xff, 0x12, 0xff, 0xff, 0x04, 0xff, 0x02, 0xff, + 0xff, 0x04, 0xff, 0xff, 0x0c, 0xff, 0x05, 0xff, 0xff, 0x01, 0x01, 0x80, 0xff, 0xff, 0x04, 0xff, + 0xff, 0x0c, 0xff, 0x05, 0xff, 0x80, 0xff, 0xff, 0x01, 0x01, 0x80, 0xff, 0x80, 0x80, 0x80, 0x80, + 0x80, 0x80, 0xff, 0x01, 0x80, 0xff, 0x01, 0x80, 0x80, 0xff, 0x04, 0xff, 0xff, 0x02, 0xff, 0x16, + 0xff, 0xff, 0x04, 0xff, 0x02, 0xff, 0xff, 0x04, 0xff, 0x09, 0xff, 0x80, 0x80, 0x80, 0x80, 0xff, + 0x0d, 0x80, 0xff, 0xff, 0x04, 0xff, 0x09, 0xff, 0xff, 0x04, 0xff, 0xff, 0x02, 0xff, 0x1e, 0xff, + 0xff, 0x04, 0xff, 0x02, 0xff, 0xff, 0x04, 0xff, 0x15, 0xff, 0x80, 0x80, 0x80, 0x80, 0xff, 0xff, + 0x04, 0xff, 0x2d, 0xff, 0xff, 0x04, 0xff, 0xff, 0x02, 0xff, 0x15, 0xff, 0x5d, 0x80, 0xff, 0x7d, + 0x80, 0x80, 0x80, 0x80, 0xff, 0xff, 0x02, 0xff, 0xff, 0x03, 0xff, 0x05, 0xff, 0xff, 0x01, 0xff, + 0x04, 0xff, 0xff, 0x02, 0xff, 0x0a, 0xff, 0xff, 0x04, 0xff, 0x02, 0xff, 0xff, 0x04, 0xff, 0x09, + 0xff, 0x80, 0x80, 0x80, 0x80, 0xff, 0xff, 0x02, 0xff, 0x16, 0xff, 0xff, 0x04, 0xff, 0x02, 0xff, + 0xff, 0x04, 0xff, 0x0d, 0xff, 0x80, 0x80, 0x80, 0x80, 0x80, 0xff, 0x80, 0x80, 0xff, 0x01, 0x80, + 0xff, 0x02, 0xff, 0xff, 0x03, 0xff, 0xff, 0x07, 0xff, 0x05, 0x80, 0xff, 0xff, 0x01, 0xff, 0x0b, + 0xff, 0xff, 0x01, 0x02, 0xff, 0xff, 0x02, 0xff, 0x1e, 0xff, 0xff, 0x04, 0xff, 0x02, 0xff, 0xff, + 0x04, 0xff, 0x09, 0xff, 0x80, 0x80, 0x80, 0x80, 0xff, 0xff, 0x02, 0xff, 0x1e, 0xff, 0xff, 0x04, + 0xff, 0x02, 0xff, 0xff, 0x04, 0xff, 0x0d, 0xff, 0x80, 0x80, 0x80, 0x80, 0x80, 0xff, 0xff, 0x01, + 0xff, 0x0b, 0xff, 0xff, 0x01, 0x01, 0xff, 0x05, 0x80, 0x80, 0xff, 0x01, 0x80, 0xff, 0x01, 0x80, + 0x80, +]; + +// constant from the main chia blockchain: +// https://github.com/Chia-Network/chia-blockchain/blob/main/chia/consensus/default_constants.py +pub const COST_PER_BYTE: u64 = 12000; diff --git a/src/lib.rs b/src/lib.rs index deed25c4a..e42fb1341 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,5 +1,6 @@ pub mod compression; pub mod gen; +pub mod generator_rom; pub mod merkle_set; #[cfg(fuzzing)] diff --git a/tests/run_gen.py b/tests/run_gen.py index 4b8e40c7f..20137b7b5 100755 --- a/tests/run_gen.py +++ b/tests/run_gen.py @@ -1,6 +1,6 @@ #!/usr/bin/env python3 -from chia_rs import run_generator +from chia_rs import run_block_generator from time import time from clvm_tools import binutils from clvm.serialize import atom_to_byte_iterator @@ -16,70 +16,27 @@ def serialize_atom(blob: bytes) -> bytes: def run_gen(fn: str, flags: int = 0, args: Optional[str] = None): - # the generator ROM from: - # https://github.com/Chia-Network/chia-blockchain/blob/main/chia/wallet/puzzles/rom_bootstrap_generator.clvm.hex - program_data = bytes.fromhex( - "ff02ffff01ff02ff0cffff04ff02ffff04ffff02ff05ffff04ff08ffff04ff13" - "ff80808080ff80808080ffff04ffff01ffffff02ffff01ff05ffff02ff3effff" - "04ff02ffff04ff05ff8080808080ffff04ffff01ffffff81ff7fff81df81bfff" - "ffff02ffff03ffff09ff0bffff01818080ffff01ff04ff80ffff04ff05ff8080" - "80ffff01ff02ffff03ffff0aff0bff1880ffff01ff02ff1affff04ff02ffff04" - "ffff02ffff03ffff0aff0bff1c80ffff01ff02ffff03ffff0aff0bff1480ffff" - "01ff0880ffff01ff04ffff0effff18ffff011fff0b80ffff0cff05ff80ffff01" - "018080ffff04ffff0cff05ffff010180ff80808080ff0180ffff01ff04ffff18" - "ffff013fff0b80ffff04ff05ff80808080ff0180ff80808080ffff01ff04ff0b" - "ffff04ff05ff80808080ff018080ff0180ff04ffff0cff15ff80ff0980ffff04" - "ffff0cff15ff0980ff808080ffff04ffff04ff05ff1380ffff04ff2bff808080" - "ffff02ff16ffff04ff02ffff04ff09ffff04ffff02ff3effff04ff02ffff04ff" - "15ff80808080ff8080808080ff02ffff03ffff09ffff0cff05ff80ffff010180" - "ff1080ffff01ff02ff2effff04ff02ffff04ffff02ff3effff04ff02ffff04ff" - "ff0cff05ffff010180ff80808080ff80808080ffff01ff02ff12ffff04ff02ff" - "ff04ffff0cff05ffff010180ffff04ffff0cff05ff80ffff010180ff80808080" - "8080ff0180ff018080ff04ffff02ff16ffff04ff02ffff04ff09ff80808080ff" - "0d80ffff04ff09ffff04ffff02ff1effff04ff02ffff04ff15ff80808080ffff" - "04ff2dffff04ffff02ff15ff5d80ff7d80808080ffff02ffff03ff05ffff01ff" - "04ffff02ff0affff04ff02ffff04ff09ff80808080ffff02ff16ffff04ff02ff" - "ff04ff0dff8080808080ff8080ff0180ff02ffff03ffff07ff0580ffff01ff0b" - "ffff0102ffff02ff1effff04ff02ffff04ff09ff80808080ffff02ff1effff04" - "ff02ffff04ff0dff8080808080ffff01ff0bffff0101ff058080ff0180ff0180" - "80") - # constants from the main chia blockchain: # https://github.com/Chia-Network/chia-blockchain/blob/main/chia/consensus/default_constants.py max_cost = 11000000000 cost_per_byte = 12000 - env_data = binutils.assemble(open(fn, "r").read()).as_bin() - - byte_cost = len(env_data) * cost_per_byte - - # we don't charge for the size of the generator ROM. However, we do charge - # cost for the operations it executes - max_cost -= len(env_data) * cost_per_byte + generator = binutils.assemble(open(fn, "r").read()).as_bin() # add the block program arguments - block_program_args = b"\x80" + block_refs = [] if args and args != "": with open(args, "r") as f: - block_ref = bytes.fromhex(f.read()) - block_program_args = b"\xff" + serialize_atom(block_ref) + block_program_args - env_data = b"\xff" + env_data + b"\xff\xff" + block_program_args + b"\x80\x80" + block_refs = [bytes.fromhex(f.read())] try: - err, result = run_generator( - program_data, - env_data, - max_cost, - flags, - ) - cost = 0 if result is None else result.cost + byte_cost - return (err, result, cost) + return run_block_generator(generator, block_refs, max_cost, flags) except Exception as e: # GENERATOR_RUNTIME_ERROR - return (117, None, 0) + return (117, None) -def print_spend_bundle_conditions(result, cost: int) -> str: +def print_spend_bundle_conditions(result) -> str: ret = "" if result.reserve_fee > 0: ret += f"RESERVE_FEE: {result.reserve_fee}\n" @@ -104,14 +61,14 @@ def print_spend_bundle_conditions(result, cost: int) -> str: ret += f" CREATE_COIN: ph: {a[0].hex()} amount: {a[1]}\n" for a in sorted(s.agg_sig_me): ret += f" AGG_SIG_ME pk: {a[0].hex()} msg: {a[1].hex()}\n" - ret += f"cost: {cost}\n" + ret += f"cost: {result.cost}\n" return ret if __name__ == "__main__": try: start_time = time() - error_code, result, cost = run_gen(sys.argv[1], + error_code, result = run_gen(sys.argv[1], 0 if len(sys.argv) < 3 else int(sys.argv[2]), None if len(sys.argv) < 4 else sys.argv[3]) run_time = time() - start_time @@ -121,7 +78,7 @@ def print_spend_bundle_conditions(result, cost: int) -> str: sys.exit(1) start_time = time() print("Spend bundle:") - print(print_spend_bundle_conditions(result, cost)) + print(print_spend_bundle_conditions(result)) print_time = time() - start_time print(f"run-time: {run_time:.2f}s") print(f"print-time: {print_time:.2f}s") diff --git a/tests/test-generators.py b/tests/test-generators.py index f5f7e4cd1..487b2adff 100755 --- a/tests/test-generators.py +++ b/tests/test-generators.py @@ -18,26 +18,26 @@ def compare_output(output, expected, title): failed = 1 -def parse_output(result, error_code, cost) -> str: +def parse_output(result, error_code) -> str: if error_code: return f"FAILED: {error_code}\n" else: - return print_spend_bundle_conditions(result, cost) + return print_spend_bundle_conditions(result) for g in sorted(glob.glob('generators/*.clvm')): print(f"{g}") sys.stdout.write("running generator...\r") start_time = perf_counter() - error_code, result, cost = run_gen(g, LIMIT_STACK, "generators/block-834752.hex") + error_code, result = run_gen(g, LIMIT_STACK, "generators/block-834752.hex") run_time = perf_counter() - start_time - output = parse_output(result, error_code, cost) + output = parse_output(result, error_code) sys.stdout.write("running generator (mempool mode) ...\r") sys.stdout.flush() start_time = perf_counter() - error_code2, result2, cost = run_gen(g, MEMPOOL_MODE, "generators/block-834752.hex") + error_code2, result2 = run_gen(g, MEMPOOL_MODE, "generators/block-834752.hex") run_time2 = perf_counter() - start_time - output2 = parse_output(result2, error_code2, cost) + output2 = parse_output(result2, error_code2) with open(g) as f: expected = f.read().split('\n', 1)[1] diff --git a/wheel/chia_rs.pyi b/wheel/chia_rs.pyi index 53ef6ba7b..c69a2f639 100644 --- a/wheel/chia_rs.pyi +++ b/wheel/chia_rs.pyi @@ -14,6 +14,10 @@ def run_generator( program: bytes, args: bytes, max_cost: int, flags: int ) -> Tuple[Optional[int], Optional[SpendBundleConditions]]: ... +def run_block_generator( + program: bytes, args: List[bytes], max_cost: int, flags: int +) -> Tuple[Optional[int], Optional[SpendBundleConditions]]: ... + COND_ARGS_NIL: int = ... NO_UNKNOWN_CONDS: int = ... STRICT_ARGS_COUNT: int = ... diff --git a/wheel/generate_type_stubs.py b/wheel/generate_type_stubs.py index 9d785ad4e..1877384b8 100644 --- a/wheel/generate_type_stubs.py +++ b/wheel/generate_type_stubs.py @@ -158,6 +158,10 @@ def run_generator( program: bytes, args: bytes, max_cost: int, flags: int ) -> Tuple[Optional[int], Optional[SpendBundleConditions]]: ... +def run_block_generator( + program: bytes, args: List[bytes], max_cost: int, flags: int +) -> Tuple[Optional[int], Optional[SpendBundleConditions]]: ... + COND_ARGS_NIL: int = ... NO_UNKNOWN_CONDS: int = ... STRICT_ARGS_COUNT: int = ... diff --git a/wheel/src/api.rs b/wheel/src/api.rs index 41c3f8da3..af79fc31c 100644 --- a/wheel/src/api.rs +++ b/wheel/src/api.rs @@ -1,5 +1,5 @@ use crate::compression; -use crate::run_generator::{PySpend, PySpendBundleConditions, __pyo3_get_function_run_generator}; +use crate::run_generator::{PySpend, PySpendBundleConditions, __pyo3_get_function_run_generator, __pyo3_get_function_run_block_generator}; use chia::gen::flags::COND_ARGS_NIL; use chia::gen::flags::NO_UNKNOWN_CONDS; use chia::gen::flags::STRICT_ARGS_COUNT; @@ -112,6 +112,7 @@ pub fn get_puzzle_and_solution_for_coin<'py>( pub fn chia_rs(py: Python, m: &PyModule) -> PyResult<()> { // generator functions m.add_function(wrap_pyfunction!(run_generator, m)?)?; + m.add_function(wrap_pyfunction!(run_block_generator, m)?)?; m.add_class::()?; m.add("ELIGIBLE_FOR_DEDUP", chia::gen::conditions::ELIGIBLE_FOR_DEDUP)?; m.add_class::()?; diff --git a/wheel/src/run_generator.rs b/wheel/src/run_generator.rs index ee5500a7c..b184044da 100644 --- a/wheel/src/run_generator.rs +++ b/wheel/src/run_generator.rs @@ -4,6 +4,7 @@ use chia_protocol::to_json_dict::ToJsonDict; use chia::gen::conditions::{parse_spends, Spend, SpendBundleConditions}; use chia::gen::validation_error::{ErrorCode, ValidationErr}; +use chia::gen::run_block_generator::run_block_generator as native_run_block_generator; use chia_protocol::bytes::{Bytes, Bytes32, Bytes48}; use clvmr::allocator::Allocator; @@ -15,6 +16,7 @@ use clvmr::run_program::run_program; use clvmr::serde::node_from_bytes; use pyo3::prelude::*; +use pyo3::types::PyList; use chia_protocol::chia_error; use chia_protocol::streamable::Streamable; @@ -100,7 +102,6 @@ fn convert_spend_bundle_conds(a: &Allocator, sb: SpendBundleConditions) -> PySpe // returns the cost of running the CLVM program along with conditions and the list of // spends #[pyfunction] -#[allow(clippy::borrow_deref_ref)] pub fn run_generator( py: Python, program: &[u8], @@ -142,7 +143,7 @@ pub fn run_generator( None, Some(convert_spend_bundle_conds(&allocator, spend_bundle_conds)), )) - } + }, Ok((error_code, _)) => { // a validation error occurred Ok((error_code.map(|x| x.into()), None)) @@ -150,3 +151,37 @@ pub fn run_generator( Err(eval_err) => eval_err_to_pyresult(py, eval_err, allocator), } } + +#[pyfunction] +pub fn run_block_generator( + _py: Python, + program: &[u8], + block_refs: &PyList, + max_cost: Cost, + flags: u32, +) -> PyResult<(Option, Option)> { + let mut allocator = if flags & LIMIT_HEAP != 0 { + Allocator::new_limited(500000000, 62500000, 62500000) + } else { + Allocator::new() + }; + + let mut refs = Vec::>::new(); + for g in block_refs { + refs.push(g.extract::>()?); + } + + match native_run_block_generator(&mut allocator, program, &refs, max_cost, flags) { + Ok(spend_bundle_conds) => { + // everything was successful + Ok(( + None, + Some(convert_spend_bundle_conds(&allocator, spend_bundle_conds)), + )) + }, + Err(ValidationErr(_, error_code)) => { + // a validation error occurred + Ok((Some(error_code.into()), None)) + } + } +}