1use chia_protocol::Bytes32;
2use chia_puzzle_types::{cat::CatArgs, singleton::SingletonArgs};
3use chia_sdk_types::{
4 MerkleTree,
5 puzzles::{
6 ActionLayerArgs, DefaultReserveAmountFromStateProgramArgs, P2DelegatedBySingletonLayerArgs,
7 RESERVE_FINALIZER_DEFAULT_RESERVE_AMOUNT_FROM_STATE_PROGRAM_HASH,
8 ReserveFinalizer2ndCurryArgs,
9 },
10};
11use clvm_traits::{ClvmDecoder, ClvmEncoder, FromClvm, FromClvmError, Raw, ToClvm, ToClvmError};
12use clvm_utils::{ToTreeHash, TreeHash};
13use clvmr::{Allocator, NodePtr};
14
15use crate::{
16 ActionLayer, DriverError, Finalizer, Layer, Puzzle, RewardDistributorAddEntryAction,
17 RewardDistributorAddIncentivesAction, RewardDistributorCommitIncentivesAction,
18 RewardDistributorInitiatePayoutAction, RewardDistributorNewEpochAction,
19 RewardDistributorRefreshAction, RewardDistributorRemoveEntryAction,
20 RewardDistributorStakeAction, RewardDistributorSyncAction, RewardDistributorUnstakeAction,
21 RewardDistributorWithdrawIncentivesAction, SingletonAction, SingletonLayer, SpendContext,
22};
23
24use super::Reserveful;
25
26pub type RewardDistributorLayers = SingletonLayer<ActionLayer<RewardDistributorState, NodePtr>>;
27
28#[derive(
29 Debug, Clone, PartialEq, Eq, ToClvm, FromClvm, Copy, serde::Serialize, serde::Deserialize,
30)]
31#[clvm(list)]
32pub struct RoundRewardInfo {
33 pub cumulative_payout: u128,
34 #[clvm(rest)]
35 pub remaining_rewards: u128,
36}
37
38#[derive(
39 Debug, Clone, PartialEq, Eq, ToClvm, FromClvm, Copy, serde::Serialize, serde::Deserialize,
40)]
41#[clvm(list)]
42pub struct RoundTimeInfo {
43 pub last_update: u64,
44 #[clvm(rest)]
45 pub epoch_end: u64,
46}
47
48#[must_use]
49#[derive(
50 Debug, Clone, PartialEq, Eq, ToClvm, FromClvm, Copy, serde::Serialize, serde::Deserialize,
51)]
52#[clvm(list)]
53pub struct RewardDistributorState {
54 pub total_reserves: u64,
55 pub active_shares: u64,
56 pub round_reward_info: RoundRewardInfo,
57 #[clvm(rest)]
58 pub round_time_info: RoundTimeInfo,
59}
60
61impl RewardDistributorState {
62 pub fn initial(first_epoch_start: u64) -> Self {
63 Self {
64 total_reserves: 0,
65 active_shares: 0,
66 round_reward_info: RoundRewardInfo {
67 cumulative_payout: 0,
68 remaining_rewards: 0,
69 },
70 round_time_info: RoundTimeInfo {
71 last_update: first_epoch_start,
72 epoch_end: first_epoch_start,
73 },
74 }
75 }
76}
77
78impl Reserveful for RewardDistributorState {
79 fn reserve_amount(&self, index: u64) -> u64 {
80 if index == 0 { self.total_reserves } else { 0 }
81 }
82}
83
84#[derive(Debug, Clone, Copy, PartialEq, Eq)]
85pub enum RewardDistributorType {
86 Managed {
87 manager_singleton_launcher_id: Bytes32,
88 },
89 NftCollection {
90 collection_did_launcher_id: Bytes32,
91 },
92 CuratedNft {
93 store_launcher_id: Bytes32,
94 refreshable: bool,
95 },
96 Cat {
97 asset_id: Bytes32,
98 hidden_puzzle_hash: Option<Bytes32>,
99 },
100}
101
102impl<N, D: ClvmDecoder<Node = N>> FromClvm<D> for RewardDistributorType {
103 fn from_clvm(decoder: &D, node: N) -> Result<Self, FromClvmError> {
104 let type_pair: (u8, Raw<N>) = FromClvm::from_clvm(decoder, node)?;
105
106 match type_pair.0 {
107 1 => Ok(RewardDistributorType::Managed {
108 manager_singleton_launcher_id: FromClvm::from_clvm(decoder, type_pair.1.0)?,
109 }),
110 2 => Ok(RewardDistributorType::NftCollection {
111 collection_did_launcher_id: FromClvm::from_clvm(decoder, type_pair.1.0)?,
112 }),
113 3 => {
114 let (store_launcher_id, refreshable): (Bytes32, bool) =
115 FromClvm::from_clvm(decoder, type_pair.1.0)?;
116 Ok(RewardDistributorType::CuratedNft {
117 store_launcher_id,
118 refreshable,
119 })
120 }
121 4 => {
122 let (asset_id, hidden_puzzle_hash): (Bytes32, Option<Bytes32>) =
123 FromClvm::from_clvm(decoder, type_pair.1.0)?;
124 Ok(RewardDistributorType::Cat {
125 asset_id,
126 hidden_puzzle_hash,
127 })
128 }
129 _ => Err(FromClvmError::Custom(format!(
130 "Invalid RewardDistributorType: {}",
131 type_pair.0
132 ))),
133 }
134 }
135}
136
137impl<N, E: ClvmEncoder<Node = N>> ToClvm<E> for RewardDistributorType {
138 fn to_clvm(&self, encoder: &mut E) -> Result<N, ToClvmError> {
139 match self {
140 RewardDistributorType::Managed {
141 manager_singleton_launcher_id,
142 } => (1, manager_singleton_launcher_id).to_clvm(encoder),
143 RewardDistributorType::NftCollection {
144 collection_did_launcher_id,
145 } => (2, collection_did_launcher_id).to_clvm(encoder),
146 RewardDistributorType::CuratedNft {
147 store_launcher_id,
148 refreshable,
149 } => (3, (store_launcher_id, refreshable)).to_clvm(encoder),
150 RewardDistributorType::Cat {
151 asset_id,
152 hidden_puzzle_hash,
153 } => (4, (asset_id, hidden_puzzle_hash)).to_clvm(encoder),
154 }
155 }
156}
157
158#[must_use]
159#[derive(Debug, Clone, PartialEq, Eq, Copy, ToClvm, FromClvm)]
160#[clvm(list)]
161pub struct RewardDistributorConstants {
162 pub launcher_id: Bytes32,
163 pub reward_distributor_type: RewardDistributorType,
164 pub fee_payout_puzzle_hash: Bytes32,
165 pub epoch_seconds: u64,
166 pub precision: u64,
167 pub max_seconds_offset: u64,
168 pub payout_threshold: u64,
169 pub require_payout_approval: bool,
170 pub fee_bps: u64,
171 pub withdrawal_share_bps: u64,
172 pub reserve_asset_id: Bytes32,
173 pub reserve_inner_puzzle_hash: Bytes32,
174 pub reserve_full_puzzle_hash: Bytes32,
175}
176
177impl RewardDistributorConstants {
178 #[allow(clippy::too_many_arguments)]
179 pub fn without_launcher_id(
180 reward_distributor_type: RewardDistributorType,
181 fee_payout_puzzle_hash: Bytes32,
182 epoch_seconds: u64,
183 precision: u64,
184 max_seconds_offset: u64,
185 payout_threshold: u64,
186 require_payout_approval: bool,
187 fee_bps: u64,
188 withdrawal_share_bps: u64,
189 reserve_asset_id: Bytes32,
190 ) -> Self {
191 Self {
192 launcher_id: Bytes32::default(),
193 reward_distributor_type,
194 fee_payout_puzzle_hash,
195 epoch_seconds,
196 precision,
197 max_seconds_offset,
198 payout_threshold,
199 require_payout_approval,
200 fee_bps,
201 withdrawal_share_bps,
202 reserve_asset_id,
203 reserve_inner_puzzle_hash: Bytes32::default(),
204 reserve_full_puzzle_hash: Bytes32::default(),
205 }
206 }
207
208 pub fn with_launcher_id(mut self, launcher_id: Bytes32) -> Self {
209 self.launcher_id = launcher_id;
210 self.reserve_inner_puzzle_hash =
211 P2DelegatedBySingletonLayerArgs::curry_tree_hash_with_launcher_id(launcher_id, 0)
212 .into();
213 self.reserve_full_puzzle_hash =
214 CatArgs::curry_tree_hash(self.reserve_asset_id, self.reserve_inner_puzzle_hash.into())
215 .into();
216 self
217 }
218}
219
220#[must_use]
221#[derive(Debug, Clone, PartialEq, Eq, Copy)]
222pub struct RewardDistributorInfo {
223 pub state: RewardDistributorState,
224
225 pub constants: RewardDistributorConstants,
226}
227
228impl RewardDistributorInfo {
229 pub fn new(state: RewardDistributorState, constants: RewardDistributorConstants) -> Self {
230 Self { state, constants }
231 }
232
233 pub fn with_state(mut self, state: RewardDistributorState) -> Self {
234 self.state = state;
235 self
236 }
237
238 pub fn action_puzzle_hashes(constants: &RewardDistributorConstants) -> Vec<Bytes32> {
239 let mut action_puzzle_hashes = vec![
240 RewardDistributorAddIncentivesAction::from_constants(constants)
241 .tree_hash()
242 .into(),
243 RewardDistributorCommitIncentivesAction::from_constants(constants)
244 .tree_hash()
245 .into(),
246 RewardDistributorInitiatePayoutAction::from_constants(constants)
247 .tree_hash()
248 .into(),
249 RewardDistributorNewEpochAction::from_constants(constants)
250 .tree_hash()
251 .into(),
252 RewardDistributorSyncAction::from_constants(constants)
253 .tree_hash()
254 .into(),
255 RewardDistributorWithdrawIncentivesAction::from_constants(constants)
256 .tree_hash()
257 .into(),
258 match constants.reward_distributor_type {
259 RewardDistributorType::Managed {
260 manager_singleton_launcher_id: _,
261 } => RewardDistributorAddEntryAction::from_constants(constants)
262 .tree_hash()
263 .into(),
264 RewardDistributorType::NftCollection {
265 collection_did_launcher_id: _,
266 }
267 | RewardDistributorType::CuratedNft {
268 store_launcher_id: _,
269 refreshable: _,
270 }
271 | RewardDistributorType::Cat {
272 asset_id: _,
273 hidden_puzzle_hash: _,
274 } => RewardDistributorStakeAction::from_constants(constants)
275 .tree_hash()
276 .into(),
277 },
278 match constants.reward_distributor_type {
279 RewardDistributorType::Managed {
280 manager_singleton_launcher_id: _,
281 } => RewardDistributorRemoveEntryAction::from_constants(constants)
282 .tree_hash()
283 .into(),
284 RewardDistributorType::NftCollection {
285 collection_did_launcher_id: _,
286 }
287 | RewardDistributorType::CuratedNft {
288 store_launcher_id: _,
289 refreshable: _,
290 }
291 | RewardDistributorType::Cat {
292 asset_id: _,
293 hidden_puzzle_hash: _,
294 } => RewardDistributorUnstakeAction::from_constants(constants)
295 .tree_hash()
296 .into(),
297 },
298 ];
299
300 if let RewardDistributorType::CuratedNft { refreshable, .. } =
301 constants.reward_distributor_type
302 && refreshable
303 {
304 action_puzzle_hashes.push(
305 RewardDistributorRefreshAction::from_constants(constants)
306 .tree_hash()
307 .into(),
308 );
309 }
310
311 action_puzzle_hashes
312 }
313
314 pub fn into_layers(
315 self,
316 ctx: &mut SpendContext,
317 ) -> Result<RewardDistributorLayers, DriverError> {
318 Ok(SingletonLayer::new(
319 self.constants.launcher_id,
320 ActionLayer::<RewardDistributorState, NodePtr>::from_action_puzzle_hashes(
321 &Self::action_puzzle_hashes(&self.constants),
322 self.state,
323 Finalizer::Reserve {
324 reserve_full_puzzle_hash: self.constants.reserve_full_puzzle_hash,
325 reserve_inner_puzzle_hash: self.constants.reserve_inner_puzzle_hash,
326 reserve_amount_from_state_program: ctx
327 .alloc_mod::<DefaultReserveAmountFromStateProgramArgs>(
328 )?,
329 hint: self.constants.launcher_id,
330 },
331 ),
332 ))
333 }
334
335 pub fn parse(
336 allocator: &mut Allocator,
337 puzzle: Puzzle,
338 constants: RewardDistributorConstants,
339 ) -> Result<Option<Self>, DriverError> {
340 let Some(layers) = RewardDistributorLayers::parse_puzzle(allocator, puzzle)? else {
341 return Ok(None);
342 };
343
344 let action_puzzle_hashes = Self::action_puzzle_hashes(&constants);
345 let merkle_root = MerkleTree::new(&action_puzzle_hashes).root();
346 if layers.inner_puzzle.merkle_root != merkle_root {
347 return Ok(None);
348 }
349
350 Ok(Some(Self::from_layers(&layers, constants)))
351 }
352
353 pub fn from_layers(
354 layers: &RewardDistributorLayers,
355 constants: RewardDistributorConstants,
356 ) -> Self {
357 Self {
358 state: layers.inner_puzzle.state,
359 constants,
360 }
361 }
362
363 pub fn puzzle_hash(&self) -> TreeHash {
364 SingletonArgs::curry_tree_hash(self.constants.launcher_id, self.inner_puzzle_hash())
365 }
366
367 pub fn inner_puzzle_hash(&self) -> TreeHash {
368 ActionLayerArgs::curry_tree_hash(
369 ReserveFinalizer2ndCurryArgs::curry_tree_hash(
370 self.constants.reserve_full_puzzle_hash,
371 self.constants.reserve_inner_puzzle_hash,
372 RESERVE_FINALIZER_DEFAULT_RESERVE_AMOUNT_FROM_STATE_PROGRAM_HASH,
373 self.constants.launcher_id,
374 ),
375 MerkleTree::new(&Self::action_puzzle_hashes(&self.constants)).root(),
376 self.state.tree_hash(),
377 )
378 }
379}