[jsonrpc] fix estimated rewards during safe mode

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