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