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::{FromClvm, ToClvm};
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 RewardDistributorRemoveEntryAction, RewardDistributorStakeAction, RewardDistributorSyncAction,
20 RewardDistributorUnstakeAction, RewardDistributorWithdrawIncentivesAction, SingletonAction,
21 SingletonLayer, SpendContext,
22};
23
24use super::Reserveful;
25
26pub type RewardDistributorLayers = SingletonLayer<ActionLayer<RewardDistributorState, NodePtr>>;
27
28#[derive(Debug, Clone, PartialEq, Eq, ToClvm, FromClvm, Copy)]
29#[clvm(list)]
30pub struct RoundRewardInfo {
31 pub cumulative_payout: u64,
32 #[clvm(rest)]
33 pub remaining_rewards: u64,
34}
35
36#[derive(Debug, Clone, PartialEq, Eq, ToClvm, FromClvm, Copy)]
37#[clvm(list)]
38pub struct RoundTimeInfo {
39 pub last_update: u64,
40 #[clvm(rest)]
41 pub epoch_end: u64,
42}
43
44#[must_use]
45#[derive(Debug, Clone, PartialEq, Eq, ToClvm, FromClvm, Copy)]
46#[clvm(list)]
47pub struct RewardDistributorState {
48 pub total_reserves: u64,
49 pub active_shares: u64,
50 pub round_reward_info: RoundRewardInfo,
51 pub round_time_info: RoundTimeInfo,
52}
53
54impl RewardDistributorState {
55 pub fn initial(first_epoch_start: u64) -> Self {
56 Self {
57 total_reserves: 0,
58 active_shares: 0,
59 round_reward_info: RoundRewardInfo {
60 cumulative_payout: 0,
61 remaining_rewards: 0,
62 },
63 round_time_info: RoundTimeInfo {
64 last_update: first_epoch_start,
65 epoch_end: first_epoch_start,
66 },
67 }
68 }
69}
70
71impl Reserveful for RewardDistributorState {
72 fn reserve_amount(&self, index: u64) -> u64 {
73 if index == 0 { self.total_reserves } else { 0 }
74 }
75}
76
77#[derive(Debug, Clone, Copy, PartialEq, Eq, ToClvm, FromClvm)]
78#[repr(u8)]
79#[clvm(atom)]
80pub enum RewardDistributorType {
81 Manager = 1,
82 Nft = 2,
83}
84
85#[must_use]
86#[derive(Debug, Clone, PartialEq, Eq, Copy, ToClvm, FromClvm)]
87#[clvm(list)]
88pub struct RewardDistributorConstants {
89 pub launcher_id: Bytes32,
90 pub reward_distributor_type: RewardDistributorType,
91 pub manager_or_collection_did_launcher_id: Bytes32,
92 pub fee_payout_puzzle_hash: Bytes32,
93 pub epoch_seconds: u64,
94 pub max_seconds_offset: u64,
95 pub payout_threshold: u64,
96 pub fee_bps: u64,
97 pub withdrawal_share_bps: u64,
98 pub reserve_asset_id: Bytes32,
99 pub reserve_inner_puzzle_hash: Bytes32,
100 pub reserve_full_puzzle_hash: Bytes32,
101}
102
103impl RewardDistributorConstants {
104 #[allow(clippy::too_many_arguments)]
105 pub fn without_launcher_id(
106 reward_distributor_type: RewardDistributorType,
107 manager_or_collection_did_launcher_id: Bytes32,
108 fee_payout_puzzle_hash: Bytes32,
109 epoch_seconds: u64,
110 max_seconds_offset: u64,
111 payout_threshold: u64,
112 fee_bps: u64,
113 withdrawal_share_bps: u64,
114 reserve_asset_id: Bytes32,
115 ) -> Self {
116 Self {
117 launcher_id: Bytes32::default(),
118 reward_distributor_type,
119 manager_or_collection_did_launcher_id,
120 fee_payout_puzzle_hash,
121 epoch_seconds,
122 max_seconds_offset,
123 payout_threshold,
124 fee_bps,
125 withdrawal_share_bps,
126 reserve_asset_id,
127 reserve_inner_puzzle_hash: Bytes32::default(),
128 reserve_full_puzzle_hash: Bytes32::default(),
129 }
130 }
131
132 pub fn with_launcher_id(mut self, launcher_id: Bytes32) -> Self {
133 self.launcher_id = launcher_id;
134 self.reserve_inner_puzzle_hash =
135 P2DelegatedBySingletonLayerArgs::curry_tree_hash_with_launcher_id(launcher_id, 0)
136 .into();
137 self.reserve_full_puzzle_hash =
138 CatArgs::curry_tree_hash(self.reserve_asset_id, self.reserve_inner_puzzle_hash.into())
139 .into();
140 self
141 }
142}
143
144#[must_use]
145#[derive(Debug, Clone, PartialEq, Eq, Copy)]
146pub struct RewardDistributorInfo {
147 pub state: RewardDistributorState,
148
149 pub constants: RewardDistributorConstants,
150}
151
152impl RewardDistributorInfo {
153 pub fn new(state: RewardDistributorState, constants: RewardDistributorConstants) -> Self {
154 Self { state, constants }
155 }
156
157 pub fn with_state(mut self, state: RewardDistributorState) -> Self {
158 self.state = state;
159 self
160 }
161
162 pub fn action_puzzle_hashes(constants: &RewardDistributorConstants) -> [Bytes32; 8] {
163 [
164 RewardDistributorAddIncentivesAction::from_constants(constants)
165 .tree_hash()
166 .into(),
167 RewardDistributorCommitIncentivesAction::from_constants(constants)
168 .tree_hash()
169 .into(),
170 RewardDistributorInitiatePayoutAction::from_constants(constants)
171 .tree_hash()
172 .into(),
173 RewardDistributorNewEpochAction::from_constants(constants)
174 .tree_hash()
175 .into(),
176 RewardDistributorSyncAction::from_constants(constants)
177 .tree_hash()
178 .into(),
179 RewardDistributorWithdrawIncentivesAction::from_constants(constants)
180 .tree_hash()
181 .into(),
182 match constants.reward_distributor_type {
183 RewardDistributorType::Manager => {
184 RewardDistributorAddEntryAction::from_constants(constants)
185 .tree_hash()
186 .into()
187 }
188 RewardDistributorType::Nft => {
189 RewardDistributorStakeAction::from_constants(constants)
190 .tree_hash()
191 .into()
192 }
193 },
194 match constants.reward_distributor_type {
195 RewardDistributorType::Manager => {
196 RewardDistributorRemoveEntryAction::from_constants(constants)
197 .tree_hash()
198 .into()
199 }
200 RewardDistributorType::Nft => {
201 RewardDistributorUnstakeAction::from_constants(constants)
202 .tree_hash()
203 .into()
204 }
205 },
206 ]
207 }
208
209 pub fn into_layers(
210 self,
211 ctx: &mut SpendContext,
212 ) -> Result<RewardDistributorLayers, DriverError> {
213 Ok(SingletonLayer::new(
214 self.constants.launcher_id,
215 ActionLayer::<RewardDistributorState, NodePtr>::from_action_puzzle_hashes(
216 &Self::action_puzzle_hashes(&self.constants),
217 self.state,
218 Finalizer::Reserve {
219 reserve_full_puzzle_hash: self.constants.reserve_full_puzzle_hash,
220 reserve_inner_puzzle_hash: self.constants.reserve_inner_puzzle_hash,
221 reserve_amount_from_state_program: ctx
222 .alloc_mod::<DefaultReserveAmountFromStateProgramArgs>(
223 )?,
224 hint: self.constants.launcher_id,
225 },
226 ),
227 ))
228 }
229
230 pub fn parse(
231 allocator: &mut Allocator,
232 puzzle: Puzzle,
233 constants: RewardDistributorConstants,
234 ) -> Result<Option<Self>, DriverError> {
235 let Some(layers) = RewardDistributorLayers::parse_puzzle(allocator, puzzle)? else {
236 return Ok(None);
237 };
238
239 let action_puzzle_hashes = Self::action_puzzle_hashes(&constants);
240 let merkle_root = MerkleTree::new(&action_puzzle_hashes).root();
241 if layers.inner_puzzle.merkle_root != merkle_root {
242 return Ok(None);
243 }
244
245 Ok(Some(Self::from_layers(&layers, constants)))
246 }
247
248 pub fn from_layers(
249 layers: &RewardDistributorLayers,
250 constants: RewardDistributorConstants,
251 ) -> Self {
252 Self {
253 state: layers.inner_puzzle.state,
254 constants,
255 }
256 }
257
258 pub fn puzzle_hash(&self) -> TreeHash {
259 SingletonArgs::curry_tree_hash(self.constants.launcher_id, self.inner_puzzle_hash())
260 }
261
262 pub fn inner_puzzle_hash(&self) -> TreeHash {
263 ActionLayerArgs::curry_tree_hash(
264 ReserveFinalizer2ndCurryArgs::curry_tree_hash(
265 self.constants.reserve_full_puzzle_hash,
266 self.constants.reserve_inner_puzzle_hash,
267 RESERVE_FINALIZER_DEFAULT_RESERVE_AMOUNT_FROM_STATE_PROGRAM_HASH,
268 self.constants.launcher_id,
269 ),
270 MerkleTree::new(&Self::action_puzzle_hashes(&self.constants)).root(),
271 self.state.tree_hash(),
272 )
273 }
274}