action_layer_driver/puzzle.rs
1//! Puzzle loading and currying utilities
2//!
3//! The key insight is that callers should define their own curry args types with `#[clvm(curry)]`.
4//! This module provides utilities to work with those types via `CurriedProgram`.
5
6use chia_wallet_sdk::driver::SpendContext;
7use clvm_traits::ToClvm;
8use clvm_utils::{CurriedProgram, TreeHash};
9use clvmr::serde::node_from_bytes;
10use clvmr::{Allocator, NodePtr};
11
12use crate::DriverError;
13
14/// A compiled puzzle module that can be curried with typed arguments.
15///
16/// # Example
17///
18/// ```rust,ignore
19/// // Define curry args with #[clvm(curry)]
20/// #[derive(Debug, Clone, ToClvm, FromClvm)]
21/// #[clvm(curry)]
22/// pub struct MyCurryArgs {
23/// pub inner_puzzle_hash: Bytes32,
24/// }
25///
26/// // Load puzzle
27/// let puzzle = PuzzleModule::from_hex(MY_PUZZLE_HEX)?;
28///
29/// // Compute curried hash using CurriedProgram directly
30/// let args = MyCurryArgs { inner_puzzle_hash: some_hash };
31/// let curried_hash = CurriedProgram {
32/// program: puzzle.mod_hash(),
33/// args: args.clone(),
34/// }.tree_hash();
35///
36/// // Build curried puzzle for spending
37/// let curried_ptr = puzzle.curry_puzzle(ctx, args)?;
38/// ```
39#[derive(Clone)]
40pub struct PuzzleModule {
41 mod_hash: TreeHash,
42 mod_bytes: Vec<u8>,
43}
44
45impl PuzzleModule {
46 /// Create a PuzzleModule from raw puzzle bytes
47 pub fn new(bytes: Vec<u8>) -> Result<Self, DriverError> {
48 let mut alloc = Allocator::new();
49 let ptr = node_from_bytes(&mut alloc, &bytes)
50 .map_err(|e| DriverError::PuzzleParse(format!("{:?}", e)))?;
51 let mod_hash = chia::clvm_utils::tree_hash(&alloc, ptr);
52 Ok(Self {
53 mod_hash,
54 mod_bytes: bytes,
55 })
56 }
57
58 /// Create a PuzzleModule from a hex string (with optional whitespace)
59 pub fn from_hex(hex_str: &str) -> Result<Self, DriverError> {
60 let bytes = hex::decode(hex_str.trim())
61 .map_err(|e| DriverError::PuzzleParse(format!("Invalid hex: {}", e)))?;
62 Self::new(bytes)
63 }
64
65 /// Get the module hash (uncurried puzzle hash)
66 pub fn mod_hash(&self) -> TreeHash {
67 self.mod_hash
68 }
69
70 /// Get the raw puzzle bytes
71 pub fn mod_bytes(&self) -> &[u8] {
72 &self.mod_bytes
73 }
74
75 /// Build a curried puzzle in the spend context.
76 ///
77 /// The args type should have `#[clvm(curry)]` attribute for proper serialization.
78 ///
79 /// # Example
80 ///
81 /// ```rust,ignore
82 /// #[derive(ToClvm)]
83 /// #[clvm(curry)]
84 /// struct MyCurryArgs { pub hash: Bytes32 }
85 ///
86 /// let curried = puzzle.curry_puzzle(ctx, MyCurryArgs { hash })?;
87 /// ```
88 pub fn curry_puzzle<A>(&self, ctx: &mut SpendContext, args: A) -> Result<NodePtr, DriverError>
89 where
90 A: ToClvm<Allocator>,
91 {
92 let mod_ptr = ctx
93 .puzzle(self.mod_hash, &self.mod_bytes)
94 .map_err(|e| DriverError::PuzzleLoad(format!("{:?}", e)))?;
95
96 ctx.alloc(&CurriedProgram {
97 program: mod_ptr,
98 args,
99 })
100 .map_err(|e| DriverError::Alloc(format!("{:?}", e)))
101 }
102}
103
104impl std::fmt::Debug for PuzzleModule {
105 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
106 f.debug_struct("PuzzleModule")
107 .field("mod_hash", &hex::encode(self.mod_hash.to_bytes()))
108 .field("mod_bytes_len", &self.mod_bytes.len())
109 .finish()
110 }
111}