1use chia_protocol::{Bytes, Bytes32};
2use chia_puzzle_types::Memos;
3use chia_puzzle_types::singleton::{LauncherSolution, SingletonArgs, SingletonStruct};
4use chia_sdk_types::Condition;
5use chia_sdk_types::puzzles::StateSchedulerLayerArgs;
6use clvm_traits::{FromClvm, ToClvm, clvm_quote};
7use clvm_utils::{ToTreeHash, TreeHash};
8use clvmr::{Allocator, NodePtr};
9
10use crate::{
11 DriverError, SingletonLayer, StateSchedulerLayer, XchandlesRegistryReceivedMessagePrefix,
12};
13
14#[derive(Debug, Clone, PartialEq, Eq)]
15pub struct StateSchedulerInfo<S> {
16 pub launcher_id: Bytes32,
17
18 pub receiver_singleton_launcher_id: Bytes32,
19 pub state_schedule: Vec<(u64, S)>,
21 pub generation: usize,
22 pub final_puzzle_hash: Bytes32,
23}
24
25impl<S> StateSchedulerInfo<S>
26where
27 S: ToTreeHash + Clone,
28{
29 pub fn new(
30 launcher_id: Bytes32,
31 receiver_singleton_launcher_id: Bytes32,
32 state_schedule: Vec<(u64, S)>,
33 generation: usize,
34 final_puzzle_hash: Bytes32,
35 ) -> Result<Self, DriverError> {
36 validate_state_schedule(&state_schedule)?;
37
38 Ok(Self {
39 launcher_id,
40 receiver_singleton_launcher_id,
41 state_schedule,
42 generation,
43 final_puzzle_hash,
44 })
45 }
46
47 #[must_use]
48 pub fn with_generation(&self, generation: usize) -> Self {
49 Self {
50 generation,
51 ..self.clone()
52 }
53 }
54
55 pub fn inner_puzzle_hash_for(
56 &self,
57 next_puzzle_hash: Bytes32,
58 required_timestamp: u64,
59 prefix_and_message_hash: TreeHash,
60 ) -> TreeHash {
61 StateSchedulerLayerArgs::<TreeHash, _>::curry_tree_hash(
62 SingletonStruct::new(self.receiver_singleton_launcher_id)
63 .tree_hash()
64 .into(),
65 prefix_and_message_hash,
66 &clvm_quote!(vec![
67 Condition::<()>::create_coin(next_puzzle_hash, 1, Memos::None),
68 Condition::assert_seconds_absolute(required_timestamp),
69 ]),
70 )
71 }
72
73 pub fn inner_puzzle_hash_for_generation(&self, generation: usize) -> TreeHash {
74 if generation >= self.state_schedule.len() {
75 return self.final_puzzle_hash.into();
76 }
77
78 let mut inner_puzzle_hash: TreeHash = self.final_puzzle_hash.into();
79
80 let mut i = self.state_schedule.len();
81 while i > generation {
82 let prefix_and_message_hash: Bytes =
83 XchandlesRegistryReceivedMessagePrefix::update_state(
84 self.state_schedule[i - 1].1.tree_hash(),
85 )
86 .into();
87 inner_puzzle_hash = self.inner_puzzle_hash_for(
88 inner_puzzle_hash.into(),
89 self.state_schedule[i - 1].0,
90 prefix_and_message_hash.tree_hash(),
91 );
92
93 i -= 1;
94 }
95
96 inner_puzzle_hash
97 }
98
99 pub fn inner_puzzle_hash(&self) -> TreeHash {
100 self.inner_puzzle_hash_for_generation(self.generation)
101 }
102
103 pub fn into_layers(self) -> SingletonLayer<StateSchedulerLayer> {
104 let (required_timestamp, new_state) = self.state_schedule[self.generation].clone();
105
106 SingletonLayer::new(
107 self.launcher_id,
108 StateSchedulerLayer::new(
109 SingletonStruct::new(self.receiver_singleton_launcher_id)
110 .tree_hash()
111 .into(),
112 new_state.tree_hash().into(),
113 required_timestamp,
114 self.inner_puzzle_hash_for_generation(self.generation + 1)
115 .into(),
116 ),
117 )
118 }
119
120 pub fn from_launcher_solution<H>(
121 allocator: &mut Allocator,
122 laucher_solution: LauncherSolution<NodePtr>,
123 ) -> Result<Option<(Self, H)>, DriverError>
124 where
125 S: FromClvm<Allocator>,
126 H: FromClvm<Allocator>,
127 {
128 let hints = StateSchedulerLauncherHints::<S, H>::from_clvm(
129 allocator,
130 laucher_solution.key_value_list,
131 )?;
132
133 let candidate = Self::new(
134 hints.my_launcher_id,
135 hints.receiver_singleton_launcher_id,
136 hints.state_schedule,
137 0,
138 hints.final_puzzle_hash,
139 )?;
140
141 let predicted_inner_puzzle_hash = candidate.inner_puzzle_hash();
142 let predicted_puzzle_hash =
143 SingletonArgs::curry_tree_hash(hints.my_launcher_id, predicted_inner_puzzle_hash);
144
145 if laucher_solution.amount == 1
146 && laucher_solution.singleton_puzzle_hash == predicted_puzzle_hash.into()
147 {
148 Ok(Some((candidate, hints.final_puzzle_hash_hints)))
149 } else {
150 Ok(None)
151 }
152 }
153
154 pub fn to_hints<H>(&self, final_puzzle_hash_hints: H) -> StateSchedulerLauncherHints<S, H> {
155 StateSchedulerLauncherHints {
156 my_launcher_id: self.launcher_id,
157 receiver_singleton_launcher_id: self.receiver_singleton_launcher_id,
158 final_puzzle_hash: self.final_puzzle_hash,
159 state_schedule: self.state_schedule.clone(),
160 final_puzzle_hash_hints,
161 }
162 }
163}
164
165#[derive(ToClvm, FromClvm, Debug, Clone, PartialEq, Eq)]
167#[clvm(curry)]
168pub struct StateSchedulerLauncherHints<S, H> {
169 pub my_launcher_id: Bytes32,
170 pub receiver_singleton_launcher_id: Bytes32,
171 pub final_puzzle_hash: Bytes32,
172 pub state_schedule: Vec<(u64, S)>,
173 #[clvm(rest)]
174 pub final_puzzle_hash_hints: H,
175}
176
177fn validate_state_schedule<S>(state_schedule: &[(u64, S)]) -> Result<(), DriverError> {
178 if state_schedule.is_empty() {
179 return Err(DriverError::InvalidStateSchedule);
180 }
181
182 for window in state_schedule.windows(2) {
183 if window[1].0 <= window[0].0 {
184 return Err(DriverError::InvalidStateSchedule);
185 }
186 }
187
188 Ok(())
189}
190
191#[cfg(test)]
192mod tests {
193 use chia_protocol::Bytes32;
194 use chia_puzzle_types::Memos;
195 use chia_sdk_types::Condition;
196 use clvm_traits::{FromClvm, ToClvm, clvm_quote};
197 use clvm_utils::ToTreeHash;
198 use clvmr::Allocator;
199
200 use crate::{CatalogRegistryState, DriverError};
201
202 use super::*;
203
204 fn mock_state(generator: u8) -> CatalogRegistryState {
205 CatalogRegistryState {
206 cat_maker_puzzle_hash: Bytes32::new([generator; 32]),
207 registration_price: u64::from(generator) * 1000,
208 }
209 }
210
211 #[test]
212 fn test_rejects_empty_schedule() {
213 let err = StateSchedulerInfo::new(
214 Bytes32::default(),
215 Bytes32::default(),
216 Vec::<(u64, CatalogRegistryState)>::new(),
217 0,
218 Bytes32::default(),
219 )
220 .unwrap_err();
221 assert!(matches!(err, DriverError::InvalidStateSchedule));
222 }
223
224 #[test]
225 fn test_rejects_duplicate_timestamps() {
226 let err = StateSchedulerInfo::new(
227 Bytes32::default(),
228 Bytes32::default(),
229 vec![(100, mock_state(0)), (100, mock_state(1))],
230 0,
231 Bytes32::default(),
232 )
233 .unwrap_err();
234 assert!(matches!(err, DriverError::InvalidStateSchedule));
235 }
236
237 #[test]
238 fn test_rejects_non_increasing_timestamps() {
239 let err = StateSchedulerInfo::new(
240 Bytes32::default(),
241 Bytes32::default(),
242 vec![(200, mock_state(0)), (150, mock_state(1))],
243 0,
244 Bytes32::default(),
245 )
246 .unwrap_err();
247 assert!(matches!(err, DriverError::InvalidStateSchedule));
248 }
249
250 #[test]
251 fn test_accepts_strictly_increasing_timestamps() -> anyhow::Result<()> {
252 let info = StateSchedulerInfo::new(
253 Bytes32::new([1; 32]),
254 Bytes32::new([2; 32]),
255 vec![(100, mock_state(0)), (200, mock_state(1))],
256 0,
257 Bytes32::new([3; 32]),
258 )?;
259 assert_eq!(info.state_schedule.len(), 2);
260 Ok(())
261 }
262
263 #[test]
264 fn test_inner_puzzle_hash_uses_assert_seconds_absolute() -> anyhow::Result<()> {
265 let info = StateSchedulerInfo::new(
266 Bytes32::new([1; 32]),
267 Bytes32::new([2; 32]),
268 vec![(1_700_000_000, mock_state(0))],
269 0,
270 Bytes32::new([3; 32]),
271 )?;
272
273 let seconds_hash = info.inner_puzzle_hash();
274
275 let prefix_and_message: Bytes = XchandlesRegistryReceivedMessagePrefix::update_state(
276 info.state_schedule[0].1.tree_hash(),
277 )
278 .into();
279 let height_hash = StateSchedulerLayerArgs::<TreeHash, _>::curry_tree_hash(
280 SingletonStruct::new(info.receiver_singleton_launcher_id)
281 .tree_hash()
282 .into(),
283 prefix_and_message.tree_hash(),
284 &clvm_quote!(vec![
285 Condition::<()>::create_coin(info.final_puzzle_hash, 1, Memos::None),
286 Condition::assert_height_absolute(1_700_000_000),
287 ]),
288 );
289
290 assert_ne!(hex::encode(seconds_hash), hex::encode(height_hash));
291 assert_eq!(
292 hex::encode(seconds_hash),
293 hex::encode(info.inner_puzzle_hash_for(
294 info.final_puzzle_hash,
295 1_700_000_000,
296 prefix_and_message.tree_hash(),
297 ))
298 );
299
300 Ok(())
301 }
302
303 #[test]
304 fn test_launcher_hints_roundtrip() -> anyhow::Result<()> {
305 let schedule = vec![(100, mock_state(0)), (200, mock_state(1))];
306 let info = StateSchedulerInfo::new(
307 Bytes32::new([9; 32]),
308 Bytes32::new([8; 32]),
309 schedule,
310 0,
311 Bytes32::new([7; 32]),
312 )?;
313 let hints = info.to_hints(NodePtr::NIL);
314
315 let mut allocator = Allocator::new();
316 let ptr = hints.to_clvm(&mut allocator)?;
317 let roundtrip = StateSchedulerLauncherHints::<CatalogRegistryState, NodePtr>::from_clvm(
318 &allocator, ptr,
319 )?;
320
321 assert_eq!(roundtrip.my_launcher_id, hints.my_launcher_id);
322 assert_eq!(
323 roundtrip.receiver_singleton_launcher_id,
324 hints.receiver_singleton_launcher_id
325 );
326 assert_eq!(roundtrip.final_puzzle_hash, hints.final_puzzle_hash);
327 assert_eq!(roundtrip.state_schedule, hints.state_schedule);
328 assert_eq!(roundtrip.final_puzzle_hash_hints, NodePtr::NIL);
329
330 Ok(())
331 }
332}