Custom masp epoch tracker

.changelog/unreleased/improvements/3318-masp-custom-epoch.md

@@ -0,0 +1,2 @@
+- Added a separate epoch tracker for masp to decouple its logic from the rest of
+  the protocol. ([\#3318](https://github.com/anoma/namada/pull/3318))

crates/apps_lib/src/cli.rs

@@ -2989,6 +2989,7 @@ pub mod args {
     use namada::core::token::NATIVE_MAX_DECIMAL_PLACES;
     use namada::hash::Hash;
     use namada::ibc::core::host::types::identifiers::{ChannelId, PortId};
+    use namada::masp::MaspEpoch;
     use namada::tx::data::GasLimit;
     pub use namada_sdk::args::*;
     pub use namada_sdk::tx::{
@@ -3165,6 +3166,7 @@ pub mod args {
     pub const LIST_FIND_ADDRESSES_ONLY: ArgFlag = flag("addr");
     pub const LIST_FIND_KEYS_ONLY: ArgFlag = flag("keys");
     pub const LOCALHOST: ArgFlag = flag("localhost");
+    pub const MASP_EPOCH: ArgOpt<MaspEpoch> = arg_opt("masp-epoch");
     pub const MAX_COMMISSION_RATE_CHANGE: Arg<Dec> =
         arg("max-commission-rate-change");
     pub const MAX_ETH_GAS: ArgOpt<u64> = arg_opt("max_eth-gas");
@@ -5504,7 +5506,7 @@ pub mod args {
         fn parse(matches: &ArgMatches) -> Self {
             let query = Query::parse(matches);
             let token = TOKEN_OPT.parse(matches);
-            let epoch = EPOCH.parse(matches);
+            let epoch = MASP_EPOCH.parse(matches);
             Self {
                 query,
                 epoch,
@@ -5514,11 +5516,9 @@ pub mod args {
 
         fn def(app: App) -> App {
             app.add_args::<Query<CliTypes>>()
-                .arg(
-                    EPOCH.def().help(wrap!(
-                        "The epoch for which to query conversions."
-                    )),
-                )
+                .arg(MASP_EPOCH.def().help(wrap!(
+                    "The masp epoch for which to query conversions."
+                )))
                 .arg(TOKEN_OPT.def().help(wrap!(
                     "The token address for which to query conversions."
                 )))

crates/apps_lib/src/client/rpc.rs

@@ -32,6 +32,7 @@ use namada::ledger::parameters::{storage as param_storage, EpochDuration};
 use namada::ledger::pos::types::{CommissionPair, Slash};
 use namada::ledger::pos::PosParams;
 use namada::ledger::queries::RPC;
+use namada::masp::MaspEpoch;
 use namada::proof_of_stake::types::{
     ValidatorState, ValidatorStateInfo, WeightedValidator,
 };
@@ -72,6 +73,13 @@ pub async fn query_and_print_epoch(context: &impl Namada) -> Epoch {
     epoch
 }
 
+/// Query and print the masp epoch of the last committed block
+pub async fn query_and_print_masp_epoch(context: &impl Namada) -> MaspEpoch {
+    let epoch = rpc::query_masp_epoch(context.client()).await.unwrap();
+    display_line!(context.io(), "Last committed masp epoch: {}", epoch);
+    epoch
+}
+
 /// Query and print some information to help discern when the next epoch will
 /// begin.
 pub async fn query_and_print_next_epoch_info(context: &impl Namada) {
@@ -372,7 +380,7 @@ async fn query_shielded_balance(
     }
 
     // The epoch is required to identify timestamped tokens
-    let epoch = query_and_print_epoch(context).await;
+    let masp_epoch = query_and_print_masp_epoch(context).await;
 
     // Query the token alias in the wallet for pretty printing token balances
     let token_alias = lookup_token_alias(context, &token, &MASP).await;
@@ -400,7 +408,7 @@ async fn query_shielded_balance(
                 context.client(),
                 context.io(),
                 &viewing_key,
-                epoch,
+                masp_epoch,
             )
             .await
             .unwrap()
@@ -412,7 +420,7 @@ async fn query_shielded_balance(
     };
 
     let total_balance = shielded
-        .decode_combine_sum_to_epoch(context.client(), balance, epoch)
+        .decode_combine_sum_to_epoch(context.client(), balance, masp_epoch)
         .await
         .0
         .get(&token);
@@ -1765,7 +1773,7 @@ pub async fn query_conversion<C: namada::ledger::queries::Client + Sync>(
     Address,
     token::Denomination,
     MaspDigitPos,
-    Epoch,
+    MaspEpoch,
     I128Sum,
     MerklePath<Node>,
 )> {

crates/apps_lib/src/client/tx.rs

@@ -770,9 +770,9 @@ pub async fn submit_transfer(
                     // And it is rejected by a VP
                     matches!(resp.batch_result().get(&cmt_hash), Some(InnerTxResult::VpsRejected(_))) =>
                 {
-                    let submission_epoch = rpc::query_and_print_epoch(namada).await;
+                    let submission_masp_epoch = rpc::query_and_print_masp_epoch(namada).await;
                     // And its submission epoch doesn't match construction epoch
-                    if tx_epoch.unwrap() != submission_epoch {
+                    if tx_epoch.unwrap() != submission_masp_epoch {
                         // Then we probably straddled an epoch boundary. Let's retry...
                         edisplay_line!(namada.io(),
                             "MASP transaction rejected and this may be due to the \

crates/apps_lib/src/config/genesis.rs

@@ -309,6 +309,8 @@ pub struct Parameters {
     pub implicit_vp_sha256: [u8; 32],
     /// Expected number of epochs per year (read only)
     pub epochs_per_year: u64,
+    /// How many epochs it takes to transition to the next masp epoch
+    pub masp_epoch_multiplier: u64,
     /// Maximum amount of signatures per transaction
     pub max_signatures_per_transaction: u8,
     /// Fee unshielding gas limit

crates/apps_lib/src/config/genesis/chain.rs

@@ -302,6 +302,7 @@ impl Finalized {
             tx_allowlist,
             implicit_vp,
             epochs_per_year,
+            masp_epoch_multiplier,
             max_signatures_per_transaction,
             fee_unshielding_gas_limit,
             max_block_gas,
@@ -346,6 +347,7 @@ impl Finalized {
             tx_allowlist,
             implicit_vp_code_hash,
             epochs_per_year,
+            masp_epoch_multiplier,
             max_proposal_bytes,
             max_signatures_per_transaction,
             fee_unshielding_gas_limit,

crates/apps_lib/src/config/genesis/templates.rs

@@ -293,6 +293,8 @@ pub struct ChainParams<T: TemplateValidation> {
     pub implicit_vp: String,
     /// Expected number of epochs per year
     pub epochs_per_year: u64,
+    /// How many epochs it takes to transition to the next masp epoch
+    pub masp_epoch_multiplier: u64,
     /// Maximum number of signature per transaction
     pub max_signatures_per_transaction: u8,
     /// Max gas for block
@@ -319,6 +321,7 @@ impl ChainParams<Unvalidated> {
             tx_allowlist,
             implicit_vp,
             epochs_per_year,
+            masp_epoch_multiplier,
             max_signatures_per_transaction,
             max_block_gas,
             fee_unshielding_gas_limit,
@@ -364,6 +367,7 @@ impl ChainParams<Unvalidated> {
             tx_allowlist,
             implicit_vp,
             epochs_per_year,
+            masp_epoch_multiplier,
             max_signatures_per_transaction,
             max_block_gas,
             fee_unshielding_gas_limit,

crates/core/src/masp.rs

@@ -2,6 +2,7 @@
 
 use std::collections::BTreeMap;
 use std::fmt::Display;
+use std::num::ParseIntError;
 use std::str::FromStr;
 
 use borsh::{BorshDeserialize, BorshSchema, BorshSerialize};
@@ -23,6 +24,72 @@ use crate::string_encoding::{
 };
 use crate::token::{Denomination, MaspDigitPos};
 
+/// Wrapper type around `Epoch` for type safe operations involving the masp
+/// epoch
+#[derive(
+    BorshSerialize,
+    BorshDeserialize,
+    BorshDeserializer,
+    BorshSchema,
+    Clone,
+    Copy,
+    Debug,
+    PartialOrd,
+    Ord,
+    PartialEq,
+    Eq,
+    Hash,
+    Serialize,
+    Deserialize,
+)]
+pub struct MaspEpoch(Epoch);
+
+impl Display for MaspEpoch {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "{}", self.0)
+    }
+}
+
+impl FromStr for MaspEpoch {
+    type Err = ParseIntError;
+
+    fn from_str(s: &str) -> std::result::Result<Self, Self::Err> {
+        let raw: u64 = u64::from_str(s)?;
+        Ok(Self(Epoch(raw)))
+    }
+}
+
+impl MaspEpoch {
+    /// Converts and `Epoch` into a `MaspEpoch` based on the provided conversion
+    /// rate
+    pub fn try_from_epoch(
+        epoch: Epoch,
+        masp_epoch_multiplier: u64,
+    ) -> Result<Self, &'static str> {
+        Ok(Self(
+            epoch
+                .checked_div(masp_epoch_multiplier)
+                .ok_or("Masp epoch multiplier cannot be 0")?,
+        ))
+    }
+
+    /// Returns a 0 masp epoch
+    pub const fn zero() -> Self {
+        Self(Epoch(0))
+    }
+
+    /// Change to the previous masp epoch.
+    pub fn prev(&self) -> Option<Self> {
+        Some(Self(self.0.checked_sub(1)?))
+    }
+
+    /// Initialize a new masp epoch from the provided one
+    #[cfg(any(test, feature = "testing"))]
+    pub const fn new(epoch: u64) -> Self {
+        Self(Epoch(epoch))
+    }
+}
+
 /// The plain representation of a MASP aaset
 #[derive(
     BorshSerialize,
@@ -47,7 +114,7 @@ pub struct AssetData {
     /// The digit position covered by this asset type
     pub position: MaspDigitPos,
     /// The epoch of the asset type, if any
-    pub epoch: Option<Epoch>,
+    pub epoch: Option<MaspEpoch>,
 }
 
 impl AssetData {
@@ -66,7 +133,7 @@ impl AssetData {
 
     /// Give this pre-asset type the given epoch if already has an epoch. Return
     /// the replaced value.
-    pub fn redate(&mut self, to: Epoch) -> Option<Epoch> {
+    pub fn redate(&mut self, to: MaspEpoch) -> Option<MaspEpoch> {
         if self.epoch.is_some() {
             self.epoch.replace(to)
         } else {
@@ -85,7 +152,7 @@ pub fn encode_asset_type(
     token: Address,
     denom: Denomination,
     position: MaspDigitPos,
-    epoch: Option<Epoch>,
+    epoch: Option<MaspEpoch>,
 ) -> Result<AssetType, std::io::Error> {
     AssetData {
         token,

crates/core/src/parameters.rs

@@ -46,6 +46,9 @@ pub struct Parameters {
     pub implicit_vp_code_hash: Option<Hash>,
     /// Expected number of epochs per year (read only)
     pub epochs_per_year: u64,
+    /// The multiplier for masp epochs (it requires this amount of epochs to
+    /// transition to the next masp epoch)
+    pub masp_epoch_multiplier: u64,
     /// Maximum number of signature per transaction
     pub max_signatures_per_transaction: u8,
     /// Fee unshielding gas limit

crates/namada/src/ledger/native_vp/masp.rs

@@ -17,11 +17,10 @@ use namada_core::address::InternalAddress::Masp;
 use namada_core::arith::{checked, CheckedAdd, CheckedSub};
 use namada_core::booleans::BoolResultUnitExt;
 use namada_core::collections::HashSet;
-use namada_core::masp::encode_asset_type;
+use namada_core::masp::{encode_asset_type, MaspEpoch};
 use namada_core::storage::Key;
 use namada_gas::GasMetering;
 use namada_governance::storage::is_proposal_accepted;
-use namada_proof_of_stake::Epoch;
 use namada_sdk::masp::verify_shielded_tx;
 use namada_state::{ConversionState, OptionExt, ResultExt, StateRead};
 use namada_token::read_denom;
@@ -348,7 +347,17 @@ where
         tx_data: &BatchedTxRef<'_>,
         keys_changed: &BTreeSet<Key>,
     ) -> Result<()> {
-        let epoch = self.ctx.get_block_epoch()?;
+        let masp_epoch_multiplier =
+            namada_parameters::read_masp_epoch_multiplier_parameter(
+                self.ctx.state,
+            )?;
+        let masp_epoch = MaspEpoch::try_from_epoch(
+            self.ctx.get_block_epoch()?,
+            masp_epoch_multiplier,
+        )
+        .map_err(|msg| {
+            Error::NativeVpError(native_vp::Error::new_const(msg))
+        })?;
         let conversion_state = self.ctx.state.in_mem().get_conversion_state();
 
         // Get the Transaction object from the actions
@@ -402,7 +411,7 @@ where
                 .get(&masp_address_hash)
                 .unwrap_or(&ValueSum::zero()),
             &shielded_tx.sapling_value_balance(),
-            epoch,
+            masp_epoch,
             &changed_balances.tokens,
             conversion_state,
         )?;
@@ -428,7 +437,7 @@ where
             &shielded_tx,
             &mut changed_balances,
             &mut transparent_tx_pool,
-            epoch,
+            masp_epoch,
             conversion_state,
             &mut signers,
         )?;
@@ -567,7 +576,7 @@ fn validate_transparent_input<A: Authorization>(
     vin: &TxIn<A>,
     changed_balances: &mut ChangedBalances,
     transparent_tx_pool: &mut I128Sum,
-    epoch: Epoch,
+    epoch: MaspEpoch,
     conversion_state: &ConversionState,
     signers: &mut BTreeSet<TransparentAddress>,
 ) -> Result<()> {
@@ -654,7 +663,7 @@ fn validate_transparent_output(
     out: &TxOut,
     changed_balances: &mut ChangedBalances,
     transparent_tx_pool: &mut I128Sum,
-    epoch: Epoch,
+    epoch: MaspEpoch,
     conversion_state: &ConversionState,
 ) -> Result<()> {
     // Non-masp destinations subtract from transparent tx pool
@@ -720,7 +729,7 @@ fn validate_transparent_bundle(
     shielded_tx: &Transaction,
     changed_balances: &mut ChangedBalances,
     transparent_tx_pool: &mut I128Sum,
-    epoch: Epoch,
+    epoch: MaspEpoch,
     conversion_state: &ConversionState,
     signers: &mut BTreeSet<TransparentAddress>,
 ) -> Result<()> {
@@ -779,7 +788,7 @@ fn verify_sapling_balancing_value(
     pre: &ValueSum<Address, Amount>,
     post: &ValueSum<Address, Amount>,
     sapling_value_balance: &I128Sum,
-    target_epoch: Epoch,
+    target_epoch: MaspEpoch,
     tokens: &BTreeMap<AssetType, (Address, token::Denomination, MaspDigitPos)>,
     conversion_state: &ConversionState,
 ) -> Result<()> {

crates/node/src/bench_utils.rs

@@ -441,8 +441,12 @@ impl BenchShell {
         )
         .unwrap();
 
-        namada::token::conversion::update_allowed_conversions(&mut self.state)
+        if self.state.is_masp_new_epoch(true).unwrap() {
+            namada::token::conversion::update_allowed_conversions(
+                &mut self.state,
+            )
             .unwrap();
+        }
     }
 
     pub fn init_ibc_client_state(&mut self, addr_key: Key) -> ClientId {