1use std::{collections::HashMap, fmt::Debug};
2
3use chia_protocol::Bytes32;
4use chia_sdk_types::{
5 puzzles::{
6 ActionLayerArgs, DefaultFinalizer1stCurryArgs, DefaultFinalizer2ndCurryArgs,
7 RawActionLayerSolution, ReserveFinalizer1stCurryArgs, ReserveFinalizer2ndCurryArgs,
8 ACTION_LAYER_PUZZLE_HASH, DEFAULT_FINALIZER_PUZZLE_HASH, RESERVE_FINALIZER_PUZZLE_HASH,
9 },
10 run_puzzle, MerkleProof, MerkleTree, Mod,
11};
12use clvm_traits::{clvm_list, match_tuple, FromClvm, ToClvm};
13use clvm_utils::{tree_hash, CurriedProgram, TreeHash};
14use clvmr::{Allocator, NodePtr};
15
16use crate::{DriverError, Layer, Puzzle, Spend, SpendContext};
17
18#[derive(Debug, Clone, PartialEq, Eq)]
19pub enum Finalizer<P> {
20 Default {
21 hint: Bytes32,
22 },
23 Reserve {
24 reserve_full_puzzle_hash: Bytes32,
25 reserve_inner_puzzle_hash: Bytes32,
26 reserve_amount_from_state_program: P,
27 hint: Bytes32,
28 },
29}
30
31#[derive(Debug, Clone, PartialEq, Eq)]
32pub struct ActionLayer<S, P = ()> {
33 pub merkle_root: Bytes32,
34 pub state: S,
35 pub finalizer: Finalizer<P>,
36}
37
38#[derive(Debug, Clone)]
39pub struct ActionLayerSolution<F> {
40 pub proofs: Vec<MerkleProof>,
41 pub action_spends: Vec<Spend>,
42 pub finalizer_solution: F,
43}
44
45impl<S, P> ActionLayer<S, P> {
46 pub fn new(merkle_root: Bytes32, state: S, finalizer: Finalizer<P>) -> Self {
47 Self {
48 merkle_root,
49 state,
50 finalizer,
51 }
52 }
53
54 pub fn from_action_puzzle_hashes(
55 leaves: &[Bytes32],
56 state: S,
57 finalizer: Finalizer<P>,
58 ) -> Self {
59 let merkle_root = MerkleTree::new(leaves).root();
60
61 Self {
62 merkle_root,
63 state,
64 finalizer,
65 }
66 }
67
68 pub fn get_proofs(
69 &self,
70 action_puzzle_hashes: &[Bytes32],
71 action_spends_puzzle_hashes: &[Bytes32],
72 ) -> Option<Vec<MerkleProof>> {
73 let merkle_tree = MerkleTree::new(action_puzzle_hashes);
74
75 let proofs = action_spends_puzzle_hashes
76 .iter()
77 .filter_map(|puzzle_hash| {
78 let proof = merkle_tree.proof(*puzzle_hash)?;
79
80 Some(proof)
81 })
82 .collect::<Vec<_>>();
83
84 if proofs.len() != action_spends_puzzle_hashes.len() {
85 return None;
86 }
87
88 Some(proofs)
89 }
90
91 pub fn extract_merkle_root_and_state(
92 allocator: &Allocator,
93 inner_puzzle: Puzzle,
94 ) -> Result<Option<(Bytes32, S)>, DriverError>
95 where
96 S: FromClvm<Allocator>,
97 {
98 let Some(puzzle) = inner_puzzle.as_curried() else {
99 return Ok(None);
100 };
101
102 if inner_puzzle.mod_hash() != ACTION_LAYER_PUZZLE_HASH {
103 return Ok(None);
104 }
105
106 let args = ActionLayerArgs::<NodePtr, S>::from_clvm(allocator, puzzle.args)?;
107
108 Ok(Some((args.merkle_root, args.state)))
109 }
110
111 pub fn get_new_state(
112 allocator: &mut Allocator,
113 initial_state: S,
114 action_layer_solution: NodePtr,
115 ) -> Result<S, DriverError>
116 where
117 S: ToClvm<Allocator> + FromClvm<Allocator> + Clone,
118 {
119 let solution = ActionLayer::<S, NodePtr>::parse_solution(allocator, action_layer_solution)?;
120
121 let mut state_incl_ephemeral: (NodePtr, S) = (NodePtr::NIL, initial_state);
122 for raw_action in solution.action_spends {
123 let actual_solution =
124 clvm_list!(state_incl_ephemeral, raw_action.solution).to_clvm(allocator)?;
125
126 let output = run_puzzle(allocator, raw_action.puzzle, actual_solution)?;
127
128 (state_incl_ephemeral, _) =
129 <match_tuple!((NodePtr, S), NodePtr)>::from_clvm(allocator, output)?;
130 }
131
132 Ok(state_incl_ephemeral.1)
133 }
134}
135
136impl<S, P> Layer for ActionLayer<S, P>
137where
138 S: ToClvm<Allocator> + FromClvm<Allocator> + Clone,
139 P: ToClvm<Allocator> + FromClvm<Allocator> + Clone,
140{
141 type Solution = ActionLayerSolution<NodePtr>;
142
143 fn parse_puzzle(allocator: &Allocator, puzzle: Puzzle) -> Result<Option<Self>, DriverError> {
144 let Some(puzzle) = puzzle.as_curried() else {
145 return Ok(None);
146 };
147
148 if puzzle.mod_hash != ACTION_LAYER_PUZZLE_HASH {
149 return Ok(None);
150 }
151
152 let args = ActionLayerArgs::<NodePtr, S>::from_clvm(allocator, puzzle.args)?;
153 let finalizer_2nd_curry =
154 CurriedProgram::<NodePtr, NodePtr>::from_clvm(allocator, args.finalizer);
155 let Ok(finalizer_2nd_curry) = finalizer_2nd_curry else {
156 return Ok(None);
157 };
158
159 let finalizer_1st_curry = Puzzle::from_clvm(allocator, finalizer_2nd_curry.program)?;
160 let Some(finalizer_1st_curry) = finalizer_1st_curry.as_curried() else {
161 return Ok(None);
162 };
163
164 match finalizer_1st_curry.mod_hash {
165 DEFAULT_FINALIZER_PUZZLE_HASH => {
166 let finalizer_2nd_curry_args =
167 DefaultFinalizer2ndCurryArgs::from_clvm(allocator, finalizer_2nd_curry.args)?;
168 let finalizer_1st_curry_args =
169 DefaultFinalizer1stCurryArgs::from_clvm(allocator, finalizer_1st_curry.args)?;
170
171 let expected_self_hash = DefaultFinalizer1stCurryArgs {
172 action_layer_mod_hash: ACTION_LAYER_PUZZLE_HASH.into(),
173 hint: finalizer_1st_curry_args.hint,
174 }
175 .curry_tree_hash()
176 .into();
177 if finalizer_1st_curry.mod_hash != DEFAULT_FINALIZER_PUZZLE_HASH
178 || finalizer_1st_curry_args.action_layer_mod_hash
179 != ACTION_LAYER_PUZZLE_HASH.into()
180 || finalizer_2nd_curry_args.finalizer_self_hash != expected_self_hash
181 {
182 return Err(DriverError::NonStandardLayer);
183 }
184
185 Ok(Some(Self {
186 merkle_root: args.merkle_root,
187 state: args.state,
188 finalizer: Finalizer::Default {
189 hint: finalizer_1st_curry_args.hint,
190 },
191 }))
192 }
193 RESERVE_FINALIZER_PUZZLE_HASH => {
194 let finalizer_2nd_curry_args =
195 ReserveFinalizer2ndCurryArgs::from_clvm(allocator, finalizer_2nd_curry.args)?;
196 let finalizer_1st_curry_args = ReserveFinalizer1stCurryArgs::<NodePtr>::from_clvm(
197 allocator,
198 finalizer_1st_curry.args,
199 )?;
200
201 let reserve_amount_from_state_program_hash = tree_hash(
202 allocator,
203 finalizer_1st_curry_args.reserve_amount_from_state_program,
204 );
205
206 if finalizer_1st_curry.mod_hash != RESERVE_FINALIZER_PUZZLE_HASH
207 || finalizer_1st_curry_args.action_layer_mod_hash
208 != ACTION_LAYER_PUZZLE_HASH.into()
209 || finalizer_2nd_curry_args.finalizer_self_hash
210 != ReserveFinalizer1stCurryArgs::<TreeHash>::curry_tree_hash(
211 finalizer_1st_curry_args.reserve_full_puzzle_hash,
212 finalizer_1st_curry_args.reserve_inner_puzzle_hash,
213 reserve_amount_from_state_program_hash,
214 finalizer_1st_curry_args.hint,
215 )
216 .into()
217 {
218 return Err(DriverError::NonStandardLayer);
219 }
220
221 let reserve_amount_from_state_program = <P>::from_clvm(
222 allocator,
223 finalizer_1st_curry_args.reserve_amount_from_state_program,
224 )?;
225
226 Ok(Some(Self {
227 merkle_root: args.merkle_root,
228 state: args.state,
229 finalizer: Finalizer::Reserve {
230 reserve_full_puzzle_hash: finalizer_1st_curry_args.reserve_full_puzzle_hash,
231 reserve_inner_puzzle_hash: finalizer_1st_curry_args
232 .reserve_inner_puzzle_hash,
233 reserve_amount_from_state_program,
234 hint: finalizer_1st_curry_args.hint,
235 },
236 }))
237 }
238 _ => Err(DriverError::NonStandardLayer),
239 }
240 }
241
242 fn parse_solution(
243 allocator: &Allocator,
244 solution: NodePtr,
245 ) -> Result<Self::Solution, DriverError> {
246 let solution =
247 RawActionLayerSolution::<NodePtr, NodePtr, NodePtr>::from_clvm(allocator, solution)?;
248
249 let mut actions = Vec::<NodePtr>::with_capacity(solution.solutions.len());
250 let mut proofs = Vec::<MerkleProof>::with_capacity(solution.solutions.len());
251 let mut selector_proofs = HashMap::<u32, MerkleProof>::new();
252
253 for (selector, proof) in solution.selectors_and_proofs {
254 let proof = if let Some(existing_proof) = selector_proofs.get(&selector) {
255 existing_proof.clone()
256 } else {
257 let proof = proof.ok_or(DriverError::InvalidMerkleProof)?;
258 selector_proofs.insert(selector, proof.clone());
259 proof
260 };
261
262 proofs.push(proof);
263
264 let mut index = 0;
265 let mut remaining_selector = selector;
266 while remaining_selector > 2 {
267 index += 1;
268 remaining_selector /= 2;
269 }
270 actions.push(solution.puzzles[index]);
271 }
272
273 let action_spends = solution
274 .solutions
275 .iter()
276 .zip(actions.into_iter().rev())
277 .map(|(action_solution, action_puzzle)| Spend::new(action_puzzle, *action_solution))
278 .collect();
279 let proofs = proofs.into_iter().rev().collect();
280
281 Ok(ActionLayerSolution {
282 proofs,
283 action_spends,
284 finalizer_solution: solution.finalizer_solution,
285 })
286 }
287
288 fn construct_puzzle(&self, ctx: &mut SpendContext) -> Result<NodePtr, DriverError> {
289 let finalizer_1st_curry = match &self.finalizer {
290 Finalizer::Default { hint } => ctx.curry(DefaultFinalizer1stCurryArgs::new(*hint))?,
291 Finalizer::Reserve {
292 reserve_full_puzzle_hash,
293 reserve_inner_puzzle_hash,
294 reserve_amount_from_state_program,
295 hint,
296 } => ctx.curry(ReserveFinalizer1stCurryArgs::<P>::new(
297 *reserve_full_puzzle_hash,
298 *reserve_inner_puzzle_hash,
299 reserve_amount_from_state_program.clone(),
300 *hint,
301 ))?,
302 };
303
304 let finalizer = match &self.finalizer {
305 Finalizer::Default { hint } => CurriedProgram {
306 program: finalizer_1st_curry,
307 args: DefaultFinalizer2ndCurryArgs::new(*hint),
308 }
309 .to_clvm(ctx)?,
310 Finalizer::Reserve {
311 reserve_full_puzzle_hash,
312 reserve_inner_puzzle_hash,
313 reserve_amount_from_state_program,
314 hint,
315 } => {
316 let reserve_amount_from_state_program =
317 ctx.alloc(&reserve_amount_from_state_program)?;
318 let reserve_amount_from_state_program_hash =
319 ctx.tree_hash(reserve_amount_from_state_program);
320
321 CurriedProgram {
322 program: finalizer_1st_curry,
323 args: ReserveFinalizer2ndCurryArgs::new(
324 *reserve_full_puzzle_hash,
325 *reserve_inner_puzzle_hash,
326 &reserve_amount_from_state_program_hash,
327 *hint,
328 ),
329 }
330 .to_clvm(ctx)?
331 }
332 };
333
334 ctx.curry(ActionLayerArgs::<NodePtr, S>::new(
335 finalizer,
336 self.merkle_root,
337 self.state.clone(),
338 ))
339 }
340
341 fn construct_solution(
342 &self,
343 ctx: &mut SpendContext,
344 solution: Self::Solution,
345 ) -> Result<NodePtr, DriverError> {
346 let mut puzzle_to_selector = HashMap::<Bytes32, u32>::new();
347 let mut next_selector = 2;
348
349 let mut puzzles = Vec::new();
350 let mut selectors_and_proofs = Vec::new();
351 let mut solutions = Vec::new();
352
353 for (spend, proof) in solution.action_spends.into_iter().zip(solution.proofs) {
354 let puzzle_hash = ctx.tree_hash(spend.puzzle).into();
355 if let Some(selector) = puzzle_to_selector.get(&puzzle_hash) {
356 selectors_and_proofs.push((*selector, Some(proof.clone())));
357 } else {
358 puzzles.push(spend.puzzle);
359 selectors_and_proofs.push((next_selector, Some(proof.clone())));
360 puzzle_to_selector.insert(puzzle_hash, next_selector);
361
362 next_selector = next_selector * 2 + 1;
363 }
364
365 solutions.push(spend.solution);
366 }
367
368 let mut proven_selectors = Vec::<u32>::new();
369 let mut selectors_and_proofs: Vec<(u32, Option<MerkleProof>)> =
370 selectors_and_proofs.into_iter().rev().collect();
371 #[allow(clippy::needless_range_loop)]
372 for i in 0..selectors_and_proofs.len() {
373 let selector = selectors_and_proofs[i].0;
374
375 if proven_selectors.contains(&selector) {
376 selectors_and_proofs[i].1 = None;
377 } else {
378 proven_selectors.push(selector);
379 }
380 }
381
382 Ok(RawActionLayerSolution {
383 puzzles,
384 selectors_and_proofs,
385 solutions,
386 finalizer_solution: solution.finalizer_solution,
387 }
388 .to_clvm(ctx)?)
389 }
390}