[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

@@ -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);
+    }
+}

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)]