chia_sdk_driver/
puzzle.rs1use 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}