1use std::{collections::HashSet, iter::once, str::FromStr};
2
3use alloy::primitives::{
4 utils::{parse_units, ParseUnits},
5 Address, U256,
6};
7use anyhow::{bail, ensure, Context};
8use ark_serialize::{
9 CanonicalDeserialize, CanonicalSerialize, Compress, Read, SerializationError, Valid, Validate,
10};
11use committable::{Commitment, Committable, RawCommitmentBuilder};
12use hotshot::types::BLSPubKey;
13use hotshot_types::{
14 data::{EpochNumber, ViewNumber},
15 traits::{election::Membership, node_implementation::ConsensusTime},
16 utils::epoch_from_block_number,
17};
18use jf_merkle_tree::{
19 ForgetableMerkleTreeScheme, ForgetableUniversalMerkleTreeScheme, LookupResult,
20 MerkleCommitment, MerkleTreeScheme, PersistentUniversalMerkleTreeScheme, ToTraversalPath,
21 UniversalMerkleTreeScheme,
22};
23use num_traits::CheckedSub;
24use sequencer_utils::{
25 impl_serde_from_string_or_integer, impl_to_fixed_bytes, ser::FromStringOrInteger,
26};
27
28use super::{
29 v0_1::{
30 RewardAccount, RewardAccountProof, RewardAccountQueryData, RewardAmount, RewardInfo,
31 RewardMerkleCommitment, RewardMerkleProof, RewardMerkleTree, COMMISSION_BASIS_POINTS,
32 },
33 v0_3::Validator,
34 Leaf2, NodeState, ValidatedState,
35};
36use crate::{eth_signature_key::EthKeyPair, Delta, FeeAccount};
37
38impl Committable for RewardInfo {
39 fn commit(&self) -> Commitment<Self> {
40 RawCommitmentBuilder::new(&Self::tag())
41 .fixed_size_field("account", &self.account.to_fixed_bytes())
42 .fixed_size_field("amount", &self.amount.to_fixed_bytes())
43 .finalize()
44 }
45 fn tag() -> String {
46 "REWARD_INFO".into()
47 }
48}
49
50impl_serde_from_string_or_integer!(RewardAmount);
51impl_to_fixed_bytes!(RewardAmount, U256);
52
53impl From<u64> for RewardAmount {
54 fn from(amt: u64) -> Self {
55 Self(U256::from(amt))
56 }
57}
58
59impl CheckedSub for RewardAmount {
60 fn checked_sub(&self, v: &Self) -> Option<Self> {
61 self.0.checked_sub(v.0).map(RewardAmount)
62 }
63}
64
65impl FromStr for RewardAmount {
66 type Err = <U256 as FromStr>::Err;
67
68 fn from_str(s: &str) -> Result<Self, Self::Err> {
69 Ok(Self(s.parse()?))
70 }
71}
72
73impl FromStringOrInteger for RewardAmount {
74 type Binary = U256;
75 type Integer = u64;
76
77 fn from_binary(b: Self::Binary) -> anyhow::Result<Self> {
78 Ok(Self(b))
79 }
80
81 fn from_integer(i: Self::Integer) -> anyhow::Result<Self> {
82 Ok(i.into())
83 }
84
85 fn from_string(s: String) -> anyhow::Result<Self> {
86 if let Some(s) = s.strip_prefix("0x") {
89 return Ok(Self(s.parse()?));
90 }
91
92 let (base, unit) = s
94 .split_once(char::is_whitespace)
95 .unwrap_or((s.as_str(), "wei"));
96 match parse_units(base, unit)? {
97 ParseUnits::U256(n) => Ok(Self(n)),
98 ParseUnits::I256(_) => bail!("amount cannot be negative"),
99 }
100 }
101
102 fn to_binary(&self) -> anyhow::Result<Self::Binary> {
103 Ok(self.0)
104 }
105
106 fn to_string(&self) -> anyhow::Result<String> {
107 Ok(format!("{self}"))
108 }
109}
110
111impl RewardAmount {
112 pub fn as_u64(&self) -> Option<u64> {
113 if self.0 <= U256::from(u64::MAX) {
114 Some(self.0.to::<u64>())
115 } else {
116 None
117 }
118 }
119}
120impl RewardAccount {
121 pub fn address(&self) -> Address {
123 self.0
124 }
125 pub fn as_bytes(&self) -> &[u8] {
127 self.0.as_slice()
128 }
129 pub fn to_fixed_bytes(self) -> [u8; 20] {
131 self.0.into_array()
132 }
133 pub fn test_key_pair() -> EthKeyPair {
134 EthKeyPair::from_mnemonic(
135 "test test test test test test test test test test test junk",
136 0u32,
137 )
138 .unwrap()
139 }
140}
141
142impl FromStr for RewardAccount {
143 type Err = anyhow::Error;
144
145 fn from_str(s: &str) -> Result<Self, Self::Err> {
146 Ok(Self(s.parse()?))
147 }
148}
149
150impl Valid for RewardAmount {
151 fn check(&self) -> Result<(), SerializationError> {
152 Ok(())
153 }
154}
155
156impl Valid for RewardAccount {
157 fn check(&self) -> Result<(), SerializationError> {
158 Ok(())
159 }
160}
161
162impl CanonicalSerialize for RewardAmount {
163 fn serialize_with_mode<W: std::io::prelude::Write>(
164 &self,
165 mut writer: W,
166 _compress: Compress,
167 ) -> Result<(), SerializationError> {
168 Ok(writer.write_all(&self.to_fixed_bytes())?)
169 }
170
171 fn serialized_size(&self, _compress: Compress) -> usize {
172 core::mem::size_of::<U256>()
173 }
174}
175impl CanonicalDeserialize for RewardAmount {
176 fn deserialize_with_mode<R: Read>(
177 mut reader: R,
178 _compress: Compress,
179 _validate: Validate,
180 ) -> Result<Self, SerializationError> {
181 let mut bytes = [0u8; core::mem::size_of::<U256>()];
182 reader.read_exact(&mut bytes)?;
183 let value = U256::from_le_slice(&bytes);
184 Ok(Self(value))
185 }
186}
187impl CanonicalSerialize for RewardAccount {
188 fn serialize_with_mode<W: std::io::prelude::Write>(
189 &self,
190 mut writer: W,
191 _compress: Compress,
192 ) -> Result<(), SerializationError> {
193 Ok(writer.write_all(self.0.as_slice())?)
194 }
195
196 fn serialized_size(&self, _compress: Compress) -> usize {
197 core::mem::size_of::<Address>()
198 }
199}
200impl CanonicalDeserialize for RewardAccount {
201 fn deserialize_with_mode<R: Read>(
202 mut reader: R,
203 _compress: Compress,
204 _validate: Validate,
205 ) -> Result<Self, SerializationError> {
206 let mut bytes = [0u8; core::mem::size_of::<Address>()];
207 reader.read_exact(&mut bytes)?;
208 let value = Address::from_slice(&bytes);
209 Ok(Self(value))
210 }
211}
212
213impl ToTraversalPath<256> for RewardAccount {
214 fn to_traversal_path(&self, height: usize) -> Vec<usize> {
215 self.0
216 .as_slice()
217 .iter()
218 .take(height)
219 .map(|i| *i as usize)
220 .collect()
221 }
222}
223
224#[allow(dead_code)]
225impl RewardAccountProof {
226 pub fn presence(
227 pos: FeeAccount,
228 proof: <RewardMerkleTree as MerkleTreeScheme>::MembershipProof,
229 ) -> Self {
230 Self {
231 account: pos.into(),
232 proof: RewardMerkleProof::Presence(proof),
233 }
234 }
235
236 pub fn absence(
237 pos: RewardAccount,
238 proof: <RewardMerkleTree as UniversalMerkleTreeScheme>::NonMembershipProof,
239 ) -> Self {
240 Self {
241 account: pos.into(),
242 proof: RewardMerkleProof::Absence(proof),
243 }
244 }
245
246 pub fn prove(tree: &RewardMerkleTree, account: Address) -> Option<(Self, U256)> {
247 match tree.universal_lookup(RewardAccount(account)) {
248 LookupResult::Ok(balance, proof) => Some((
249 Self {
250 account,
251 proof: RewardMerkleProof::Presence(proof),
252 },
253 balance.0,
254 )),
255 LookupResult::NotFound(proof) => Some((
256 Self {
257 account,
258 proof: RewardMerkleProof::Absence(proof),
259 },
260 U256::ZERO,
261 )),
262 LookupResult::NotInMemory => None,
263 }
264 }
265
266 pub fn verify(&self, comm: &RewardMerkleCommitment) -> anyhow::Result<U256> {
267 match &self.proof {
268 RewardMerkleProof::Presence(proof) => {
269 ensure!(
270 RewardMerkleTree::verify(comm.digest(), RewardAccount(self.account), proof)?
271 .is_ok(),
272 "invalid proof"
273 );
274 Ok(proof
275 .elem()
276 .context("presence proof is missing account balance")?
277 .0)
278 },
279 RewardMerkleProof::Absence(proof) => {
280 let tree = RewardMerkleTree::from_commitment(comm);
281 ensure!(
282 tree.non_membership_verify(RewardAccount(self.account), proof)?,
283 "invalid proof"
284 );
285 Ok(U256::ZERO)
286 },
287 }
288 }
289
290 pub fn remember(&self, tree: &mut RewardMerkleTree) -> anyhow::Result<()> {
291 match &self.proof {
292 RewardMerkleProof::Presence(proof) => {
293 tree.remember(
294 RewardAccount(self.account),
295 proof
296 .elem()
297 .context("presence proof is missing account balance")?,
298 proof,
299 )?;
300 Ok(())
301 },
302 RewardMerkleProof::Absence(proof) => {
303 tree.non_membership_remember(RewardAccount(self.account), proof)?;
304 Ok(())
305 },
306 }
307 }
308}
309
310impl From<(RewardAccountProof, U256)> for RewardAccountQueryData {
311 fn from((proof, balance): (RewardAccountProof, U256)) -> Self {
312 Self { balance, proof }
313 }
314}
315
316#[derive(Clone, Debug)]
317pub struct ComputedRewards {
318 leader_address: Address,
319 leader_commission: RewardAmount,
321 delegators: Vec<(Address, RewardAmount)>,
323}
324
325impl ComputedRewards {
326 pub fn new(
327 delegators: Vec<(Address, RewardAmount)>,
328 leader_address: Address,
329 leader_commission: RewardAmount,
330 ) -> Self {
331 Self {
332 delegators,
333 leader_address,
334 leader_commission,
335 }
336 }
337
338 pub fn leader_commission(&self) -> &RewardAmount {
339 &self.leader_commission
340 }
341
342 pub fn delegators(&self) -> &Vec<(Address, RewardAmount)> {
343 &self.delegators
344 }
345
346 pub fn all_rewards(self) -> Vec<(Address, RewardAmount)> {
348 self.delegators
349 .into_iter()
350 .chain(once((self.leader_address, self.leader_commission)))
351 .collect()
352 }
353}
354
355pub struct RewardDistributor {
356 validator: Validator<BLSPubKey>,
357 block_reward: RewardAmount,
358}
359
360impl RewardDistributor {
361 pub fn new(validator: Validator<BLSPubKey>, block_reward: RewardAmount) -> Self {
362 Self {
363 validator,
364 block_reward,
365 }
366 }
367
368 pub fn validator(&self) -> Validator<BLSPubKey> {
369 self.validator.clone()
370 }
371
372 pub fn block_reward(&self) -> RewardAmount {
373 self.block_reward
374 }
375
376 pub fn distribute(&self, state: &mut ValidatedState, delta: &mut Delta) -> anyhow::Result<()> {
377 let reward_state = self.apply_rewards(state.reward_merkle_tree.clone())?;
378 state.reward_merkle_tree = reward_state;
379
380 delta
382 .rewards_delta
383 .insert(RewardAccount(self.validator().account));
384 delta.rewards_delta.extend(
385 self.validator()
386 .delegators
387 .keys()
388 .map(|d| RewardAccount(*d)),
389 );
390
391 Ok(())
392 }
393
394 pub fn apply_rewards(
395 &self,
396 mut reward_state: RewardMerkleTree,
397 ) -> anyhow::Result<RewardMerkleTree> {
398 let mut update_balance = |account: &RewardAccount, amount: RewardAmount| {
399 let mut err = None;
400 reward_state = reward_state.persistent_update_with(account, |balance| {
401 let balance = balance.copied();
402 match balance.unwrap_or_default().0.checked_add(amount.0) {
403 Some(updated) => Some(updated.into()),
404 None => {
405 err = Some(format!("overflowed reward balance for account {account}"));
406 balance
407 },
408 }
409 })?;
410
411 if let Some(error) = err {
412 tracing::warn!(error);
413 bail!(error)
414 }
415 Ok::<(), anyhow::Error>(())
416 };
417 let computed_rewards = self.compute_rewards()?;
418 for (address, reward) in computed_rewards.all_rewards() {
419 update_balance(&RewardAccount(address), reward)?;
420 tracing::debug!("applied rewards address={address} reward={reward}",);
421 }
422
423 Ok(reward_state)
424 }
425
426 pub fn compute_rewards(&self) -> anyhow::Result<ComputedRewards> {
434 ensure!(
435 self.validator.commission <= COMMISSION_BASIS_POINTS,
436 "commission must not exceed {COMMISSION_BASIS_POINTS}"
437 );
438
439 ensure!(self.block_reward.0 > U256::ZERO, "block reward is zero");
440
441 let mut rewards = Vec::new();
442
443 let total_reward = self.block_reward.0;
444 let delegators_ratio_basis_points = U256::from(COMMISSION_BASIS_POINTS)
445 .checked_sub(U256::from(self.validator.commission))
446 .context("overflow")?;
447 let delegators_reward = delegators_ratio_basis_points
448 .checked_mul(total_reward)
449 .context("overflow")?;
450
451 let total_stake = self.validator.stake;
453 let mut delegators_rewards_distributed = U256::from(0);
454 for (delegator_address, delegator_stake) in &self.validator.delegators {
455 let delegator_reward = RewardAmount::from(
456 (delegator_stake
457 .checked_mul(delegators_reward)
458 .context("overflow")?
459 .checked_div(total_stake)
460 .context("overflow")?)
461 .checked_div(U256::from(COMMISSION_BASIS_POINTS))
462 .context("overflow")?,
463 );
464
465 delegators_rewards_distributed += delegator_reward.0;
466
467 rewards.push((*delegator_address, delegator_reward));
468 }
469
470 let leader_commission = total_reward
471 .checked_sub(delegators_rewards_distributed)
472 .context("overflow")?;
473
474 Ok(ComputedRewards::new(
475 rewards,
476 self.validator.account,
477 leader_commission.into(),
478 ))
479 }
480}
481pub async fn first_two_epochs(height: u64, instance_state: &NodeState) -> anyhow::Result<bool> {
488 let epoch_height = instance_state
489 .epoch_height
490 .context("epoch height not found")?;
491 let epoch = EpochNumber::new(epoch_from_block_number(height, epoch_height));
492 let coordinator = instance_state.coordinator.clone();
493 let first_epoch = coordinator
494 .membership()
495 .read()
496 .await
497 .first_epoch()
498 .context("The first epoch was not set.")?;
499
500 Ok(epoch <= first_epoch + 1)
501}
502
503pub async fn find_validator_info(
504 instance_state: &NodeState,
505 validated_state: &mut ValidatedState,
506 parent_leaf: &Leaf2,
507 view: ViewNumber,
508) -> anyhow::Result<Validator<BLSPubKey>> {
509 let parent_height = parent_leaf.height();
510 let parent_view = parent_leaf.view_number();
511 let new_height = parent_height + 1;
512
513 let epoch_height = instance_state
514 .epoch_height
515 .context("epoch height not found")?;
516 if epoch_height == 0 {
517 bail!("epoch height is 0. can not catchup reward accounts");
518 }
519 let epoch = EpochNumber::new(epoch_from_block_number(new_height, epoch_height));
520
521 let coordinator = instance_state.coordinator.clone();
522
523 let epoch_membership = coordinator.membership_for_epoch(Some(epoch)).await?;
524 let membership = epoch_membership.coordinator.membership().read().await;
525
526 let leader: BLSPubKey = membership
527 .leader(view, Some(epoch))
528 .context(format!("leader for epoch {epoch:?} not found"))?;
529
530 let validator = membership
531 .get_validator_config(&epoch, leader)
532 .context("validator not found")?;
533 drop(membership);
534
535 let mut reward_accounts = HashSet::new();
536 reward_accounts.insert(validator.account.into());
537 let delegators = validator
538 .delegators
539 .keys()
540 .cloned()
541 .map(|a| a.into())
542 .collect::<Vec<RewardAccount>>();
543
544 reward_accounts.extend(delegators.clone());
545 let missing_reward_accts = validated_state.forgotten_reward_accounts(reward_accounts);
546
547 if !missing_reward_accts.is_empty() {
548 tracing::warn!(
549 parent_height,
550 ?parent_view,
551 ?missing_reward_accts,
552 "fetching missing reward accounts from peers"
553 );
554
555 let missing_account_proofs = instance_state
556 .state_catchup
557 .fetch_reward_accounts(
558 instance_state,
559 parent_height,
560 parent_view,
561 validated_state.reward_merkle_tree.commitment(),
562 missing_reward_accts,
563 )
564 .await?;
565
566 for proof in missing_account_proofs.iter() {
567 proof
568 .remember(&mut validated_state.reward_merkle_tree)
569 .expect("proof previously verified");
570 }
571 }
572 Ok(validator)
573}
574
575#[cfg(test)]
576pub mod tests {
577
578 use super::*;
579
580 #[test]
583 fn test_reward_calculation_sanity_checks() {
584 let validator = Validator::mock();
589 let mut distributor = RewardDistributor::new(
590 validator,
591 RewardAmount(U256::from(1902000000000000000_u128)),
592 );
593 let rewards = distributor.compute_rewards().unwrap();
594 let total = |rewards: ComputedRewards| {
595 rewards
596 .all_rewards()
597 .iter()
598 .fold(U256::ZERO, |acc, (_, r)| acc + r.0)
599 };
600 assert_eq!(total(rewards.clone()), distributor.block_reward.into());
601
602 distributor.validator.commission = 0;
603 let rewards = distributor.compute_rewards().unwrap();
604 assert_eq!(total(rewards.clone()), distributor.block_reward.into());
605
606 distributor.validator.commission = 10000;
607 let rewards = distributor.compute_rewards().unwrap();
608 assert_eq!(total(rewards.clone()), distributor.block_reward.into());
609 let leader_commission = rewards.leader_commission();
610 assert_eq!(*leader_commission, distributor.block_reward);
611
612 distributor.validator.commission = 10001;
613 assert!(distributor
614 .compute_rewards()
615 .err()
616 .unwrap()
617 .to_string()
618 .contains("must not exceed"));
619 }
620
621 #[test]
622 fn test_compute_rewards_validator_commission() {
623 let validator = Validator::mock();
624 let mut distributor = RewardDistributor::new(
625 validator.clone(),
626 RewardAmount(U256::from(1902000000000000000_u128)),
627 );
628 distributor.validator.commission = 0;
629
630 let rewards = distributor.compute_rewards().unwrap();
631
632 let leader_commission = rewards.leader_commission();
633 let percentage =
634 leader_commission.0 * U256::from(COMMISSION_BASIS_POINTS) / distributor.block_reward.0;
635 assert_eq!(percentage, U256::ZERO);
636
637 distributor.validator.commission = 300;
639
640 let rewards = distributor.compute_rewards().unwrap();
641 let leader_commission = rewards.leader_commission();
642 let percentage =
643 leader_commission.0 * U256::from(COMMISSION_BASIS_POINTS) / distributor.block_reward.0;
644 println!("percentage: {percentage:?}");
645 assert_eq!(percentage, U256::from(300));
646
647 distributor.validator.commission = 10000;
649
650 let rewards = distributor.compute_rewards().unwrap();
651 let leader_commission = rewards.leader_commission();
652 assert_eq!(*leader_commission, distributor.block_reward);
653 }
654}