1use crate::{
10 error::{Result, TempoPrecompileError},
11 storage::Handler,
12 tip20::{Recipient, TIP20Token},
13};
14use alloy::primitives::{Address, U256, uint};
15use tempo_contracts::precompiles::{ITIP20, TIP20Error, TIP20Event};
16use tempo_precompiles_macros::Storable;
17use tempo_primitives::TempoAddressExt;
18
19pub const ACC_PRECISION: U256 = uint!(1000000000000000000_U256);
21
22impl TIP20Token {
23 pub fn distribute_reward(
35 &mut self,
36 msg_sender: Address,
37 call: ITIP20::distributeRewardCall,
38 ) -> Result<()> {
39 self.check_not_paused()?;
40 let token_address = self.address;
41
42 if call.amount == U256::ZERO {
43 return Err(TIP20Error::invalid_amount().into());
44 }
45
46 self.ensure_transfer_authorized(msg_sender, token_address)?;
47 self.check_and_update_spending_limit(msg_sender, call.amount)?;
48
49 self._transfer(msg_sender, &Recipient::direct(token_address), call.amount)?;
50
51 let opted_in_supply = U256::from(self.get_opted_in_supply()?);
52 if opted_in_supply.is_zero() {
53 return Err(TIP20Error::no_opted_in_supply().into());
54 }
55
56 let delta_rpt = call
57 .amount
58 .checked_mul(ACC_PRECISION)
59 .and_then(|v| v.checked_div(opted_in_supply))
60 .ok_or(TempoPrecompileError::under_overflow())?;
61 let current_rpt = self.get_global_reward_per_token()?;
62 let new_rpt = current_rpt
63 .checked_add(delta_rpt)
64 .ok_or(TempoPrecompileError::under_overflow())?;
65 self.set_global_reward_per_token(new_rpt)?;
66
67 self.emit_event(TIP20Event::RewardDistributed(ITIP20::RewardDistributed {
69 funder: msg_sender,
70 amount: call.amount,
71 }))?;
72
73 Ok(())
74 }
75
76 pub fn update_rewards(&mut self, holder: Address) -> Result<Address> {
83 let mut info = self.user_reward_info[holder].read()?;
84
85 let cached_delegate = info.reward_recipient;
86
87 let global_reward_per_token = self.get_global_reward_per_token()?;
88 let reward_per_token_delta = global_reward_per_token
89 .checked_sub(info.reward_per_token)
90 .ok_or(TempoPrecompileError::under_overflow())?;
91
92 if reward_per_token_delta != U256::ZERO {
93 if cached_delegate != Address::ZERO {
94 let holder_balance = self.get_balance(holder)?;
95 let reward = holder_balance
96 .checked_mul(reward_per_token_delta)
97 .and_then(|v| v.checked_div(ACC_PRECISION))
98 .ok_or(TempoPrecompileError::under_overflow())?;
99
100 if cached_delegate == holder {
102 info.reward_balance = info
103 .reward_balance
104 .checked_add(reward)
105 .ok_or(TempoPrecompileError::under_overflow())?;
106 } else {
107 let mut delegate_info = self.user_reward_info[cached_delegate].read()?;
108 delegate_info.reward_balance = delegate_info
109 .reward_balance
110 .checked_add(reward)
111 .ok_or(TempoPrecompileError::under_overflow())?;
112 self.user_reward_info[cached_delegate].write(delegate_info)?;
113 }
114 }
115 info.reward_per_token = global_reward_per_token;
116 self.user_reward_info[holder].write(info)?;
117 }
118
119 Ok(cached_delegate)
120 }
121
122 pub fn set_reward_recipient(
132 &mut self,
133 msg_sender: Address,
134 call: ITIP20::setRewardRecipientCall,
135 ) -> Result<()> {
136 self.check_not_paused()?;
137
138 if self.storage.spec().is_t3() && call.recipient.is_virtual() {
140 return Err(TIP20Error::invalid_recipient().into());
141 }
142
143 if call.recipient != Address::ZERO {
144 self.ensure_transfer_authorized(msg_sender, call.recipient)?;
145 }
146
147 let from_delegate = self.update_rewards(msg_sender)?;
148
149 let holder_balance = self.get_balance(msg_sender)?;
150
151 if from_delegate != Address::ZERO {
152 if call.recipient == Address::ZERO {
153 let opted_in_supply = U256::from(self.get_opted_in_supply()?)
154 .checked_sub(holder_balance)
155 .ok_or(TempoPrecompileError::under_overflow())?;
156 self.set_opted_in_supply(
157 opted_in_supply
158 .try_into()
159 .map_err(|_| TempoPrecompileError::under_overflow())?,
160 )?;
161 }
162 } else if call.recipient != Address::ZERO {
163 let opted_in_supply = U256::from(self.get_opted_in_supply()?)
164 .checked_add(holder_balance)
165 .ok_or(TempoPrecompileError::under_overflow())?;
166 self.set_opted_in_supply(
167 opted_in_supply
168 .try_into()
169 .map_err(|_| TempoPrecompileError::under_overflow())?,
170 )?;
171 }
172
173 let mut info = self.user_reward_info[msg_sender].read()?;
174 info.reward_recipient = call.recipient;
175 self.user_reward_info[msg_sender].write(info)?;
176
177 self.emit_event(TIP20Event::RewardRecipientSet(ITIP20::RewardRecipientSet {
179 holder: msg_sender,
180 recipient: call.recipient,
181 }))?;
182
183 Ok(())
184 }
185
186 pub fn claim_rewards(&mut self, msg_sender: Address) -> Result<U256> {
195 self.check_not_paused()?;
196 self.ensure_transfer_authorized(self.address, msg_sender)?;
197
198 self.update_rewards(msg_sender)?;
199
200 let mut info = self.user_reward_info[msg_sender].read()?;
201 let amount = info.reward_balance;
202 let contract_address = self.address;
203 let contract_balance = self.get_balance(contract_address)?;
204 let max_amount = amount.min(contract_balance);
205
206 let reward_recipient = info.reward_recipient;
207 info.reward_balance = amount
208 .checked_sub(max_amount)
209 .ok_or(TempoPrecompileError::under_overflow())?;
210 self.user_reward_info[msg_sender].write(info)?;
211
212 if max_amount > U256::ZERO {
213 let new_contract_balance = contract_balance
214 .checked_sub(max_amount)
215 .ok_or(TempoPrecompileError::under_overflow())?;
216 self.set_balance(contract_address, new_contract_balance)?;
217
218 let recipient_balance = self
219 .get_balance(msg_sender)?
220 .checked_add(max_amount)
221 .ok_or(TempoPrecompileError::under_overflow())?;
222 self.set_balance(msg_sender, recipient_balance)?;
223
224 if reward_recipient != Address::ZERO {
225 let opted_in_supply = U256::from(self.get_opted_in_supply()?)
226 .checked_add(max_amount)
227 .ok_or(TempoPrecompileError::under_overflow())?;
228 self.set_opted_in_supply(
229 opted_in_supply
230 .try_into()
231 .map_err(|_| TempoPrecompileError::under_overflow())?,
232 )?;
233 }
234
235 self.emit_event(TIP20Event::Transfer(ITIP20::Transfer {
236 from: contract_address,
237 to: msg_sender,
238 amount: max_amount,
239 }))?;
240 }
241
242 Ok(max_amount)
243 }
244
245 pub fn get_global_reward_per_token(&self) -> Result<U256> {
247 self.global_reward_per_token.read()
248 }
249
250 fn set_global_reward_per_token(&mut self, value: U256) -> Result<()> {
252 self.global_reward_per_token.write(value)
253 }
254
255 pub fn get_opted_in_supply(&self) -> Result<u128> {
257 self.opted_in_supply.read()
258 }
259
260 pub fn set_opted_in_supply(&mut self, value: u128) -> Result<()> {
262 self.opted_in_supply.write(value)
263 }
264
265 pub fn handle_rewards_on_transfer(
267 &mut self,
268 from: Address,
269 to: Address,
270 amount: U256,
271 ) -> Result<()> {
272 let from_delegate = self.update_rewards(from)?;
273 let to_delegate = self.update_rewards(to)?;
274
275 if !from_delegate.is_zero() {
276 if to_delegate.is_zero() {
277 let opted_in_supply = U256::from(self.get_opted_in_supply()?)
278 .checked_sub(amount)
279 .ok_or(TempoPrecompileError::under_overflow())?;
280 self.set_opted_in_supply(
281 opted_in_supply
282 .try_into()
283 .map_err(|_| TempoPrecompileError::under_overflow())?,
284 )?;
285 }
286 } else if !to_delegate.is_zero() {
287 let opted_in_supply = U256::from(self.get_opted_in_supply()?)
288 .checked_add(amount)
289 .ok_or(TempoPrecompileError::under_overflow())?;
290 self.set_opted_in_supply(
291 opted_in_supply
292 .try_into()
293 .map_err(|_| TempoPrecompileError::under_overflow())?,
294 )?;
295 }
296
297 Ok(())
298 }
299
300 pub fn handle_rewards_on_mint(&mut self, to: Address, amount: U256) -> Result<()> {
302 let to_delegate = self.update_rewards(to)?;
303
304 if !to_delegate.is_zero() {
305 let opted_in_supply = U256::from(self.get_opted_in_supply()?)
306 .checked_add(amount)
307 .ok_or(TempoPrecompileError::under_overflow())?;
308 self.set_opted_in_supply(
309 opted_in_supply
310 .try_into()
311 .map_err(|_| TempoPrecompileError::under_overflow())?,
312 )?;
313 }
314
315 Ok(())
316 }
317
318 pub fn get_user_reward_info(&self, account: Address) -> Result<UserRewardInfo> {
320 self.user_reward_info[account].read()
321 }
322
323 pub fn get_pending_rewards(&self, account: Address) -> Result<u128> {
332 let info = self.user_reward_info[account].read()?;
333
334 let mut pending = info.reward_balance;
336
337 if info.reward_recipient == account {
339 let holder_balance = self.get_balance(account)?;
340 if holder_balance > U256::ZERO {
341 let global_reward_per_token = self.get_global_reward_per_token()?;
342 let reward_per_token_delta = global_reward_per_token
343 .checked_sub(info.reward_per_token)
344 .ok_or(TempoPrecompileError::under_overflow())?;
345
346 if reward_per_token_delta > U256::ZERO {
347 let accrued = holder_balance
348 .checked_mul(reward_per_token_delta)
349 .and_then(|v| v.checked_div(ACC_PRECISION))
350 .ok_or(TempoPrecompileError::under_overflow())?;
351 pending = pending
352 .checked_add(accrued)
353 .ok_or(TempoPrecompileError::under_overflow())?;
354 }
355 }
356 }
357
358 pending
359 .try_into()
360 .map_err(|_| TempoPrecompileError::under_overflow())
361 }
362}
363
364#[derive(Debug, Clone, Storable)]
366pub struct UserRewardInfo {
367 pub reward_recipient: Address,
369 pub reward_per_token: U256,
371 pub reward_balance: U256,
373}
374
375impl From<UserRewardInfo> for ITIP20::UserRewardInfo {
376 fn from(value: UserRewardInfo) -> Self {
377 Self {
378 rewardRecipient: value.reward_recipient,
379 rewardPerToken: value.reward_per_token,
380 rewardBalance: value.reward_balance,
381 }
382 }
383}
384
385#[cfg(test)]
386mod tests {
387 use super::*;
388 use crate::{
389 address_registry::{MasterId, UserTag},
390 error::TempoPrecompileError,
391 storage::{StorageCtx, hashmap::HashMapStorageProvider},
392 test_util::TIP20Setup,
393 tip403_registry::TIP403Registry,
394 };
395 use alloy::primitives::{Address, U256};
396 use tempo_chainspec::hardfork::TempoHardfork;
397 use tempo_contracts::precompiles::{ITIP403Registry, TIP20Error};
398
399 #[test]
400 fn test_set_reward_recipient() -> eyre::Result<()> {
401 let mut storage = HashMapStorageProvider::new(1);
402 let admin = Address::random();
403 let alice = Address::random();
404 let amount = U256::random() % U256::from(u128::MAX);
405
406 StorageCtx::enter(&mut storage, || {
407 let mut token = TIP20Setup::create("Test", "TST", admin)
408 .with_issuer(admin)
409 .with_mint(alice, amount)
410 .apply()?;
411
412 token
413 .set_reward_recipient(alice, ITIP20::setRewardRecipientCall { recipient: alice })?;
414
415 let info = token.user_reward_info[alice].read()?;
416 assert_eq!(info.reward_recipient, alice);
417 assert_eq!(token.get_opted_in_supply()?, amount.to::<u128>());
418 assert_eq!(info.reward_per_token, U256::ZERO);
419
420 token.set_reward_recipient(
421 alice,
422 ITIP20::setRewardRecipientCall {
423 recipient: Address::ZERO,
424 },
425 )?;
426
427 let info = token.user_reward_info[alice].read()?;
428 assert_eq!(info.reward_recipient, Address::ZERO);
429 assert_eq!(token.get_opted_in_supply()?, 0u128);
430 assert_eq!(info.reward_per_token, U256::ZERO);
431
432 Ok(())
433 })
434 }
435
436 #[test]
437 fn test_distribute_reward() -> eyre::Result<()> {
438 let mut storage = HashMapStorageProvider::new(1);
439 let admin = Address::random();
440 let alice = Address::random();
441 let amount = U256::from(1000);
442 let reward_amount = amount / U256::from(10);
443
444 StorageCtx::enter(&mut storage, || {
445 let mut token = TIP20Setup::create("Test", "TST", admin)
446 .with_issuer(admin)
447 .with_mint(alice, amount)
448 .with_mint(admin, reward_amount)
449 .apply()?;
450
451 token
452 .set_reward_recipient(alice, ITIP20::setRewardRecipientCall { recipient: alice })?;
453
454 token.distribute_reward(
456 admin,
457 ITIP20::distributeRewardCall {
458 amount: reward_amount,
459 },
460 )?;
461
462 let expected_rpt = reward_amount * ACC_PRECISION / amount;
464 assert_eq!(token.get_global_reward_per_token()?, expected_rpt);
465
466 assert_eq!(token.get_balance(token.address)?, reward_amount);
468 assert_eq!(token.get_balance(admin)?, U256::ZERO);
469
470 token.update_rewards(alice)?;
472 let info = token.get_user_reward_info(alice)?;
473 assert_eq!(info.reward_balance, reward_amount);
474
475 let claimed = token.claim_rewards(alice)?;
477 assert_eq!(claimed, reward_amount);
478 assert_eq!(token.get_balance(alice)?, amount + reward_amount);
479 assert_eq!(token.get_balance(token.address)?, U256::ZERO);
480
481 token.mint(
483 admin,
484 ITIP20::mintCall {
485 to: admin,
486 amount: U256::from(1),
487 },
488 )?;
489 let result =
490 token.distribute_reward(admin, ITIP20::distributeRewardCall { amount: U256::ZERO });
491 assert!(result.is_err());
492
493 Ok(())
494 })
495 }
496
497 #[test]
498 fn test_get_pending_rewards() -> eyre::Result<()> {
499 let mut storage = HashMapStorageProvider::new(1);
500 let admin = Address::random();
501 let alice = Address::random();
502
503 StorageCtx::enter(&mut storage, || {
504 let alice_balance = U256::from(1000e18);
505 let reward_amount = U256::from(100e18);
506
507 let mut token = TIP20Setup::create("Test", "TST", admin)
508 .with_issuer(admin)
509 .with_mint(alice, alice_balance)
510 .with_mint(admin, reward_amount)
511 .apply()?;
512
513 token
514 .set_reward_recipient(alice, ITIP20::setRewardRecipientCall { recipient: alice })?;
515
516 let pending_before = token.get_pending_rewards(alice)?;
518 assert_eq!(pending_before, 0u128);
519
520 token.distribute_reward(
522 admin,
523 ITIP20::distributeRewardCall {
524 amount: reward_amount,
525 },
526 )?;
527
528 let pending_after = token.get_pending_rewards(alice)?;
530 assert_eq!(U256::from(pending_after), reward_amount);
531
532 let user_info = token.get_user_reward_info(alice)?;
534 assert_eq!(
535 user_info.reward_balance,
536 U256::ZERO,
537 "get_pending_rewards should not modify state"
538 );
539
540 Ok(())
541 })
542 }
543
544 #[test]
545 fn test_get_pending_rewards_includes_stored_balance() -> eyre::Result<()> {
546 let mut storage = HashMapStorageProvider::new(1);
547 let admin = Address::random();
548 let alice = Address::random();
549
550 StorageCtx::enter(&mut storage, || {
551 let alice_balance = U256::from(1000e18);
552 let reward_amount = U256::from(50e18);
553
554 let mut token = TIP20Setup::create("Test", "TST", admin)
555 .with_issuer(admin)
556 .with_mint(alice, alice_balance)
557 .with_mint(admin, reward_amount * U256::from(2))
558 .apply()?;
559
560 token
561 .set_reward_recipient(alice, ITIP20::setRewardRecipientCall { recipient: alice })?;
562
563 token.distribute_reward(
565 admin,
566 ITIP20::distributeRewardCall {
567 amount: reward_amount,
568 },
569 )?;
570
571 token.update_rewards(alice)?;
573 let user_info = token.get_user_reward_info(alice)?;
574 assert_eq!(user_info.reward_balance, reward_amount);
575
576 token.distribute_reward(
578 admin,
579 ITIP20::distributeRewardCall {
580 amount: reward_amount,
581 },
582 )?;
583
584 let pending = token.get_pending_rewards(alice)?;
586 assert_eq!(U256::from(pending), reward_amount * U256::from(2));
587
588 Ok(())
589 })
590 }
591
592 #[test]
593 fn test_get_pending_rewards_with_delegation() -> eyre::Result<()> {
594 let mut storage = HashMapStorageProvider::new(1);
595 let admin = Address::random();
596 let alice = Address::random();
597 let bob = Address::random();
598
599 StorageCtx::enter(&mut storage, || {
600 let alice_balance = U256::from(1000e18);
601 let reward_amount = U256::from(100e18);
602
603 let mut token = TIP20Setup::create("Test", "TST", admin)
604 .with_issuer(admin)
605 .with_mint(alice, alice_balance)
606 .with_mint(admin, reward_amount)
607 .apply()?;
608
609 token.set_reward_recipient(alice, ITIP20::setRewardRecipientCall { recipient: bob })?;
611
612 token.distribute_reward(
614 admin,
615 ITIP20::distributeRewardCall {
616 amount: reward_amount,
617 },
618 )?;
619
620 let alice_pending = token.get_pending_rewards(alice)?;
622 assert_eq!(alice_pending, 0u128);
623
624 let bob_pending_before_update = token.get_pending_rewards(bob)?;
628 assert_eq!(bob_pending_before_update, 0u128);
629
630 token.update_rewards(alice)?;
632 let bob_pending_after_update = token.get_pending_rewards(bob)?;
633 assert_eq!(U256::from(bob_pending_after_update), reward_amount);
634
635 Ok(())
636 })
637 }
638
639 #[test]
640 fn test_get_pending_rewards_not_opted_in() -> eyre::Result<()> {
641 let mut storage = HashMapStorageProvider::new(1);
642 let admin = Address::random();
643 let alice = Address::random();
644 let bob = Address::random();
645
646 StorageCtx::enter(&mut storage, || {
647 let balance = U256::from(1000e18);
648 let reward_amount = U256::from(100e18);
649
650 let mut token = TIP20Setup::create("Test", "TST", admin)
651 .with_issuer(admin)
652 .with_mint(alice, balance)
653 .with_mint(bob, balance)
654 .with_mint(admin, reward_amount)
655 .apply()?;
656
657 token
659 .set_reward_recipient(alice, ITIP20::setRewardRecipientCall { recipient: alice })?;
660
661 token.distribute_reward(
663 admin,
664 ITIP20::distributeRewardCall {
665 amount: reward_amount,
666 },
667 )?;
668
669 let alice_pending = token.get_pending_rewards(alice)?;
671 assert_eq!(U256::from(alice_pending), reward_amount);
672
673 let bob_pending = token.get_pending_rewards(bob)?;
675 assert_eq!(bob_pending, 0u128);
676
677 Ok(())
678 })
679 }
680
681 #[test]
682 fn test_claim_rewards_unauthorized() -> eyre::Result<()> {
683 let mut storage = HashMapStorageProvider::new(1);
684 let admin = Address::random();
685 let alice = Address::random();
686
687 StorageCtx::enter(&mut storage, || {
688 let mut registry = TIP403Registry::new();
689 registry.initialize()?;
690
691 let policy_id = registry.create_policy(
692 admin,
693 ITIP403Registry::createPolicyCall {
694 admin,
695 policyType: ITIP403Registry::PolicyType::BLACKLIST,
696 },
697 )?;
698
699 registry.modify_policy_blacklist(
700 admin,
701 ITIP403Registry::modifyPolicyBlacklistCall {
702 policyId: policy_id,
703 account: alice,
704 restricted: true,
705 },
706 )?;
707
708 let mut token = TIP20Setup::create("Test", "TST", admin).apply()?;
709
710 token.change_transfer_policy_id(
711 admin,
712 ITIP20::changeTransferPolicyIdCall {
713 newPolicyId: policy_id,
714 },
715 )?;
716
717 let err = token.claim_rewards(alice).unwrap_err();
718 assert!(
719 matches!(
720 err,
721 TempoPrecompileError::TIP20(TIP20Error::PolicyForbids(_))
722 ),
723 "Expected PolicyForbids error, got: {err:?}"
724 );
725
726 Ok(())
727 })
728 }
729
730 #[test]
731 fn test_set_reward_recipient_rejects_virtual_on_t3() -> eyre::Result<()> {
732 let virtual_addr = Address::new_virtual(MasterId::ZERO, UserTag::ZERO);
733
734 for hardfork in [TempoHardfork::T2, TempoHardfork::T3] {
735 let mut storage = HashMapStorageProvider::new_with_spec(1, hardfork);
736 let admin = Address::random();
737 let alice = Address::random();
738
739 StorageCtx::enter(&mut storage, || {
740 let mut token = TIP20Setup::create("Test", "TST", admin)
741 .with_issuer(admin)
742 .with_mint(alice, U256::from(1000))
743 .apply()?;
744
745 let result = token.set_reward_recipient(
746 alice,
747 ITIP20::setRewardRecipientCall {
748 recipient: virtual_addr,
749 },
750 );
751
752 if hardfork.is_t3() {
753 assert!(matches!(
754 result.unwrap_err(),
755 TempoPrecompileError::TIP20(TIP20Error::InvalidRecipient(_))
756 ));
757 } else {
758 assert!(result.is_ok());
760 }
761
762 Ok::<_, TempoPrecompileError>(())
763 })?;
764 }
765 Ok(())
766 }
767}