chia_sdk_driver/
spend_context.rs

1use std::{
2    collections::HashMap,
3    ops::{Deref, DerefMut},
4};
5
6use chia_protocol::{Bytes32, Coin, CoinSpend, Program};
7use chia_sdk_types::{conditions::Memos, run_puzzle, Conditions, Mod};
8use clvm_traits::{clvm_quote, FromClvm, ToClvm};
9use clvm_utils::{tree_hash, CurriedProgram, TreeHash};
10use clvmr::{
11    serde::{node_from_bytes, node_to_bytes, node_to_bytes_backrefs},
12    Allocator, NodePtr,
13};
14
15use crate::{DriverError, HashedPtr, Spend};
16
17/// A wrapper around [`Allocator`] that caches puzzles and keeps track of a list of [`CoinSpend`].
18/// It's used to construct spend bundles in an easy and efficient way.
19#[derive(Debug, Default)]
20pub struct SpendContext {
21    allocator: Allocator,
22    puzzles: HashMap<TreeHash, NodePtr>,
23    coin_spends: Vec<CoinSpend>,
24}
25
26impl SpendContext {
27    pub fn new() -> Self {
28        Self::default()
29    }
30
31    pub fn iter(&self) -> impl Iterator<Item = &CoinSpend> {
32        self.coin_spends.iter()
33    }
34
35    /// Remove all of the [`CoinSpend`] that have been collected so far.
36    pub fn take(&mut self) -> Vec<CoinSpend> {
37        std::mem::take(&mut self.coin_spends)
38    }
39
40    /// Adds a [`CoinSpend`] to the collection.
41    pub fn insert(&mut self, coin_spend: CoinSpend) {
42        self.coin_spends.push(coin_spend);
43    }
44
45    /// Serializes a [`Spend`] and adds it to the list of [`CoinSpend`].
46    pub fn spend(&mut self, coin: Coin, spend: Spend) -> Result<(), DriverError> {
47        let puzzle_reveal = self.serialize(&spend.puzzle)?;
48        let solution = self.serialize(&spend.solution)?;
49        self.insert(CoinSpend::new(coin, puzzle_reveal, solution));
50        Ok(())
51    }
52
53    /// Allocate a new node and return its pointer.
54    pub fn alloc<T>(&mut self, value: &T) -> Result<NodePtr, DriverError>
55    where
56        T: ToClvm<Allocator>,
57    {
58        Ok(value.to_clvm(&mut self.allocator)?)
59    }
60
61    /// Allocate a new node and return its pointer pre-hashed.
62    pub fn alloc_hashed<T>(&mut self, value: &T) -> Result<HashedPtr, DriverError>
63    where
64        T: ToClvm<Allocator>,
65    {
66        let ptr = value.to_clvm(&mut self.allocator)?;
67        Ok(HashedPtr::from_ptr(self, ptr))
68    }
69
70    /// Extract a value from a node pointer.
71    pub fn extract<T>(&self, ptr: NodePtr) -> Result<T, DriverError>
72    where
73        T: FromClvm<Allocator>,
74    {
75        Ok(T::from_clvm(&self.allocator, ptr)?)
76    }
77
78    /// Compute the tree hash of a node pointer.
79    pub fn tree_hash(&self, ptr: NodePtr) -> TreeHash {
80        tree_hash(&self.allocator, ptr)
81    }
82
83    /// Run a puzzle with a solution and return the result.
84    pub fn run(&mut self, puzzle: NodePtr, solution: NodePtr) -> Result<NodePtr, DriverError> {
85        Ok(run_puzzle(&mut self.allocator, puzzle, solution)?)
86    }
87
88    /// Allocate a value and serialize it into a [`Program`].
89    pub fn serialize<T>(&mut self, value: &T) -> Result<Program, DriverError>
90    where
91        T: ToClvm<Allocator>,
92    {
93        let ptr = value.to_clvm(&mut self.allocator)?;
94        Ok(node_to_bytes(&self.allocator, ptr)?.into())
95    }
96
97    /// Allocate a value and serialize it into a [`Program`] with back references enabled.
98    pub fn serialize_with_backrefs<T>(&mut self, value: &T) -> Result<Program, DriverError>
99    where
100        T: ToClvm<Allocator>,
101    {
102        let ptr = value.to_clvm(&mut self.allocator)?;
103        Ok(node_to_bytes_backrefs(&self.allocator, ptr)?.into())
104    }
105
106    pub fn memos<T>(&mut self, value: &T) -> Result<Memos<NodePtr>, DriverError>
107    where
108        T: ToClvm<Allocator>,
109    {
110        Ok(Memos::Some(self.alloc(value)?))
111    }
112
113    pub fn hint(&mut self, hint: Bytes32) -> Result<Memos<NodePtr>, DriverError> {
114        self.memos(&[hint])
115    }
116
117    pub fn alloc_mod<T>(&mut self) -> Result<NodePtr, DriverError>
118    where
119        T: Mod,
120    {
121        self.puzzle(T::mod_hash(), T::mod_reveal().as_ref())
122    }
123
124    pub fn curry<T>(&mut self, args: T) -> Result<NodePtr, DriverError>
125    where
126        T: Mod + ToClvm<Allocator>,
127    {
128        let mod_ptr = self.alloc_mod::<T>()?;
129        self.alloc(&CurriedProgram {
130            program: mod_ptr,
131            args,
132        })
133    }
134
135    pub fn get_puzzle(&self, puzzle_hash: &TreeHash) -> Option<NodePtr> {
136        self.puzzles.get(puzzle_hash).copied()
137    }
138
139    pub fn puzzle(
140        &mut self,
141        puzzle_hash: TreeHash,
142        puzzle_bytes: &[u8],
143    ) -> Result<NodePtr, DriverError> {
144        if let Some(puzzle) = self.puzzles.get(&puzzle_hash) {
145            Ok(*puzzle)
146        } else {
147            let puzzle = node_from_bytes(&mut self.allocator, puzzle_bytes)?;
148            self.puzzles.insert(puzzle_hash, puzzle);
149            Ok(puzzle)
150        }
151    }
152
153    pub fn delegated_spend(&mut self, conditions: Conditions) -> Result<Spend, DriverError> {
154        let puzzle = self.alloc(&clvm_quote!(conditions))?;
155        Ok(Spend::new(puzzle, NodePtr::NIL))
156    }
157}
158
159impl Deref for SpendContext {
160    type Target = Allocator;
161
162    fn deref(&self) -> &Self::Target {
163        &self.allocator
164    }
165}
166
167impl DerefMut for SpendContext {
168    fn deref_mut(&mut self) -> &mut Self::Target {
169        &mut self.allocator
170    }
171}
172
173impl IntoIterator for SpendContext {
174    type Item = CoinSpend;
175    type IntoIter = std::vec::IntoIter<Self::Item>;
176
177    fn into_iter(self) -> Self::IntoIter {
178        self.coin_spends.into_iter()
179    }
180}
181
182impl From<Allocator> for SpendContext {
183    fn from(allocator: Allocator) -> Self {
184        Self {
185            allocator,
186            puzzles: HashMap::new(),
187            coin_spends: Vec::new(),
188        }
189    }
190}
chia_sdk_driver/spend_context.rs

chia_sdk_driver/
spend_context.rs
