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}