From 8d23b61a16274e5e3a6e6d6543cb46eae96fc02f Mon Sep 17 00:00:00 2001
From: Emma Zhong <emma@mystenlabs.com>
Date: Tue, 5 Nov 2024 16:13:13 -0800
Subject: [PATCH] [jsonrpc] fix estimated rewards during safe mode (#20182)

## Description

Currently if the exchange rate for an epoch is not found, we would
assign it the default exchange rate while we should've used the exchange
rate from the previous existing epoch instead. This PR fixes that.

## Test plan

Added some unit tests.

---

## Release notes

Check each box that your changes affect. If none of the boxes relate to
your changes, release notes aren't required.

For each box you select, include information after the relevant heading
that describes the impact of your changes that a user might notice and
any actions they must take to implement updates.

- [ ] Protocol:
- [ ] Nodes (Validators and Full nodes):
- [ ] Indexer:
- [x] JSON-RPC: fixed reward estimation in `getStakes` API so that we
don't overestimate stake rewards of stakes activated during safe mode
epochs.
- [ ] GraphQL:
- [ ] CLI:
- [ ] Rust SDK:
- [ ] REST API:
---
 crates/sui-json-rpc/src/governance_api.rs    | 77 +++++++++++++++++++-
 crates/sui-types/src/sui_system_state/mod.rs |  7 ++
 2 files changed, 83 insertions(+), 1 deletion(-)

diff --git a/crates/sui-json-rpc/src/governance_api.rs b/crates/sui-json-rpc/src/governance_api.rs
index a8a23a41952ed..f08d4366264ce 100644
--- a/crates/sui-json-rpc/src/governance_api.rs
+++ b/crates/sui-json-rpc/src/governance_api.rs
@@ -166,6 +166,7 @@ impl GovernanceReadApi {
                 let status = if !exists {
                     StakeStatus::Unstaked
                 } else if system_state_summary.epoch >= stake.activation_epoch() {
+                    // TODO: use dev_inspect to call a move function to get the estimated reward
                     let estimated_reward = if let Some(current_rate) = current_rate {
                         let stake_rate = rate_table
                             .rates
@@ -431,7 +432,8 @@ async fn exchange_rates(
             })
             .collect::<Result<Vec<_>, _>>()?;
 
-        rates.sort_by(|(a, _), (b, _)| a.cmp(b).reverse());
+        // Rates for some epochs might be missing due to safe mode, we need to backfill them.
+        rates = backfill_rates(rates);
 
         exchange_rates.push(ValidatorExchangeRates {
             address,
@@ -451,6 +453,38 @@ pub struct ValidatorExchangeRates {
     pub rates: Vec<(EpochId, PoolTokenExchangeRate)>,
 }
 
+/// Backfill missing rates for some epochs due to safe mode. If a rate is missing for epoch e,
+/// we will use the rate for epoch e-1 to fill it.
+/// Rates returned are in descending order by epoch.
+fn backfill_rates(
+    rates: Vec<(EpochId, PoolTokenExchangeRate)>,
+) -> Vec<(EpochId, PoolTokenExchangeRate)> {
+    if rates.is_empty() {
+        return rates;
+    }
+
+    let min_epoch = *rates.iter().map(|(e, _)| e).min().unwrap();
+    let max_epoch = *rates.iter().map(|(e, _)| e).max().unwrap();
+    let mut filled_rates = Vec::new();
+    let mut prev_rate = None;
+
+    for epoch in min_epoch..=max_epoch {
+        match rates.iter().find(|(e, _)| *e == epoch) {
+            Some((e, rate)) => {
+                prev_rate = Some(rate.clone());
+                filled_rates.push((*e, rate.clone()));
+            }
+            None => {
+                if let Some(rate) = prev_rate.clone() {
+                    filled_rates.push((epoch, rate));
+                }
+            }
+        }
+    }
+    filled_rates.reverse();
+    filled_rates
+}
+
 impl SuiRpcModule for GovernanceReadApi {
     fn rpc(self) -> RpcModule<Self> {
         self.into_rpc()
@@ -460,3 +494,44 @@ impl SuiRpcModule for GovernanceReadApi {
         GovernanceReadApiOpenRpc::module_doc()
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use sui_types::sui_system_state::PoolTokenExchangeRate;
+
+    #[test]
+    fn test_backfill_rates_empty() {
+        let rates = vec![];
+        assert_eq!(backfill_rates(rates), vec![]);
+    }
+
+    #[test]
+    fn test_backfill_rates_no_gaps() {
+        let rate1 = PoolTokenExchangeRate::new(100, 100);
+        let rate2 = PoolTokenExchangeRate::new(200, 220);
+        let rate3 = PoolTokenExchangeRate::new(300, 330);
+        let rates = vec![(2, rate2.clone()), (3, rate3.clone()), (1, rate1.clone())];
+
+        let expected: Vec<(u64, PoolTokenExchangeRate)> =
+            vec![(3, rate3.clone()), (2, rate2), (1, rate1)];
+        assert_eq!(backfill_rates(rates), expected);
+    }
+
+    #[test]
+    fn test_backfill_rates_with_gaps() {
+        let rate1 = PoolTokenExchangeRate::new(100, 100);
+        let rate3 = PoolTokenExchangeRate::new(300, 330);
+        let rate5 = PoolTokenExchangeRate::new(500, 550);
+        let rates = vec![(3, rate3.clone()), (1, rate1.clone()), (5, rate5.clone())];
+
+        let expected = vec![
+            (5, rate5.clone()),
+            (4, rate3.clone()),
+            (3, rate3.clone()),
+            (2, rate1.clone()),
+            (1, rate1),
+        ];
+        assert_eq!(backfill_rates(rates), expected);
+    }
+}
diff --git a/crates/sui-types/src/sui_system_state/mod.rs b/crates/sui-types/src/sui_system_state/mod.rs
index d647e7da495b6..7f91dbd292e98 100644
--- a/crates/sui-types/src/sui_system_state/mod.rs
+++ b/crates/sui-types/src/sui_system_state/mod.rs
@@ -416,6 +416,13 @@ impl PoolTokenExchangeRate {
             self.pool_token_amount as f64 / self.sui_amount as f64
         }
     }
+
+    pub fn new(sui_amount: u64, pool_token_amount: u64) -> Self {
+        Self {
+            sui_amount,
+            pool_token_amount,
+        }
+    }
 }
 
 #[derive(Debug, Serialize, Deserialize, Clone, Eq, PartialEq)]