From d493de48dc23a5b39cbcb340f888d5519bd5f337 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Tom=C3=A1=C5=A1=20Zemanovi=C4=8D?= <tomas@heliax.dev>
Date: Wed, 5 Jun 2024 16:59:20 +0200
Subject: [PATCH] shielded_token: store the total rewards minted by MASP

---
 crates/shielded_token/src/conversion.rs  | 11 +++------
 crates/shielded_token/src/storage.rs     | 31 +++++++++++++++++++++++-
 crates/shielded_token/src/storage_key.rs |  9 +++++++
 3 files changed, 43 insertions(+), 8 deletions(-)

diff --git a/crates/shielded_token/src/conversion.rs b/crates/shielded_token/src/conversion.rs
index a8e009ac784..e53c72814a6 100644
--- a/crates/shielded_token/src/conversion.rs
+++ b/crates/shielded_token/src/conversion.rs
@@ -254,13 +254,14 @@ where
     use namada_parameters as parameters;
     use namada_storage::conversion_state::ConversionLeaf;
     use namada_storage::{Error, OptionExt, ResultExt};
-    use namada_trans_token::storage_key::balance_key;
     use namada_trans_token::{MaspDigitPos, NATIVE_MAX_DECIMAL_PLACES};
     use rayon::iter::{
         IndexedParallelIterator, IntoParallelIterator, ParallelIterator,
     };
     use rayon::prelude::ParallelSlice;
 
+    use crate::mint_rewards;
+
     // The derived conversions will be placed in MASP address space
     let masp_addr = MASP;
 
@@ -566,12 +567,8 @@ where
 
     // Update the MASP's transparent reward token balance to ensure that it
     // is sufficiently backed to redeem rewards
-    let reward_key = balance_key(&native_token, &masp_addr);
-    let addr_bal: Amount = storage.read(&reward_key)?.unwrap_or_default();
-    let new_bal = addr_bal
-        .checked_add(total_reward)
-        .ok_or_else(|| Error::new_const("Balance with reward overflow"))?;
-    storage.write(&reward_key, new_bal)?;
+    mint_rewards(storage, total_reward)?;
+
     // Try to distribute Merkle tree construction as evenly as possible
     // across multiple cores
     // Merkle trees must have exactly 2^n leaves to be mergeable
diff --git a/crates/shielded_token/src/storage.rs b/crates/shielded_token/src/storage.rs
index 0b91da75a9a..bddfcb77310 100644
--- a/crates/shielded_token/src/storage.rs
+++ b/crates/shielded_token/src/storage.rs
@@ -1,10 +1,11 @@
-use namada_core::address::Address;
+use namada_core::address::{self, Address};
 use namada_core::arith::checked;
 use namada_core::token;
 use namada_core::token::Amount;
 use namada_core::uint::Uint;
 use namada_storage as storage;
 use namada_storage::{StorageRead, StorageWrite};
+use namada_trans_token::credit_tokens;
 use storage::ResultExt;
 
 use crate::storage_key::*;
@@ -40,3 +41,31 @@ where
     storage.write(&masp_locked_amount_target_key(address), raw_target)?;
     Ok(())
 }
+
+/// Mint MASP rewards tokens and increment the stored total rewards.
+pub fn mint_rewards<S>(
+    storage: &mut S,
+    amount: token::Amount,
+) -> storage::Result<()>
+where
+    S: StorageRead + StorageWrite,
+{
+    let native_token = storage.get_native_token()?;
+    credit_tokens(storage, &native_token, &address::MASP, amount)?;
+
+    let total_rewards_key = masp_total_rewards();
+    let mut total_rewards = read_total_rewards(storage)?;
+    checked!(total_rewards += amount)?;
+    storage.write(&total_rewards_key, total_rewards)
+}
+
+/// Read the total rewards minted by MASP.
+pub fn read_total_rewards<S>(storage: &S) -> storage::Result<token::Amount>
+where
+    S: StorageRead,
+{
+    let total_rewards_key = masp_total_rewards();
+    let total_rewards: token::Amount =
+        storage.read(&total_rewards_key)?.unwrap_or_default();
+    Ok(total_rewards)
+}
diff --git a/crates/shielded_token/src/storage_key.rs b/crates/shielded_token/src/storage_key.rs
index f584ab7fff2..df558b2e9e6 100644
--- a/crates/shielded_token/src/storage_key.rs
+++ b/crates/shielded_token/src/storage_key.rs
@@ -32,6 +32,8 @@ pub const MASP_KD_GAIN_KEY: &str = "derivative_gain";
 pub const MASP_LOCKED_AMOUNT_TARGET_KEY: &str = "locked_amount_target";
 /// The key for the max reward rate for a given asset
 pub const MASP_MAX_REWARD_RATE_KEY: &str = "max_reward_rate";
+/// The key for the total inflation rewards minted by MASP
+pub const MASP_TOTAL_REWARDS: &str = "max_total_rewards";
 
 /// Obtain the nominal proportional key for the given token
 pub fn masp_kp_gain_key(token_addr: &Address) -> storage::Key {
@@ -163,3 +165,10 @@ pub fn masp_assets_hash_key() -> storage::Key {
         .push(&MASP_ASSETS_HASH_KEY.to_owned())
         .expect("Cannot obtain a storage key")
 }
+
+/// The max reward rate key for the given token
+pub fn masp_total_rewards() -> storage::Key {
+    storage::Key::from(address::MASP.to_db_key())
+        .push(&MASP_TOTAL_REWARDS.to_owned())
+        .expect("Cannot obtain a storage key")
+}