chia_sdk_driver/
puzzle.rs

1use clvm_traits::{FromClvm, FromClvmError};
2use clvm_utils::{tree_hash, CurriedProgram, ToTreeHash, TreeHash};
3use clvmr::{Allocator, NodePtr};
4
5use crate::{DriverError, Layer, SpendContext};
6
7#[derive(Debug, Clone, Copy)]
8pub enum Puzzle {
9    Curried(CurriedPuzzle),
10    Raw(RawPuzzle),
11}
12
13impl Puzzle {
14    pub fn parse(allocator: &Allocator, puzzle: NodePtr) -> Self {
15        CurriedPuzzle::parse(allocator, puzzle).map_or_else(
16            || {
17                Self::Raw(RawPuzzle {
18                    puzzle_hash: tree_hash(allocator, puzzle),
19                    ptr: puzzle,
20                })
21            },
22            Self::Curried,
23        )
24    }
25
26    pub fn curried_puzzle_hash(&self) -> TreeHash {
27        match self {
28            Self::Curried(curried) => curried.curried_puzzle_hash,
29            Self::Raw(raw) => raw.puzzle_hash,
30        }
31    }
32
33    pub fn mod_hash(&self) -> TreeHash {
34        match self {
35            Self::Curried(curried) => curried.mod_hash,
36            Self::Raw(raw) => raw.puzzle_hash,
37        }
38    }
39
40    pub fn ptr(&self) -> NodePtr {
41        match self {
42            Self::Curried(curried) => curried.curried_ptr,
43            Self::Raw(raw) => raw.ptr,
44        }
45    }
46
47    pub fn is_curried(&self) -> bool {
48        matches!(self, Self::Curried(_))
49    }
50
51    pub fn is_raw(&self) -> bool {
52        matches!(self, Self::Raw(_))
53    }
54
55    pub fn as_curried(&self) -> Option<CurriedPuzzle> {
56        match self {
57            Self::Curried(curried) => Some(*curried),
58            Self::Raw(_raw) => None,
59        }
60    }
61
62    pub fn as_raw(&self) -> Option<RawPuzzle> {
63        match self {
64            Self::Curried(_curried) => None,
65            Self::Raw(raw) => Some(*raw),
66        }
67    }
68}
69
70impl PartialEq for Puzzle {
71    fn eq(&self, other: &Self) -> bool {
72        self.curried_puzzle_hash() == other.curried_puzzle_hash()
73    }
74}
75
76impl Eq for Puzzle {}
77
78impl FromClvm<Allocator> for Puzzle {
79    fn from_clvm(allocator: &Allocator, puzzle: NodePtr) -> Result<Self, FromClvmError> {
80        Ok(Self::parse(allocator, puzzle))
81    }
82}
83
84#[derive(Debug, Clone, Copy)]
85pub struct CurriedPuzzle {
86    pub curried_puzzle_hash: TreeHash,
87    pub curried_ptr: NodePtr,
88    pub mod_hash: TreeHash,
89    pub args: NodePtr,
90}
91
92impl CurriedPuzzle {
93    pub fn parse(allocator: &Allocator, puzzle: NodePtr) -> Option<Self> {
94        let curried = CurriedProgram::from_clvm(allocator, puzzle).ok()?;
95        let mod_hash = tree_hash(allocator, curried.program);
96        let curried_puzzle_hash = CurriedProgram {
97            program: mod_hash,
98            args: tree_hash(allocator, curried.args),
99        }
100        .tree_hash();
101
102        Some(Self {
103            curried_puzzle_hash,
104            curried_ptr: puzzle,
105            mod_hash,
106            args: curried.args,
107        })
108    }
109}
110
111#[derive(Debug, Clone, Copy)]
112pub struct RawPuzzle {
113    pub puzzle_hash: TreeHash,
114    pub ptr: NodePtr,
115}
116
117impl ToTreeHash for Puzzle {
118    fn tree_hash(&self) -> TreeHash {
119        self.curried_puzzle_hash()
120    }
121}
122
123impl Layer for Puzzle {
124    type Solution = NodePtr;
125
126    fn parse_puzzle(_allocator: &Allocator, puzzle: Puzzle) -> Result<Option<Self>, DriverError>
127    where
128        Self: Sized,
129    {
130        Ok(Some(puzzle))
131    }
132
133    fn parse_solution(
134        _allocator: &Allocator,
135        solution: NodePtr,
136    ) -> Result<Self::Solution, DriverError> {
137        Ok(solution)
138    }
139
140    fn construct_puzzle(&self, _ctx: &mut SpendContext) -> Result<NodePtr, DriverError> {
141        Ok(self.ptr())
142    }
143
144    fn construct_solution(
145        &self,
146        _ctx: &mut SpendContext,
147        solution: Self::Solution,
148    ) -> Result<NodePtr, DriverError> {
149        Ok(solution)
150    }
151}