1use std::sync::{Arc, Mutex};
2
3use bindy::{Error, Result};
4use chia_bls::{SecretKey, Signature};
5use chia_protocol::{Bytes32, Coin};
6use chia_puzzle_types::{LineageProof, singleton::SingletonStruct};
7use chia_sdk_driver::{
8 Cat, Reserve, RewardDistributor as SdkRewardDistributor, RewardDistributorAddEntryAction,
9 RewardDistributorAddIncentivesAction, RewardDistributorCommitIncentivesAction,
10 RewardDistributorConstants, RewardDistributorInitiatePayoutAction,
11 RewardDistributorNewEpochAction, RewardDistributorRemoveEntryAction,
12 RewardDistributorStakeAction, RewardDistributorState, RewardDistributorSyncAction,
13 RewardDistributorType, RewardDistributorUnstakeAction,
14 RewardDistributorWithdrawIncentivesAction, RoundRewardInfo, RoundTimeInfo, SpendContext,
15};
16use chia_sdk_types::{
17 Conditions, Mod,
18 puzzles::{
19 IntermediaryCoinProof, NftLauncherProof, NonceWrapperArgs, RewardDistributorSlotNonce,
20 },
21};
22use clvm_utils::{ToTreeHash, TreeHash};
23
24use crate::{
25 AsProgram, AsPtr, CatSpend, CommitmentSlot, EntrySlot, Nft, NotarizedPayment, Program, Proof,
26 RewardSlot,
27};
28
29pub trait RewardDistributorTypeExt {}
30
31impl RewardDistributorTypeExt for RewardDistributorType {}
32
33pub trait RewardDistributorConstantsExt
34where
35 Self: Sized,
36{
37 #[allow(clippy::too_many_arguments)]
38 fn without_launcher_id(
39 reward_distributor_type: RewardDistributorType,
40 manager_or_collection_did_launcher_id: Bytes32,
41 fee_payout_puzzle_hash: Bytes32,
42 epoch_seconds: u64,
43 max_seconds_offset: u64,
44 payout_threshold: u64,
45 fee_bps: u64,
46 withdrawal_share_bps: u64,
47 reserve_asset_id: Bytes32,
48 ) -> Result<Self>;
49
50 fn with_launcher_id(&self, launcher_id: Bytes32) -> Result<Self>;
51}
52
53impl RewardDistributorConstantsExt for RewardDistributorConstants {
54 #[allow(clippy::too_many_arguments)]
55 fn without_launcher_id(
56 reward_distributor_type: RewardDistributorType,
57 manager_or_collection_did_launcher_id: Bytes32,
58 fee_payout_puzzle_hash: Bytes32,
59 epoch_seconds: u64,
60 max_seconds_offset: u64,
61 payout_threshold: u64,
62 fee_bps: u64,
63 withdrawal_share_bps: u64,
64 reserve_asset_id: Bytes32,
65 ) -> Result<Self> {
66 Ok(RewardDistributorConstants::without_launcher_id(
67 reward_distributor_type,
68 manager_or_collection_did_launcher_id,
69 fee_payout_puzzle_hash,
70 epoch_seconds,
71 max_seconds_offset,
72 payout_threshold,
73 fee_bps,
74 withdrawal_share_bps,
75 reserve_asset_id,
76 ))
77 }
78
79 fn with_launcher_id(&self, launcher_id: Bytes32) -> Result<Self> {
80 Ok(RewardDistributorConstants::with_launcher_id(
81 *self,
82 launcher_id,
83 ))
84 }
85}
86
87pub trait RoundRewardInfoExt {}
88
89impl RoundRewardInfoExt for RoundRewardInfo {}
90
91pub trait RoundTimeInfoExt {}
92
93impl RoundTimeInfoExt for RoundTimeInfo {}
94
95pub trait RewardDistributorStateExt
96where
97 Self: Sized,
98{
99 fn initial(first_epoch_start: u64) -> Result<Self>;
100}
101
102impl RewardDistributorStateExt for RewardDistributorState {
103 fn initial(first_epoch_start: u64) -> Result<Self> {
104 Ok(RewardDistributorState::initial(first_epoch_start))
105 }
106}
107
108pub trait RewardDistributorLauncherSolutionInfoExt {}
109
110impl RewardDistributorLauncherSolutionInfoExt for RewardDistributorLauncherSolutionInfo {}
111
112#[derive(Clone, Copy)]
113pub struct RewardDistributorLauncherSolutionInfo {
114 pub constants: RewardDistributorConstants,
115 pub initial_state: RewardDistributorState,
116 pub coin: Coin,
117}
118
119#[derive(Clone)]
120pub struct RewardDistributorFinishedSpendResult {
121 pub new_distributor: RewardDistributor,
122 pub signature: Signature,
123}
124
125#[derive(Clone)]
126pub struct RewardDistributorInitiatePayoutResult {
127 pub conditions: Vec<Program>,
128 pub payout_amount: u64,
129}
130
131#[derive(Clone)]
132pub struct RewardDistributorNewEpochResult {
133 pub conditions: Vec<Program>,
134 pub epoch_fee: u64,
135}
136
137#[derive(Clone)]
138pub struct RewardDistributorWithdrawIncentivesResult {
139 pub conditions: Vec<Program>,
140 pub withdrawn_amount: u64,
141}
142
143#[derive(Clone)]
144pub struct RewardDistributorRemoveEntryResult {
145 pub conditions: Vec<Program>,
146 pub last_payment_amount: u64,
147}
148
149pub trait IntermediaryCoinProofExt {}
150
151impl IntermediaryCoinProofExt for IntermediaryCoinProof {}
152
153pub trait NftLauncherProofExt {}
154
155impl NftLauncherProofExt for NftLauncherProof {}
156
157#[derive(Clone)]
158pub struct RewardDistributorStakeResult {
159 pub conditions: Vec<Program>,
160 pub notarized_payment: NotarizedPayment,
161 pub new_nft: Nft,
162}
163
164#[derive(Clone)]
165pub struct RewardDistributorUnstakeResult {
166 pub conditions: Vec<Program>,
167 pub payment_amount: u64,
168}
169
170#[derive(Clone)]
171pub struct RewardDistributorLaunchResult {
172 pub security_signature: Signature,
173 pub security_secret_key: SecretKey,
174 pub reward_distributor: RewardDistributor,
175 pub first_epoch_slot: RewardSlot,
176 pub refunded_cat: Cat,
177}
178
179#[derive(Clone)]
180pub struct RewardDistributorInfoFromLauncher {
181 pub constants: RewardDistributorConstants,
182 pub initial_state: RewardDistributorState,
183 pub eve_singleton: Coin,
184}
185
186#[derive(Clone)]
187pub struct RewardDistributorInfoFromEveCoin {
188 pub distributor: RewardDistributor,
189 pub first_reward_slot: RewardSlot,
190}
191
192#[derive(Clone)]
193pub struct RewardDistributor {
194 pub(crate) clvm: Arc<Mutex<SpendContext>>,
195 pub(crate) distributor: Arc<Mutex<SdkRewardDistributor>>,
196}
197
198impl RewardDistributor {
199 pub fn coin(&self) -> Result<Coin> {
200 Ok(self.distributor.lock().unwrap().coin)
201 }
202
203 pub fn proof(&self) -> Result<Proof> {
204 Ok(self.distributor.lock().unwrap().proof.into())
205 }
206
207 pub fn state(&self) -> Result<RewardDistributorState> {
208 Ok(self.distributor.lock().unwrap().info.state)
209 }
210
211 pub fn constants(&self) -> Result<RewardDistributorConstants> {
212 Ok(self.distributor.lock().unwrap().info.constants)
213 }
214
215 pub fn inner_puzzle_hash(&self) -> Result<TreeHash> {
216 Ok(self.distributor.lock().unwrap().info.inner_puzzle_hash())
217 }
218
219 pub fn puzzle_hash(&self) -> Result<TreeHash> {
220 Ok(self.distributor.lock().unwrap().info.puzzle_hash())
221 }
222
223 pub fn reserve_coin(&self) -> Result<Coin> {
224 Ok(self.distributor.lock().unwrap().reserve.coin)
225 }
226
227 pub fn reserve_asset_id(&self) -> Result<Bytes32> {
228 Ok(self.distributor.lock().unwrap().reserve.asset_id)
229 }
230
231 pub fn reserve_proof(&self) -> Result<LineageProof> {
232 Ok(self.distributor.lock().unwrap().reserve.proof)
233 }
234
235 pub fn pending_created_reward_slots(&self) -> Result<Vec<RewardSlot>> {
236 let distributor = self.distributor.lock().unwrap();
237
238 Ok(distributor
239 .pending_spend
240 .created_reward_slots
241 .clone()
242 .into_iter()
243 .map(|slot_value| {
244 RewardSlot::from_slot(
245 distributor
246 .created_slot_value_to_slot(slot_value, RewardDistributorSlotNonce::REWARD),
247 )
248 })
249 .collect())
250 }
251
252 pub fn pending_created_commitment_slots(&self) -> Result<Vec<CommitmentSlot>> {
253 let distributor = self.distributor.lock().unwrap();
254
255 Ok(distributor
256 .pending_spend
257 .created_commitment_slots
258 .clone()
259 .into_iter()
260 .map(|slot_value| {
261 CommitmentSlot::from_slot(
262 distributor.created_slot_value_to_slot(
263 slot_value,
264 RewardDistributorSlotNonce::COMMITMENT,
265 ),
266 )
267 })
268 .collect())
269 }
270
271 pub fn pending_created_entry_slots(&self) -> Result<Vec<EntrySlot>> {
272 let distributor = self.distributor.lock().unwrap();
273
274 Ok(distributor
275 .pending_spend
276 .created_entry_slots
277 .clone()
278 .into_iter()
279 .map(|slot_value| {
280 EntrySlot::from_slot(
281 distributor
282 .created_slot_value_to_slot(slot_value, RewardDistributorSlotNonce::ENTRY),
283 )
284 })
285 .collect())
286 }
287
288 pub fn pending_signature(&self) -> Result<Signature> {
289 Ok(self
290 .distributor
291 .lock()
292 .unwrap()
293 .pending_spend
294 .signature
295 .clone())
296 }
297
298 pub fn reserve_full_puzzle_hash(
299 asset_id: Bytes32,
300 distributor_launcher_id: Bytes32,
301 nonce: u64,
302 ) -> Result<TreeHash> {
303 Ok(Reserve::puzzle_hash(
304 asset_id,
305 SingletonStruct::new(distributor_launcher_id)
306 .tree_hash()
307 .into(),
308 nonce,
309 ))
310 }
311
312 pub fn parse_launcher_solution(
313 launcher_coin: Coin,
314 launcher_solution: Program,
315 ) -> Result<Option<RewardDistributorInfoFromLauncher>> {
316 let mut ctx = launcher_solution.0.lock().unwrap();
317
318 Ok(SdkRewardDistributor::from_launcher_solution(
319 &mut ctx,
320 launcher_coin,
321 launcher_solution.1,
322 )?
323 .map(|(constants, initial_state, eve_singleton)| {
324 RewardDistributorInfoFromLauncher {
325 constants,
326 initial_state,
327 eve_singleton,
328 }
329 }))
330 }
331
332 pub fn finish_spend(
333 &self,
334 other_cat_spends: Vec<CatSpend>,
335 ) -> Result<RewardDistributorFinishedSpendResult> {
336 let mut ctx = self.clvm.lock().unwrap();
337
338 let (distributor, signature) = self.distributor.lock().unwrap().clone().finish_spend(
339 &mut ctx,
340 other_cat_spends.into_iter().map(Into::into).collect(),
341 )?;
342
343 Ok(RewardDistributorFinishedSpendResult {
344 new_distributor: RewardDistributor {
345 clvm: self.clvm.clone(),
346 distributor: Arc::new(Mutex::new(distributor)),
347 },
348 signature,
349 })
350 }
351
352 fn sdk_conditions_to_program_list(
353 &self,
354 ctx: &mut SpendContext,
355 conditions: Conditions,
356 ) -> Result<Vec<Program>> {
357 let mut result = Vec::with_capacity(conditions.len());
358
359 for condition in conditions {
360 result.push(Program(self.clvm.clone(), ctx.alloc(&condition)?));
361 }
362
363 Ok(result)
364 }
365
366 pub fn add_incentives(&self, amount: u64) -> Result<Vec<Program>> {
367 let mut ctx = self.clvm.lock().unwrap();
368 let mut distributor = self.distributor.lock().unwrap();
369
370 let conditions = distributor
371 .new_action::<RewardDistributorAddIncentivesAction>()
372 .spend(&mut ctx, &mut distributor, amount)?;
373
374 self.sdk_conditions_to_program_list(&mut ctx, conditions)
375 }
376
377 pub fn commit_incentives(
378 &self,
379 reward_slot: RewardSlot,
380 epoch_start: u64,
381 clawback_ph: Bytes32,
382 rewards_to_add: u64,
383 ) -> Result<Vec<Program>> {
384 let mut ctx = self.clvm.lock().unwrap();
385 let mut distributor = self.distributor.lock().unwrap();
386
387 let conditions = distributor
388 .new_action::<RewardDistributorCommitIncentivesAction>()
389 .spend(
390 &mut ctx,
391 &mut distributor,
392 reward_slot.to_slot(),
393 epoch_start,
394 clawback_ph,
395 rewards_to_add,
396 )?;
397
398 self.sdk_conditions_to_program_list(&mut ctx, conditions)
399 }
400
401 pub fn initiate_payout(
402 &self,
403 entry_slot: EntrySlot,
404 ) -> Result<RewardDistributorInitiatePayoutResult> {
405 let mut ctx = self.clvm.lock().unwrap();
406 let mut distributor = self.distributor.lock().unwrap();
407
408 let (conditions, payout_amount) = distributor
409 .new_action::<RewardDistributorInitiatePayoutAction>()
410 .spend(&mut ctx, &mut distributor, entry_slot.to_slot())?;
411
412 Ok(RewardDistributorInitiatePayoutResult {
413 conditions: self.sdk_conditions_to_program_list(&mut ctx, conditions)?,
414 payout_amount,
415 })
416 }
417
418 pub fn new_epoch(&self, reward_slot: RewardSlot) -> Result<RewardDistributorNewEpochResult> {
419 let mut ctx = self.clvm.lock().unwrap();
420 let mut distributor = self.distributor.lock().unwrap();
421
422 let (conditions, epoch_fee) = distributor
423 .new_action::<RewardDistributorNewEpochAction>()
424 .spend(&mut ctx, &mut distributor, reward_slot.to_slot())?;
425
426 Ok(RewardDistributorNewEpochResult {
427 conditions: self.sdk_conditions_to_program_list(&mut ctx, conditions)?,
428 epoch_fee,
429 })
430 }
431
432 pub fn sync(&self, update_time: u64) -> Result<Vec<Program>> {
433 let mut ctx = self.clvm.lock().unwrap();
434 let mut distributor = self.distributor.lock().unwrap();
435
436 let conditions = distributor
437 .new_action::<RewardDistributorSyncAction>()
438 .spend(&mut ctx, &mut distributor, update_time)?;
439
440 self.sdk_conditions_to_program_list(&mut ctx, conditions)
441 }
442
443 pub fn withdraw_incentives(
444 &self,
445 commitment_slot: CommitmentSlot,
446 reward_slot: RewardSlot,
447 ) -> Result<RewardDistributorWithdrawIncentivesResult> {
448 let mut ctx = self.clvm.lock().unwrap();
449 let mut distributor = self.distributor.lock().unwrap();
450
451 let (conditions, withdrawn_amount) = distributor
452 .new_action::<RewardDistributorWithdrawIncentivesAction>()
453 .spend(
454 &mut ctx,
455 &mut distributor,
456 commitment_slot.to_slot(),
457 reward_slot.to_slot(),
458 )?;
459
460 Ok(RewardDistributorWithdrawIncentivesResult {
461 conditions: self.sdk_conditions_to_program_list(&mut ctx, conditions)?,
462 withdrawn_amount,
463 })
464 }
465
466 pub fn add_entry(
467 &self,
468 payout_puzzle_hash: Bytes32,
469 shares: u64,
470 manager_singleton_inner_puzzle_hash: Bytes32,
471 ) -> Result<Vec<Program>> {
472 let mut ctx = self.clvm.lock().unwrap();
473 let mut distributor = self.distributor.lock().unwrap();
474
475 if distributor.info.constants.reward_distributor_type != RewardDistributorType::Manager {
476 return Err(Error::Custom(
477 "Reward distributor is not managed".to_string(),
478 ));
479 }
480
481 let conditions = distributor
482 .new_action::<RewardDistributorAddEntryAction>()
483 .spend(
484 &mut ctx,
485 &mut distributor,
486 payout_puzzle_hash,
487 shares,
488 manager_singleton_inner_puzzle_hash,
489 )?;
490
491 self.sdk_conditions_to_program_list(&mut ctx, conditions)
492 }
493
494 pub fn remove_entry(
495 &self,
496 entry_slot: EntrySlot,
497 manager_singleton_inner_puzzle_hash: Bytes32,
498 ) -> Result<RewardDistributorRemoveEntryResult> {
499 let mut ctx = self.clvm.lock().unwrap();
500 let mut distributor = self.distributor.lock().unwrap();
501
502 if distributor.info.constants.reward_distributor_type != RewardDistributorType::Manager {
503 return Err(Error::Custom(
504 "Reward distributor is not managed".to_string(),
505 ));
506 }
507
508 let (conditions, last_payment_amount) = distributor
509 .new_action::<RewardDistributorRemoveEntryAction>()
510 .spend(
511 &mut ctx,
512 &mut distributor,
513 entry_slot.to_slot(),
514 manager_singleton_inner_puzzle_hash,
515 )?;
516
517 Ok(RewardDistributorRemoveEntryResult {
518 conditions: self.sdk_conditions_to_program_list(&mut ctx, conditions)?,
519 last_payment_amount,
520 })
521 }
522
523 pub fn stake(
524 &self,
525 current_nft: Nft,
526 nft_launcher_proof: NftLauncherProof,
527 entry_custody_puzzle_hash: Bytes32,
528 ) -> Result<RewardDistributorStakeResult> {
529 let mut ctx = self.clvm.lock().unwrap();
530 let mut distributor = self.distributor.lock().unwrap();
531
532 if distributor.info.constants.reward_distributor_type != RewardDistributorType::Nft {
533 return Err(Error::Custom(
534 "Reward distributor is not an NFT one".to_string(),
535 ));
536 }
537
538 let sdk_nft = current_nft.as_ptr(&ctx);
539 let (conditions, notarized_payment, new_nft) = distributor
540 .new_action::<RewardDistributorStakeAction>()
541 .spend(
542 &mut ctx,
543 &mut distributor,
544 sdk_nft,
545 nft_launcher_proof,
546 entry_custody_puzzle_hash,
547 )?;
548
549 Ok(RewardDistributorStakeResult {
550 conditions: self.sdk_conditions_to_program_list(&mut ctx, conditions)?,
551 notarized_payment: notarized_payment.as_program(&self.clvm),
552 new_nft: new_nft.as_program(&self.clvm),
553 })
554 }
555
556 pub fn unstake(
557 &self,
558 entry_slot: EntrySlot,
559 locked_nft: Nft,
560 ) -> Result<RewardDistributorUnstakeResult> {
561 let mut ctx = self.clvm.lock().unwrap();
562 let mut distributor = self.distributor.lock().unwrap();
563
564 let sdk_locked_nft = locked_nft.as_ptr(&ctx);
565 let (conditions, payment_amount) = distributor
566 .new_action::<RewardDistributorUnstakeAction>()
567 .spend(
568 &mut ctx,
569 &mut distributor,
570 entry_slot.to_slot(),
571 sdk_locked_nft,
572 )?;
573
574 Ok(RewardDistributorUnstakeResult {
575 conditions: self.sdk_conditions_to_program_list(&mut ctx, conditions)?,
576 payment_amount,
577 })
578 }
579
580 pub fn locked_nft_hint(
581 distributor_launcher_id: Bytes32,
582 custody_puzzle_hash: Bytes32,
583 ) -> Result<Bytes32> {
584 Ok(NonceWrapperArgs::<Bytes32, TreeHash> {
585 nonce: custody_puzzle_hash,
586 inner_puzzle: RewardDistributorStakeAction::my_p2_puzzle_hash(distributor_launcher_id)
587 .into(),
588 }
589 .curry_tree_hash()
590 .into())
591 }
592}