From 826f3efce026b4eedf2f128215cddb036f89c634 Mon Sep 17 00:00:00 2001
From: Marco Granelli <marco@heliax.dev>
Date: Thu, 4 Apr 2024 23:44:56 +0200
Subject: [PATCH 1/7] Run masp verification in parallel

---
 crates/gas/src/lib.rs  |  34 +++-
 crates/sdk/src/masp.rs | 444 +++++++++++++++++++----------------------
 2 files changed, 228 insertions(+), 250 deletions(-)

diff --git a/crates/gas/src/lib.rs b/crates/gas/src/lib.rs
index c015542008..88c73b7d8d 100644
--- a/crates/gas/src/lib.rs
+++ b/crates/gas/src/lib.rs
@@ -84,14 +84,32 @@ pub const WASM_MEMORY_PAGE_GAS: u32 =
 pub const IBC_ACTION_VALIDATE_GAS: u64 = 1_472_023;
 /// The cost to execute an Ibc action
 pub const IBC_ACTION_EXECUTE_GAS: u64 = 3_678_745;
-/// The cost to verify a masp spend note
-pub const MASP_VERIFY_SPEND_GAS: u64 = 66_822_000;
-/// The cost to verify a masp convert note
-pub const MASP_VERIFY_CONVERT_GAS: u64 = 45_240_000;
-/// The cost to verify a masp output note
-pub const MASP_VERIFY_OUTPUT_GAS: u64 = 55_023_000;
-/// The cost to run the final masp verification
-pub const MASP_VERIFY_FINAL_GAS: u64 = 3_475_200;
+// FIXME: update costs
+/// The cost of masp sig verification
+pub const MASP_VERIFY_SIG_GAS: u64 = 1;
+/// The fixed cost of spend note verification
+pub const MASP_FIXED_SPEND_GAS: u64 = 1;
+/// The variable cost of spend note verification
+pub const MASP_VARIABLE_SPEND_GAS: u64 = 1;
+/// The fixed cost of convert note verification
+pub const MASP_FIXED_CONVERT_GAS: u64 = 1;
+/// The variable cost of convert note verification
+pub const MASP_VARIABLE_CONVERT_GAS: u64 = 1;
+/// The fixed cost of output note verification
+pub const MASP_FIXED_OUTPUT_GAS: u64 = 1;
+/// The variable cost of output note verification
+pub const MASP_VARIABLE_OUTPUT_GAS: u64 = 1;
+/// The cost to process a masp spend note in the bundle
+pub const MASP_SPEND_CHECK_GAS: u64 = 1;
+/// The cost to process a masp convert note in the bundle
+pub const MASP_CONVERT_CHECK_GAS: u64 = 1;
+/// The cost to process a masp output note in the bundle
+pub const MASP_OUTPUT_CHECK_GAS: u64 = 1;
+/// The cost to run the final masp check in the bundle
+pub const MASP_FINAL_CHECK_GAS: u64 = 1;
+/// Gas divider specific for the masp vp. Only allocates half of the cores to
+/// the masp vp since we can expect the other half to be busy with other vps
+pub const MASP_PARALLEL_GAS_DIVIDER: u64 = PARALLEL_GAS_DIVIDER / 2;
 
 /// Gas module result for functions that may fail
 pub type Result<T> = std::result::Result<T, Error>;
diff --git a/crates/sdk/src/masp.rs b/crates/sdk/src/masp.rs
index 7247d3e2c5..caa9a19d5a 100644
--- a/crates/sdk/src/masp.rs
+++ b/crates/sdk/src/masp.rs
@@ -17,14 +17,12 @@ use masp_primitives::consensus::MainNetwork as Network;
 use masp_primitives::consensus::TestNetwork as Network;
 use masp_primitives::convert::AllowedConversion;
 use masp_primitives::ff::PrimeField;
-use masp_primitives::group::GroupEncoding;
 use masp_primitives::memo::MemoBytes;
 use masp_primitives::merkle_tree::{
     CommitmentTree, IncrementalWitness, MerklePath,
 };
 use masp_primitives::sapling::keys::FullViewingKey;
 use masp_primitives::sapling::note_encryption::*;
-use masp_primitives::sapling::redjubjub::PublicKey;
 use masp_primitives::sapling::{
     Diversifier, Node, Note, Nullifier, ViewingKey,
 };
@@ -34,8 +32,7 @@ use masp_primitives::transaction::components::sapling::builder::{
 };
 use masp_primitives::transaction::components::transparent::builder::TransparentBuilder;
 use masp_primitives::transaction::components::{
-    ConvertDescription, I128Sum, OutputDescription, SpendDescription, TxOut,
-    U64Sum, ValueSum,
+    I128Sum, OutputDescription, TxOut, U64Sum, ValueSum,
 };
 use masp_primitives::transaction::fees::fixed::FeeRule;
 use masp_primitives::transaction::sighash::{signature_hash, SignableInput};
@@ -45,11 +42,10 @@ use masp_primitives::transaction::{
     TransparentAddress, Unauthorized,
 };
 use masp_primitives::zip32::{ExtendedFullViewingKey, ExtendedSpendingKey};
-use masp_proofs::bellman::groth16::PreparedVerifyingKey;
+use masp_proofs::bellman::groth16::VerifyingKey;
 use masp_proofs::bls12_381::Bls12;
 use masp_proofs::prover::LocalTxProver;
-#[cfg(not(feature = "testing"))]
-use masp_proofs::sapling::SaplingVerificationContext;
+use masp_proofs::sapling::BatchValidator;
 use namada_core::address::Address;
 use namada_core::collections::{HashMap, HashSet};
 use namada_core::dec::Dec;
@@ -70,7 +66,8 @@ use namada_migrations::*;
 use namada_state::StorageError;
 use namada_token::{self as token, Denomination, MaspDigitPos, Transfer};
 use namada_tx::{IndexedTx, Tx, TxCommitments};
-use rand_core::{CryptoRng, OsRng, RngCore};
+use rand::rngs::StdRng;
+use rand_core::{CryptoRng, OsRng, RngCore, SeedableRng};
 use ripemd::Digest as RipemdDigest;
 use sha2::Digest;
 use thiserror::Error;
@@ -148,11 +145,11 @@ struct ExtractedMaspTxs(Vec<(TxCommitments, Transaction)>);
 /// MASP verifying keys
 pub struct PVKs {
     /// spend verifying key
-    pub spend_vk: PreparedVerifyingKey<Bls12>,
+    pub spend_vk: VerifyingKey<Bls12>,
     /// convert verifying key
-    pub convert_vk: PreparedVerifyingKey<Bls12>,
+    pub convert_vk: VerifyingKey<Bls12>,
     /// output verifying key
-    pub output_vk: PreparedVerifyingKey<Bls12>,
+    pub output_vk: VerifyingKey<Bls12>,
 }
 
 lazy_static! {
@@ -186,9 +183,9 @@ lazy_static! {
             convert_path.as_path(),
         );
         PVKs {
-            spend_vk: params.spend_vk,
-            convert_vk: params.convert_vk,
-            output_vk: params.output_vk
+            spend_vk: params.spend_params.vk,
+            convert_vk: params.convert_params.vk,
+            output_vk: params.output_params.vk
         }
     };
 }
@@ -202,76 +199,6 @@ fn load_pvks() -> &'static PVKs {
     &VERIFIYING_KEYS
 }
 
-/// check_spend wrapper
-pub fn check_spend(
-    spend: &SpendDescription<<Authorized as Authorization>::SaplingAuth>,
-    sighash: &[u8; 32],
-    #[cfg(not(feature = "testing"))] ctx: &mut SaplingVerificationContext,
-    #[cfg(feature = "testing")]
-    ctx: &mut testing::MockSaplingVerificationContext,
-    parameters: &PreparedVerifyingKey<Bls12>,
-) -> bool {
-    let zkproof =
-        masp_proofs::bellman::groth16::Proof::read(spend.zkproof.as_slice());
-    let zkproof = match zkproof {
-        Ok(zkproof) => zkproof,
-        _ => return false,
-    };
-
-    ctx.check_spend(
-        spend.cv,
-        spend.anchor,
-        &spend.nullifier.0,
-        PublicKey(spend.rk.0),
-        sighash,
-        spend.spend_auth_sig,
-        zkproof,
-        parameters,
-    )
-}
-
-/// check_output wrapper
-pub fn check_output(
-    output: &OutputDescription<<<Authorized as Authorization>::SaplingAuth as masp_primitives::transaction::components::sapling::Authorization>::Proof>,
-    #[cfg(not(feature = "testing"))] ctx: &mut SaplingVerificationContext,
-    #[cfg(feature = "testing")]
-    ctx: &mut testing::MockSaplingVerificationContext,
-    parameters: &PreparedVerifyingKey<Bls12>,
-) -> bool {
-    let zkproof =
-        masp_proofs::bellman::groth16::Proof::read(output.zkproof.as_slice());
-    let zkproof = match zkproof {
-        Ok(zkproof) => zkproof,
-        _ => return false,
-    };
-    let epk =
-        masp_proofs::jubjub::ExtendedPoint::from_bytes(&output.ephemeral_key.0);
-    let epk = match epk.into() {
-        Some(p) => p,
-        None => return false,
-    };
-
-    ctx.check_output(output.cv, output.cmu, epk, zkproof, parameters)
-}
-
-/// check convert wrapper
-pub fn check_convert(
-    convert: &ConvertDescription<<<Authorized as Authorization>::SaplingAuth as masp_primitives::transaction::components::sapling::Authorization>::Proof>,
-    #[cfg(not(feature = "testing"))] ctx: &mut SaplingVerificationContext,
-    #[cfg(feature = "testing")]
-    ctx: &mut testing::MockSaplingVerificationContext,
-    parameters: &PreparedVerifyingKey<Bls12>,
-) -> bool {
-    let zkproof =
-        masp_proofs::bellman::groth16::Proof::read(convert.zkproof.as_slice());
-    let zkproof = match zkproof {
-        Ok(zkproof) => zkproof,
-        _ => return false,
-    };
-
-    ctx.check_convert(convert.cv, convert.anchor, zkproof, parameters)
-}
-
 /// Represents an authorization where the Sapling bundle is authorized and the
 /// transparent bundle is unauthorized.
 pub struct PartialAuthorized;
@@ -317,12 +244,12 @@ pub fn partial_deauthorize(
 /// Verify a shielded transaction.
 pub fn verify_shielded_tx<F>(
     transaction: &Transaction,
-    mut consume_verify_gas: F,
+    consume_verify_gas: F,
 ) -> Result<(), StorageError>
 where
-    F: FnMut(u64) -> std::result::Result<(), StorageError>,
+    F: Fn(u64) -> std::result::Result<(), StorageError>,
 {
-    tracing::info!("entered verify_shielded_tx()");
+    tracing::debug!("entered verify_shielded_tx()");
 
     let sapling_bundle = if let Some(bundle) = transaction.sapling_bundle() {
         bundle
@@ -347,8 +274,7 @@ where
     // for now we need to continue to compute it here.
     let sighash =
         signature_hash(&unauth_tx_data, &SignableInput::Shielded, &txid_parts);
-
-    tracing::info!("sighash computed");
+    tracing::debug!("sighash computed");
 
     let PVKs {
         spend_vk,
@@ -357,49 +283,154 @@ where
     } = load_pvks();
 
     #[cfg(not(feature = "testing"))]
-    let mut ctx = SaplingVerificationContext::new(true);
+    let mut ctx = BatchValidator::new();
     #[cfg(feature = "testing")]
-    let mut ctx = testing::MockSaplingVerificationContext::new(true);
-    for spend in &sapling_bundle.shielded_spends {
-        consume_verify_gas(namada_gas::MASP_VERIFY_SPEND_GAS)?;
-        if !check_spend(spend, sighash.as_ref(), &mut ctx, spend_vk) {
-            return Err(StorageError::SimpleMessage("Invalid shielded spend"));
-        }
-    }
-    for convert in &sapling_bundle.shielded_converts {
-        consume_verify_gas(namada_gas::MASP_VERIFY_CONVERT_GAS)?;
-        if !check_convert(convert, &mut ctx, convert_vk) {
-            return Err(StorageError::SimpleMessage(
-                "Invalid shielded conversion",
-            ));
-        }
-    }
-    for output in &sapling_bundle.shielded_outputs {
-        consume_verify_gas(namada_gas::MASP_VERIFY_OUTPUT_GAS)?;
-        if !check_output(output, &mut ctx, output_vk) {
-            return Err(StorageError::SimpleMessage("Invalid shielded output"));
-        }
+    let mut ctx = testing::MockBatchValidator::default();
+
+    // Charge gas before check bundle
+    let spends_len = sapling_bundle.shielded_spends.len() as u64;
+    let converts_len = sapling_bundle.shielded_converts.len() as u64;
+    let outputs_len = sapling_bundle.shielded_outputs.len() as u64;
+    charge_masp_check_bundle_gas(
+        spends_len,
+        converts_len,
+        outputs_len,
+        &consume_verify_gas,
+    )?;
+
+    if !ctx.check_bundle(sapling_bundle.to_owned(), sighash.as_ref().to_owned())
+    {
+        tracing::debug!("failed check bundle");
+        return Err(StorageError::SimpleMessage("Invalid sapling bundle"));
+    }
+    tracing::debug!("passed check bundle");
+
+    // Charge gas before final validation
+    charge_masp_validate_gas(
+        spends_len,
+        converts_len,
+        outputs_len,
+        consume_verify_gas,
+    )?;
+    if !ctx.validate(spend_vk, convert_vk, output_vk, OsRng) {
+        return Err(StorageError::SimpleMessage(
+            "Invalid proofs or signatures",
+        ));
     }
+    Ok(())
+}
 
-    tracing::info!("passed spend/output verification");
+// Charge gas for the check_bundle operation which does not leverage concurrency
+fn charge_masp_check_bundle_gas<F>(
+    spends_len: u64,
+    converts_len: u64,
+    outputs_len: u64,
+    consume_verify_gas: F,
+) -> Result<(), namada_state::StorageError>
+where
+    F: Fn(u64) -> std::result::Result<(), namada_state::StorageError>,
+{
+    consume_verify_gas(
+        spends_len
+            .checked_mul(namada_gas::MASP_SPEND_CHECK_GAS)
+            .ok_or(namada_state::StorageError::SimpleMessage(
+                "Overflow in gas calculation",
+            ))?,
+    )?;
+
+    consume_verify_gas(
+        converts_len
+            .checked_mul(namada_gas::MASP_CONVERT_CHECK_GAS)
+            .ok_or(namada_state::StorageError::SimpleMessage(
+                "Overflow in gas calculation",
+            ))?,
+    )?;
+
+    consume_verify_gas(
+        outputs_len
+            .checked_mul(namada_gas::MASP_OUTPUT_CHECK_GAS)
+            .ok_or(namada_state::StorageError::SimpleMessage(
+                "Overflow in gas calculation",
+            ))?,
+    )?;
+
+    Ok(())
+}
 
-    let assets_and_values: I128Sum = sapling_bundle.value_balance.clone();
+// Charge gas for the final validation, taking advtange of concurrency for
+// proofs verification but not for signature (because of the fallback)
+fn charge_masp_validate_gas<F>(
+    spends_len: u64,
+    converts_len: u64,
+    outputs_len: u64,
+    consume_verify_gas: F,
+) -> Result<(), namada_state::StorageError>
+where
+    F: Fn(u64) -> std::result::Result<(), namada_state::StorageError>,
+{
+    // Signatures pay the entire cost (no discount because of the possible
+    // fallback)
+    consume_verify_gas(
+        spends_len
+            // Add one for the binding signature
+            .checked_add(1)
+            .ok_or(namada_state::StorageError::SimpleMessage(
+                "Overflow in gas calculation",
+            ))?
+            .checked_mul(namada_gas::MASP_VERIFY_SIG_GAS)
+            .ok_or(namada_state::StorageError::SimpleMessage(
+                "Overflow in gas calculation",
+            ))?,
+    )?;
+
+    // If at least one note is present charge the fixed costs. Then charge the
+    // variable cost for every note, amortized on the fixed expected number of
+    // cores
+    if spends_len != 0 {
+        consume_verify_gas(namada_gas::MASP_FIXED_SPEND_GAS)?;
+        consume_verify_gas(
+            namada_gas::MASP_VARIABLE_SPEND_GAS
+                .checked_mul(spends_len)
+                .ok_or(namada_state::StorageError::SimpleMessage(
+                    "Overflow in gas calculation",
+                ))?
+                .checked_div(namada_gas::MASP_PARALLEL_GAS_DIVIDER)
+                .ok_or(namada_state::StorageError::SimpleMessage(
+                    "Overflow in gas calculation",
+                ))?,
+        )?;
+    }
 
-    tracing::info!(
-        "accumulated {} assets/values",
-        assets_and_values.components().len()
-    );
+    if converts_len != 0 {
+        consume_verify_gas(namada_gas::MASP_FIXED_CONVERT_GAS)?;
+        consume_verify_gas(
+            namada_gas::MASP_VARIABLE_CONVERT_GAS
+                .checked_mul(converts_len)
+                .ok_or(namada_state::StorageError::SimpleMessage(
+                    "Overflow in gas calculation",
+                ))?
+                .checked_div(namada_gas::MASP_PARALLEL_GAS_DIVIDER)
+                .ok_or(namada_state::StorageError::SimpleMessage(
+                    "Overflow in gas calculation",
+                ))?,
+        )?;
+    }
 
-    consume_verify_gas(namada_gas::MASP_VERIFY_FINAL_GAS)?;
-    let result = ctx.final_check(
-        assets_and_values,
-        sighash.as_ref(),
-        sapling_bundle.authorization.binding_sig,
-    );
-    tracing::info!("final check result {result}");
-    if !result {
-        return Err(StorageError::SimpleMessage("MASP final check failed"));
+    if outputs_len != 0 {
+        consume_verify_gas(namada_gas::MASP_FIXED_OUTPUT_GAS)?;
+        consume_verify_gas(
+            namada_gas::MASP_VARIABLE_OUTPUT_GAS
+                .checked_mul(outputs_len)
+                .ok_or(namada_state::StorageError::SimpleMessage(
+                    "Overflow in gas calculation",
+                ))?
+                .checked_div(namada_gas::MASP_PARALLEL_GAS_DIVIDER)
+                .ok_or(namada_state::StorageError::SimpleMessage(
+                    "Overflow in gas calculation",
+                ))?,
+        )?;
     }
+
     Ok(())
 }
 
@@ -1559,9 +1590,6 @@ impl<U: ShieldedUtils + MaybeSend + MaybeSync> ShieldedContext<U> {
         // No shielded components are needed when neither source nor destination
         // are shielded
 
-        use rand::rngs::StdRng;
-        use rand_core::SeedableRng;
-
         let spending_key = source.spending_key();
         let payment_address = target.payment_address();
         // No shielded components are needed when neither source nor
@@ -2216,11 +2244,13 @@ pub mod testing {
     use bls12_381::{G1Affine, G2Affine};
     use masp_primitives::consensus::testing::arb_height;
     use masp_primitives::constants::SPENDING_KEY_GENERATOR;
+    use masp_primitives::group::GroupEncoding;
     use masp_primitives::sapling::prover::TxProver;
-    use masp_primitives::sapling::redjubjub::Signature;
+    use masp_primitives::sapling::redjubjub::{PublicKey, Signature};
     use masp_primitives::sapling::{ProofGenerationKey, Rseed};
+    use masp_primitives::transaction::components::sapling::Bundle;
     use masp_primitives::transaction::components::GROTH_PROOF_SIZE;
-    use masp_proofs::bellman::groth16::Proof;
+    use masp_proofs::bellman::groth16::{self, Proof};
     use proptest::prelude::*;
     use proptest::sample::SizeRange;
     use proptest::test_runner::TestRng;
@@ -2234,129 +2264,59 @@ pub mod testing {
     use crate::masp_primitives::sapling::keys::OutgoingViewingKey;
     use crate::masp_primitives::sapling::redjubjub::PrivateKey;
     use crate::masp_primitives::transaction::components::transparent::testing::arb_transparent_address;
-    use crate::masp_proofs::sapling::SaplingVerificationContextInner;
     use crate::storage::testing::arb_epoch;
     use crate::token::testing::arb_denomination;
 
-    /// A context object for verifying the Sapling components of a single Zcash
-    /// transaction. Same as SaplingVerificationContext, but always assumes the
-    /// proofs to be valid.
-    pub struct MockSaplingVerificationContext {
-        inner: SaplingVerificationContextInner,
-        zip216_enabled: bool,
+    /// A context object for verifying the Sapling components of MASP
+    /// transactions. Same as BatchValidator, but always assumes the
+    /// proofs and signatures to be valid.
+    pub struct MockBatchValidator {
+        inner: BatchValidator,
     }
 
-    impl MockSaplingVerificationContext {
-        /// Construct a new context to be used with a single transaction.
-        pub fn new(zip216_enabled: bool) -> Self {
-            MockSaplingVerificationContext {
-                inner: SaplingVerificationContextInner::new(),
-                zip216_enabled,
+    impl Default for MockBatchValidator {
+        fn default() -> Self {
+            MockBatchValidator {
+                inner: BatchValidator::new(),
             }
         }
+    }
 
-        /// Perform consensus checks on a Sapling SpendDescription, while
-        /// accumulating its value commitment inside the context for later use.
-        #[allow(clippy::too_many_arguments)]
-        pub fn check_spend(
-            &mut self,
-            cv: jubjub::ExtendedPoint,
-            anchor: bls12_381::Scalar,
-            nullifier: &[u8; 32],
-            rk: PublicKey,
-            sighash_value: &[u8; 32],
-            spend_auth_sig: Signature,
-            zkproof: Proof<Bls12>,
-            _verifying_key: &PreparedVerifyingKey<Bls12>,
-        ) -> bool {
-            let zip216_enabled = true;
-            self.inner.check_spend(
-                cv,
-                anchor,
-                nullifier,
-                rk,
-                sighash_value,
-                spend_auth_sig,
-                zkproof,
-                &mut (),
-                |_, rk, msg, spend_auth_sig| {
-                    rk.verify_with_zip216(
-                        &msg,
-                        &spend_auth_sig,
-                        SPENDING_KEY_GENERATOR,
-                        zip216_enabled,
-                    )
-                },
-                |_, _proof, _public_inputs| true,
-            )
-        }
-
-        /// Perform consensus checks on a Sapling SpendDescription, while
-        /// accumulating its value commitment inside the context for later use.
-        #[allow(clippy::too_many_arguments)]
-        pub fn check_convert(
+    impl MockBatchValidator {
+        /// Checks the bundle against Sapling-specific consensus rules, and adds
+        /// its proof and signatures to the validator.
+        ///
+        /// Returns `false` if the bundle doesn't satisfy all of the consensus
+        /// rules. This `BatchValidator` can continue to be used
+        /// regardless, but some or all of the proofs and signatures
+        /// from this bundle may have already been added to the batch even if
+        /// it fails other consensus rules.
+        pub fn check_bundle(
             &mut self,
-            cv: jubjub::ExtendedPoint,
-            anchor: bls12_381::Scalar,
-            zkproof: Proof<Bls12>,
-            _verifying_key: &PreparedVerifyingKey<Bls12>,
+            bundle: Bundle<
+                masp_primitives::transaction::components::sapling::Authorized,
+            >,
+            sighash: [u8; 32],
         ) -> bool {
-            self.inner.check_convert(
-                cv,
-                anchor,
-                zkproof,
-                &mut (),
-                |_, _proof, _public_inputs| true,
-            )
-        }
-
-        /// Perform consensus checks on a Sapling OutputDescription, while
-        /// accumulating its value commitment inside the context for later use.
-        pub fn check_output(
-            &mut self,
-            cv: jubjub::ExtendedPoint,
-            cmu: bls12_381::Scalar,
-            epk: jubjub::ExtendedPoint,
-            zkproof: Proof<Bls12>,
-            _verifying_key: &PreparedVerifyingKey<Bls12>,
+            self.inner.check_bundle(bundle, sighash)
+        }
+
+        /// Batch-validates the accumulated bundles.
+        ///
+        /// Returns `true` if every proof and signature in every bundle added to
+        /// the batch validator is valid, or `false` if one or more are
+        /// invalid. No attempt is made to figure out which of the
+        /// accumulated bundles might be invalid; if that information is
+        /// desired, construct separate [`BatchValidator`]s for sub-batches of
+        /// the bundles.
+        pub fn validate<R: RngCore + CryptoRng>(
+            self,
+            _spend_vk: &groth16::VerifyingKey<Bls12>,
+            _convert_vk: &groth16::VerifyingKey<Bls12>,
+            _output_vk: &groth16::VerifyingKey<Bls12>,
+            mut _rng: R,
         ) -> bool {
-            self.inner.check_output(
-                cv,
-                cmu,
-                epk,
-                zkproof,
-                |_proof, _public_inputs| true,
-            )
-        }
-
-        /// Perform consensus checks on the valueBalance and bindingSig parts of
-        /// a Sapling transaction. All SpendDescriptions and
-        /// OutputDescriptions must have been checked before calling
-        /// this function.
-        pub fn final_check(
-            &self,
-            value_balance: I128Sum,
-            sighash_value: &[u8; 32],
-            binding_sig: Signature,
-        ) -> bool {
-            self.inner.final_check(
-                value_balance,
-                sighash_value,
-                binding_sig,
-                |bvk, msg, binding_sig| {
-                    // Compute the signature's message for bvk/binding_sig
-                    let mut data_to_be_signed = [0u8; 64];
-                    data_to_be_signed[0..32].copy_from_slice(&bvk.0.to_bytes());
-                    data_to_be_signed[32..64].copy_from_slice(msg);
-
-                    bvk.verify_with_zip216(
-                        &data_to_be_signed,
-                        &binding_sig,
-                        VALUE_COMMITMENT_RANDOMNESS_GENERATOR,
-                        self.zip216_enabled,
-                    )
-                },
-            )
+            true
         }
     }
 

From 6fe54a4dff6d1889021ff3798bc3e47ed85214bc Mon Sep 17 00:00:00 2001
From: Marco Granelli <marco@heliax.dev>
Date: Thu, 11 Apr 2024 20:02:28 +0200
Subject: [PATCH 2/7] Updates benchmarks for parallel masp verification

---
 Cargo.lock                        |   1 +
 crates/benches/Cargo.toml         |   1 +
 crates/benches/native_vps.rs      | 380 +++++++++++++++++++++++++-----
 crates/benches/process_wrapper.rs |  20 +-
 4 files changed, 335 insertions(+), 67 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 0a9d914fe1..3b3d63a1ee 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -4557,6 +4557,7 @@ dependencies = [
  "criterion",
  "lazy_static",
  "masp_primitives",
+ "masp_proofs",
  "namada",
  "namada_apps_lib",
  "namada_node",
diff --git a/crates/benches/Cargo.toml b/crates/benches/Cargo.toml
index ca308f9100..01d25fa777 100644
--- a/crates/benches/Cargo.toml
+++ b/crates/benches/Cargo.toml
@@ -46,6 +46,7 @@ namada = { path = "../namada", features = ["rand", "benches"] }
 namada_apps_lib = { path = "../apps_lib" }
 namada_node = { path = "../node", features = ["benches"] }
 masp_primitives.workspace = true
+masp_proofs = { workspace = true, features = ["multicore"] }
 borsh.workspace = true
 borsh-ext.workspace = true
 criterion = { version = "0.5", features = ["html_reports"] }
diff --git a/crates/benches/native_vps.rs b/crates/benches/native_vps.rs
index 0ad8a637b5..965a3dfe65 100644
--- a/crates/benches/native_vps.rs
+++ b/crates/benches/native_vps.rs
@@ -5,9 +5,13 @@ use std::rc::Rc;
 use std::str::FromStr;
 
 use criterion::{criterion_group, criterion_main, BatchSize, Criterion};
+use masp_primitives::sapling::redjubjub::PublicKey;
 use masp_primitives::sapling::Node;
 use masp_primitives::transaction::sighash::{signature_hash, SignableInput};
 use masp_primitives::transaction::txid::TxIdDigester;
+use masp_primitives::transaction::TransactionData;
+use masp_proofs::group::GroupEncoding;
+use masp_proofs::sapling::BatchValidator;
 use namada::core::address::{self, Address, InternalAddress};
 use namada::core::collections::HashMap;
 use namada::core::eth_bridge_pool::{GasFee, PendingTransfer};
@@ -46,13 +50,10 @@ use namada::ledger::pgf::PgfVp;
 use namada::ledger::pos::PosVP;
 use namada::proof_of_stake;
 use namada::proof_of_stake::KeySeg;
-use namada::sdk::masp::{
-    check_convert, check_output, check_spend, partial_deauthorize,
-    preload_verifying_keys, PVKs,
-};
+use namada::sdk::masp::{partial_deauthorize, preload_verifying_keys, PVKs};
 use namada::sdk::masp_primitives::merkle_tree::CommitmentTree;
 use namada::sdk::masp_primitives::transaction::Transaction;
-use namada::sdk::masp_proofs::sapling::SaplingVerificationContext;
+use namada::sdk::masp_proofs::sapling::SaplingVerificationContextInner;
 use namada::state::{Epoch, StorageRead, StorageWrite, TxIndex};
 use namada::token::{Amount, Transfer};
 use namada::tx::{BatchedTx, Code, Section, Tx};
@@ -63,6 +64,7 @@ use namada_node::bench_utils::{
     TX_BRIDGE_POOL_WASM, TX_IBC_WASM, TX_INIT_PROPOSAL_WASM, TX_RESIGN_STEWARD,
     TX_TRANSFER_WASM, TX_UPDATE_STEWARD_COMMISSION, TX_VOTE_PROPOSAL_WASM,
 };
+use rand_core::OsRng;
 
 fn governance(c: &mut Criterion) {
     let mut group = c.benchmark_group("vp_governance");
@@ -455,8 +457,7 @@ fn ibc(c: &mut Criterion) {
                         ibc.ctx.keys_changed,
                         ibc.ctx.verifiers,
                     )
-                    .is_ok()
-                )
+                    .is_ok())
             })
         });
     }
@@ -633,8 +634,7 @@ fn masp(c: &mut Criterion) {
                         masp.ctx.keys_changed,
                         masp.ctx.verifiers,
                     )
-                    .is_ok()
-                );
+                    .is_ok());
             })
         });
     }
@@ -642,11 +642,11 @@ fn masp(c: &mut Criterion) {
     group.finish();
 }
 
+// Instead of benchmarking BatchValidator::check_bundle we benchmark the 4
+// functions that are called internally for better resolution
 fn masp_check_spend(c: &mut Criterion) {
-    let spend_vk = &preload_verifying_keys().spend_vk;
-
     c.bench_function("vp_masp_check_spend", |b| {
-        b.iter_batched_ref(
+        b.iter_batched(
             || {
                 let (_, _verifiers_from_tx, signed_tx) =
                     setup_storage_for_masp_verification("shielded");
@@ -670,7 +670,7 @@ fn masp_check_spend(c: &mut Criterion) {
                     .first()
                     .unwrap()
                     .to_owned();
-                let ctx = SaplingVerificationContext::new(true);
+                let ctx = SaplingVerificationContextInner::new();
                 let tx_data = transaction.deref();
                 // Partially deauthorize the transparent bundle
                 let unauth_tx_data = partial_deauthorize(tx_data).unwrap();
@@ -680,11 +680,28 @@ fn masp_check_spend(c: &mut Criterion) {
                     &SignableInput::Shielded,
                     &txid_parts,
                 );
+                let zkproof = masp_proofs::bellman::groth16::Proof::read(
+                    spend.zkproof.as_slice(),
+                )
+                .unwrap();
 
-                (ctx, spend, sighash)
+                (ctx, spend, sighash, zkproof)
             },
-            |(ctx, spend, sighash)| {
-                assert!(check_spend(spend, sighash.as_ref(), ctx, spend_vk));
+            |(mut ctx, spend, sighash, zkproof)| {
+                assert!(ctx.check_spend(
+                    spend.cv,
+                    spend.anchor,
+                    &spend.nullifier.0,
+                    PublicKey(spend.rk.0),
+                    sighash.as_ref(),
+                    spend.spend_auth_sig,
+                    zkproof,
+                    &mut (),
+                    // We do sig and proofs verification in parallel, so just
+                    // use dummy verifiers here
+                    |_, _, _, _| true,
+                    |_, _, _| true
+                ));
             },
             BatchSize::SmallInput,
         )
@@ -692,10 +709,8 @@ fn masp_check_spend(c: &mut Criterion) {
 }
 
 fn masp_check_convert(c: &mut Criterion) {
-    let convert_vk = &preload_verifying_keys().convert_vk;
-
     c.bench_function("vp_masp_check_convert", |b| {
-        b.iter_batched_ref(
+        b.iter_batched(
             || {
                 let (_, _verifiers_from_tx, signed_tx) =
                     setup_storage_for_masp_verification("shielded");
@@ -719,12 +734,24 @@ fn masp_check_convert(c: &mut Criterion) {
                     .first()
                     .unwrap()
                     .to_owned();
-                let ctx = SaplingVerificationContext::new(true);
+                let ctx = SaplingVerificationContextInner::new();
+                let zkproof = masp_proofs::bellman::groth16::Proof::read(
+                    convert.zkproof.as_slice(),
+                )
+                .unwrap();
 
-                (ctx, convert)
+                (ctx, convert, zkproof)
             },
-            |(ctx, convert)| {
-                assert!(check_convert(convert, ctx, convert_vk));
+            |(mut ctx, convert, zkproof)| {
+                assert!(ctx.check_convert(
+                    convert.cv,
+                    convert.anchor,
+                    zkproof,
+                    &mut (),
+                    // We do proofs verification in parallel, so just use dummy
+                    // verifier here
+                    |_, _, _| true,
+                ));
             },
             BatchSize::SmallInput,
         )
@@ -732,10 +759,8 @@ fn masp_check_convert(c: &mut Criterion) {
 }
 
 fn masp_check_output(c: &mut Criterion) {
-    let output_vk = &preload_verifying_keys().output_vk;
-
     c.bench_function("masp_vp_check_output", |b| {
-        b.iter_batched_ref(
+        b.iter_batched(
             || {
                 let (_, _verifiers_from_tx, signed_tx) =
                     setup_storage_for_masp_verification("shielded");
@@ -759,12 +784,28 @@ fn masp_check_output(c: &mut Criterion) {
                     .first()
                     .unwrap()
                     .to_owned();
-                let ctx = SaplingVerificationContext::new(true);
+                let ctx = SaplingVerificationContextInner::new();
+                let zkproof = masp_proofs::bellman::groth16::Proof::read(
+                    output.zkproof.as_slice(),
+                )
+                .unwrap();
+                let epk = masp_proofs::jubjub::ExtendedPoint::from_bytes(
+                    &output.ephemeral_key.0,
+                )
+                .unwrap();
 
-                (ctx, output)
+                (ctx, output, epk, zkproof)
             },
-            |(ctx, output)| {
-                assert!(check_output(output, ctx, output_vk));
+            |(mut ctx, output, epk, zkproof)| {
+                assert!(ctx.check_output(
+                    output.cv,
+                    output.cmu,
+                    epk,
+                    zkproof,
+                    // We do proofs verification in parallel, so just use dummy
+                    // verifier here
+                    |_, _| true
+                ));
             },
             BatchSize::SmallInput,
         )
@@ -772,12 +813,6 @@ fn masp_check_output(c: &mut Criterion) {
 }
 
 fn masp_final_check(c: &mut Criterion) {
-    let PVKs {
-        spend_vk,
-        convert_vk,
-        output_vk,
-    } = preload_verifying_keys();
-
     let (_, _verifiers_from_tx, signed_tx) =
         setup_storage_for_masp_verification("shielded");
 
@@ -794,41 +829,274 @@ fn masp_final_check(c: &mut Criterion) {
         .unwrap()
         .to_owned();
     let sapling_bundle = transaction.sapling_bundle().unwrap();
-    let mut ctx = SaplingVerificationContext::new(true);
     // Partially deauthorize the transparent bundle
     let unauth_tx_data = partial_deauthorize(transaction.deref()).unwrap();
     let txid_parts = unauth_tx_data.digest(TxIdDigester);
     let sighash =
         signature_hash(&unauth_tx_data, &SignableInput::Shielded, &txid_parts);
+    let mut ctx = SaplingVerificationContextInner::new();
 
     // Check spends, converts and outputs before the final check
     assert!(sapling_bundle.shielded_spends.iter().all(|spend| {
-        check_spend(spend, sighash.as_ref(), &mut ctx, spend_vk)
+        let zkproof = masp_proofs::bellman::groth16::Proof::read(
+            spend.zkproof.as_slice(),
+        )
+        .unwrap();
+
+        ctx.check_spend(
+            spend.cv,
+            spend.anchor,
+            &spend.nullifier.0,
+            PublicKey(spend.rk.0),
+            sighash.as_ref(),
+            spend.spend_auth_sig,
+            zkproof,
+            &mut (),
+            |_, _, _, _| true,
+            |_, _, _| true,
+        )
+    }));
+    assert!(sapling_bundle.shielded_converts.iter().all(|convert| {
+        let zkproof = masp_proofs::bellman::groth16::Proof::read(
+            convert.zkproof.as_slice(),
+        )
+        .unwrap();
+        ctx.check_convert(
+            convert.cv,
+            convert.anchor,
+            zkproof,
+            &mut (),
+            |_, _, _| true,
+        )
+    }));
+    assert!(sapling_bundle.shielded_outputs.iter().all(|output| {
+        let zkproof = masp_proofs::bellman::groth16::Proof::read(
+            output.zkproof.as_slice(),
+        )
+        .unwrap();
+        let epk = masp_proofs::jubjub::ExtendedPoint::from_bytes(
+            &output.ephemeral_key.0,
+        )
+        .unwrap();
+        ctx.check_output(
+            output.cv,
+            output.cmu,
+            epk.to_owned(),
+            zkproof,
+            |_, _| true,
+        )
     }));
-    assert!(
-        sapling_bundle
-            .shielded_converts
-            .iter()
-            .all(|convert| check_convert(convert, &mut ctx, convert_vk))
-    );
-    assert!(
-        sapling_bundle
-            .shielded_outputs
-            .iter()
-            .all(|output| check_output(output, &mut ctx, output_vk))
-    );
 
     c.bench_function("vp_masp_final_check", |b| {
         b.iter(|| {
             assert!(ctx.final_check(
                 sapling_bundle.value_balance.clone(),
                 sighash.as_ref(),
-                sapling_bundle.authorization.binding_sig
+                sapling_bundle.authorization.binding_sig,
+                // We do sig verification in parallel, so just use dummy
+                // verifier here
+                |_, _, _| true
             ))
         })
     });
 }
 
+#[derive(Debug)]
+enum RequestedItem {
+    Signature,
+    Spend,
+    Convert,
+    Output,
+}
+
+// Removes the unneeded notes from the generated transaction and replicates the
+// remaining ones if needed
+fn customize_masp_tx_data(
+    multi: bool,
+    request: &RequestedItem,
+) -> Option<(
+    TransactionData<masp_primitives::transaction::Authorized>,
+    Transaction,
+)> {
+    let (_, tx) = setup_storage_for_masp_verification("unshielding");
+    let transaction = tx
+        .sections
+        .into_iter()
+        .filter_map(|section| match section {
+            Section::MaspTx(transaction) => Some(transaction),
+            _ => None,
+        })
+        .collect::<Vec<Transaction>>()
+        .first()
+        .unwrap()
+        .to_owned();
+    let mut sapling_bundle = transaction.sapling_bundle().unwrap().to_owned();
+
+    match request {
+        RequestedItem::Signature => {
+            if multi {
+                // No multisig benchmark
+                return None;
+            }
+            // ensure we only have one signature to verify (the binding one) and
+            // no proofs
+            sapling_bundle.shielded_spends.clear();
+            sapling_bundle.shielded_converts.clear();
+            sapling_bundle.shielded_outputs.clear();
+            assert_eq!(sapling_bundle.shielded_spends.len(), 0);
+            assert_eq!(sapling_bundle.shielded_converts.len(), 0);
+            assert_eq!(sapling_bundle.shielded_outputs.len(), 0);
+        }
+        RequestedItem::Spend => {
+            if multi {
+                // ensure we only have two spend proofs
+                sapling_bundle.shielded_spends = [
+                    sapling_bundle.shielded_spends.clone(),
+                    sapling_bundle.shielded_spends,
+                ]
+                .concat();
+                sapling_bundle.shielded_outputs.clear();
+                sapling_bundle.shielded_converts.clear();
+                assert_eq!(sapling_bundle.shielded_spends.len(), 2);
+            } else {
+                // ensure we only have one spend proof
+                sapling_bundle.shielded_outputs.clear();
+                sapling_bundle.shielded_converts.clear();
+                assert_eq!(sapling_bundle.shielded_spends.len(), 1);
+            }
+            assert_eq!(sapling_bundle.shielded_converts.len(), 0);
+            assert_eq!(sapling_bundle.shielded_outputs.len(), 0);
+        }
+        RequestedItem::Convert => {
+            if multi {
+                // ensure we only have two convert proofs
+                sapling_bundle.shielded_converts = [
+                    sapling_bundle.shielded_converts.clone(),
+                    sapling_bundle.shielded_converts,
+                ]
+                .concat();
+                sapling_bundle.shielded_spends.clear();
+                sapling_bundle.shielded_outputs.clear();
+                assert_eq!(sapling_bundle.shielded_converts.len(), 2);
+            } else {
+                // ensure we only have one convert proof
+                sapling_bundle.shielded_spends.clear();
+                sapling_bundle.shielded_outputs.clear();
+                assert_eq!(sapling_bundle.shielded_converts.len(), 1);
+            }
+            assert_eq!(sapling_bundle.shielded_spends.len(), 0);
+            assert_eq!(sapling_bundle.shielded_outputs.len(), 0);
+        }
+        RequestedItem::Output => {
+            // From the cost remove the cost of signature(s) validation
+            if multi {
+                // ensure we only have two output proofs
+                sapling_bundle.shielded_outputs = [
+                    sapling_bundle.shielded_outputs.clone(),
+                    sapling_bundle.shielded_outputs,
+                ]
+                .concat();
+                sapling_bundle.shielded_spends.clear();
+                sapling_bundle.shielded_converts.clear();
+                assert_eq!(sapling_bundle.shielded_outputs.len(), 2);
+            } else {
+                // ensure we only have one output proof
+                sapling_bundle.shielded_spends.clear();
+                sapling_bundle.shielded_converts.clear();
+                assert_eq!(sapling_bundle.shielded_outputs.len(), 1);
+            }
+            assert_eq!(sapling_bundle.shielded_spends.len(), 0);
+            assert_eq!(sapling_bundle.shielded_converts.len(), 0);
+        }
+    };
+
+    Some((
+        TransactionData::from_parts(
+            transaction.version(),
+            transaction.consensus_branch_id(),
+            transaction.lock_time(),
+            transaction.expiry_height(),
+            transaction.transparent_bundle().cloned(),
+            Some(sapling_bundle),
+        ),
+        transaction,
+    ))
+}
+
+// For signatures: it's impossible to benchmark more than one signature without
+// pulling in the cost of a proof, so we just benchmark the cost of one binding
+// signature and then use its cost. We don't discount by the numebr of cores
+// because there might be the need for a call to the fallback single
+// verification.. For proofs: benchmarks both one and two proofs and take the
+// difference as the variable cost for every proofs/sigs. Compute the cost of
+// the non parallel parts (with the diff) and charge that if at least one note
+// is present, then charge the variable cost multiplied by the number of notes
+// and divided by the number of cores
+fn masp_batch_validate(c: &mut Criterion) {
+    let mut group = c.benchmark_group("masp_batch_validate");
+    let PVKs {
+        spend_vk,
+        convert_vk,
+        output_vk,
+    } = preload_verifying_keys();
+
+    for multi in [true, false] {
+        for bench in [
+            RequestedItem::Signature,
+            RequestedItem::Spend,
+            RequestedItem::Convert,
+            RequestedItem::Output,
+        ] {
+            // From the cost of proofs remove the cost of signature(s)
+            // validation
+            let (tx_data, transaction) =
+                if let Some(data) = customize_masp_tx_data(multi, &bench) {
+                    data
+                } else {
+                    continue;
+                };
+
+            // Partially deauthorize the transparent bundle
+            let unauth_tx_data =
+                partial_deauthorize(transaction.deref()).unwrap();
+            let txid_parts = unauth_tx_data.digest(TxIdDigester);
+            let sighash = signature_hash(
+                &unauth_tx_data,
+                &SignableInput::Shielded,
+                &txid_parts,
+            )
+            .as_ref()
+            .to_owned();
+            let sapling_bundle = tx_data.sapling_bundle().unwrap();
+
+            let bench_name = if multi {
+                format!("{:#?}_multi", bench)
+            } else {
+                format!("{:#?}_single", bench)
+            };
+            group.bench_function(bench_name, |b| {
+                b.iter_batched(
+                    || {
+                        let mut ctx = BatchValidator::new();
+                        // Check bundle first
+                        if !ctx.check_bundle(sapling_bundle.to_owned(), sighash)
+                        {
+                            panic!("Failed check bundle");
+                        }
+
+                        ctx
+                    },
+                    |ctx| {
+                        assert!(ctx
+                            .validate(spend_vk, convert_vk, output_vk, OsRng))
+                    },
+                    BatchSize::SmallInput,
+                )
+            });
+        }
+    }
+}
+
 fn pgf(c: &mut Criterion) {
     let mut group = c.benchmark_group("vp_pgf");
 
@@ -909,8 +1177,7 @@ fn pgf(c: &mut Criterion) {
                         pgf.ctx.keys_changed,
                         pgf.ctx.verifiers,
                     )
-                    .is_ok()
-                )
+                    .is_ok())
             })
         });
     }
@@ -986,8 +1253,7 @@ fn eth_bridge_nut(c: &mut Criterion) {
                     nut.ctx.keys_changed,
                     nut.ctx.verifiers,
                 )
-                .is_ok()
-            )
+                .is_ok())
         })
     });
 }
@@ -1303,8 +1569,7 @@ fn pos(c: &mut Criterion) {
                         pos.ctx.keys_changed,
                         pos.ctx.verifiers,
                     )
-                    .is_ok()
-                )
+                    .is_ok())
             })
         });
     }
@@ -1439,6 +1704,7 @@ criterion_group!(
     masp_check_convert,
     masp_check_output,
     masp_final_check,
+    masp_batch_validate,
     vp_multitoken,
     pgf,
     eth_bridge_nut,
diff --git a/crates/benches/process_wrapper.rs b/crates/benches/process_wrapper.rs
index 26ad6d2e74..2d1c1241ba 100644
--- a/crates/benches/process_wrapper.rs
+++ b/crates/benches/process_wrapper.rs
@@ -56,7 +56,7 @@ fn process_tx(c: &mut Criterion) {
     let datetime = DateTimeUtc::now();
 
     c.bench_function("wrapper_tx_validation", |b| {
-        b.iter_batched(
+        b.iter_batched_ref(
             || {
                 (
                     // Prevent block out of gas and replay protection
@@ -68,10 +68,10 @@ fn process_tx(c: &mut Criterion) {
                 )
             },
             |(
-                mut temp_state,
-                mut validation_meta,
-                mut vp_wasm_cache,
-                mut tx_wasm_cache,
+                temp_state,
+                validation_meta,
+                vp_wasm_cache,
+                tx_wasm_cache,
                 block_proposer,
             )| {
                 assert_eq!(
@@ -79,12 +79,12 @@ fn process_tx(c: &mut Criterion) {
                     shell
                         .check_proposal_tx(
                             &wrapper,
-                            &mut validation_meta,
-                            &mut temp_state,
+                            validation_meta,
+                            temp_state,
                             datetime,
-                            &mut vp_wasm_cache,
-                            &mut tx_wasm_cache,
-                            &block_proposer
+                            vp_wasm_cache,
+                            tx_wasm_cache,
+                            block_proposer
                         )
                         .code,
                     0

From 87472f07d65a253e46056750a9d29387b60ed5a4 Mon Sep 17 00:00:00 2001
From: Marco Granelli <marco@heliax.dev>
Date: Wed, 15 May 2024 11:58:10 +0200
Subject: [PATCH 3/7] Updates masp dep. Patch reddsa dep

---
 Cargo.lock                   | 158 +++++++++++++++++++++++++++++++++--
 Cargo.toml                   |  10 ++-
 crates/benches/native_vps.rs |  24 ++++--
 wasm/Cargo.lock              |   6 +-
 wasm_for_tests/Cargo.lock    |   6 +-
 5 files changed, 182 insertions(+), 22 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 3b3d63a1ee..aa335fa313 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -310,6 +310,15 @@ dependencies = [
  "rustc_version 0.4.0",
 ]
 
+[[package]]
+name = "atomic-polyfill"
+version = "1.0.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "8cf2bce30dfe09ef0bfaef228b9d414faaf7e563035494d7fe092dba54b300f4"
+dependencies = [
+ "critical-section",
+]
+
 [[package]]
 name = "auto_impl"
 version = "1.1.0"
@@ -1023,6 +1032,12 @@ name = "clru"
 version = "0.5.0"
 source = "git+https://github.com/marmeladema/clru-rs.git?rev=71ca566#71ca566915f21f3c308091ca7756a91b0f8b5afc"
 
+[[package]]
+name = "cobs"
+version = "0.2.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "67ba02a97a2bd10f4b59b25c7973101c79642302776489e030cd13cdab09ed15"
+
 [[package]]
 name = "coins-bip32"
 version = "0.8.7"
@@ -1143,6 +1158,12 @@ dependencies = [
  "windows",
 ]
 
+[[package]]
+name = "const-crc32"
+version = "1.3.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "68d13f542d70e5b339bf46f6f74704ac052cfd526c58cd87996bd1ef4615b9a0"
+
 [[package]]
 name = "const-hex"
 version = "1.10.0"
@@ -1396,6 +1417,12 @@ dependencies = [
  "itertools 0.10.5",
 ]
 
+[[package]]
+name = "critical-section"
+version = "1.1.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7059fff8937831a9ae6f0fe4d658ffabf58f2ca96aa9dec1c889f936f705f216"
+
 [[package]]
 name = "crossbeam-channel"
 version = "0.5.12"
@@ -1599,6 +1626,12 @@ version = "2.5.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "7e962a19be5cfc3f3bf6dd8f61eb50107f356ad6270fbb3ed41476571db78be5"
 
+[[package]]
+name = "debugless-unwrap"
+version = "0.0.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f400d0750c0c069e8493f2256cb4da6f604b6d2eeb69a0ca8863acde352f8400"
+
 [[package]]
 name = "der"
 version = "0.7.8"
@@ -1635,6 +1668,17 @@ dependencies = [
  "syn 1.0.109",
 ]
 
+[[package]]
+name = "derive-getters"
+version = "0.3.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7a2c35ab6e03642397cdda1dd58abbc05d418aef8e36297f336d5aba060fe8df"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn 1.0.109",
+]
+
 [[package]]
 name = "derive_more"
 version = "0.99.17"
@@ -1716,6 +1760,15 @@ version = "0.3.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "fea41bba32d969b513997752735605054bc0dfa92b4c56bf1189f2e174be7a10"
 
+[[package]]
+name = "document-features"
+version = "0.2.8"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ef5282ad69563b5fc40319526ba27e0e7363d552a896f0297d54f767717f9b95"
+dependencies = [
+ "litrs",
+]
+
 [[package]]
 name = "drain_filter_polyfill"
 version = "0.1.3"
@@ -1875,6 +1928,12 @@ dependencies = [
  "zeroize",
 ]
 
+[[package]]
+name = "embedded-io"
+version = "0.4.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ef1a6892d9eef45c8fa6b9e0086428a2cca8491aca8f787c534a3d6d0bcb3ced"
+
 [[package]]
 name = "encoding_rs"
 version = "0.8.33"
@@ -2454,6 +2513,40 @@ dependencies = [
  "num-traits 0.2.17",
 ]
 
+[[package]]
+name = "frost-core"
+version = "1.0.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "45d6280625f1603d160df24b23e4984a6a7286f41455ae606823d0104c32e834"
+dependencies = [
+ "byteorder",
+ "const-crc32",
+ "debugless-unwrap",
+ "derive-getters",
+ "document-features",
+ "hex",
+ "itertools 0.12.1",
+ "postcard",
+ "rand_core 0.6.4",
+ "serde 1.0.193",
+ "serdect",
+ "thiserror",
+ "visibility",
+ "zeroize",
+]
+
+[[package]]
+name = "frost-rerandomized"
+version = "1.0.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "52c58f58ea009000db490efd9a3936d0035647a2b00c7ba8f3868c2ed0306b0b"
+dependencies = [
+ "derive-getters",
+ "document-features",
+ "frost-core",
+ "rand_core 0.6.4",
+]
+
 [[package]]
 name = "fs_extra"
 version = "1.3.0"
@@ -2726,6 +2819,15 @@ version = "1.8.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "eabb4a44450da02c90444cf74558da904edde8fb4e9035a9a6a4e15445af0bd7"
 
+[[package]]
+name = "hash32"
+version = "0.2.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b0c35f58762feb77d74ebe43bdbc3210f09be9fe6742234d573bacc26ed92b67"
+dependencies = [
+ "byteorder",
+]
+
 [[package]]
 name = "hashbrown"
 version = "0.11.2"
@@ -2793,6 +2895,20 @@ dependencies = [
  "http 0.2.11",
 ]
 
+[[package]]
+name = "heapless"
+version = "0.7.17"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "cdc6457c0eb62c71aac4bc17216026d8410337c4126773b9c5daba343f17964f"
+dependencies = [
+ "atomic-polyfill",
+ "hash32",
+ "rustc_version 0.4.0",
+ "serde 1.0.193",
+ "spin 0.9.8",
+ "stable_deref_trait",
+]
+
 [[package]]
 name = "heck"
 version = "0.4.1"
@@ -4078,6 +4194,12 @@ version = "0.4.12"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "c4cd1a83af159aa67994778be9070f0ae1bd732942279cabb14f86f986a21456"
 
+[[package]]
+name = "litrs"
+version = "0.4.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b4ce301924b7887e9d637144fdade93f9dfff9b60981d4ac161db09720d39aa5"
+
 [[package]]
 name = "lock_api"
 version = "0.4.11"
@@ -4141,7 +4263,7 @@ dependencies = [
 [[package]]
 name = "masp_note_encryption"
 version = "1.0.0"
-source = "git+https://github.com/anoma/masp?rev=3aacc707c5948e7423589ac617305448bead9842#3aacc707c5948e7423589ac617305448bead9842"
+source = "git+https://github.com/anoma/masp?rev=3689da58610bb65aaa303912750f1fa89f8c4a34#3689da58610bb65aaa303912750f1fa89f8c4a34"
 dependencies = [
  "borsh 1.2.1",
  "chacha20",
@@ -4154,7 +4276,7 @@ dependencies = [
 [[package]]
 name = "masp_primitives"
 version = "1.0.0"
-source = "git+https://github.com/anoma/masp?rev=3aacc707c5948e7423589ac617305448bead9842#3aacc707c5948e7423589ac617305448bead9842"
+source = "git+https://github.com/anoma/masp?rev=3689da58610bb65aaa303912750f1fa89f8c4a34#3689da58610bb65aaa303912750f1fa89f8c4a34"
 dependencies = [
  "aes",
  "bip0039",
@@ -4186,7 +4308,7 @@ dependencies = [
 [[package]]
 name = "masp_proofs"
 version = "1.0.0"
-source = "git+https://github.com/anoma/masp?rev=3aacc707c5948e7423589ac617305448bead9842#3aacc707c5948e7423589ac617305448bead9842"
+source = "git+https://github.com/anoma/masp?rev=3689da58610bb65aaa303912750f1fa89f8c4a34#3689da58610bb65aaa303912750f1fa89f8c4a34"
 dependencies = [
  "bellman",
  "blake2b_simd",
@@ -5980,6 +6102,18 @@ dependencies = [
  "universal-hash",
 ]
 
+[[package]]
+name = "postcard"
+version = "1.0.8"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a55c51ee6c0db07e68448e336cf8ea4131a620edefebf9893e759b2d793420f8"
+dependencies = [
+ "cobs",
+ "embedded-io",
+ "heapless",
+ "serde 1.0.193",
+]
+
 [[package]]
 name = "powerfmt"
 version = "0.2.0"
@@ -6391,11 +6525,11 @@ dependencies = [
 [[package]]
 name = "reddsa"
 version = "0.5.1"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "78a5191930e84973293aa5f532b513404460cd2216c1cfb76d08748c15b40b02"
+source = "git+https://github.com/heliaxdev/reddsa?rev=46d363b929e1b940688fa0c53d637e304a755185#46d363b929e1b940688fa0c53d637e304a755185"
 dependencies = [
  "blake2b_simd",
  "byteorder",
+ "frost-rerandomized",
  "group",
  "hex",
  "jubjub",
@@ -7376,6 +7510,9 @@ name = "spin"
 version = "0.9.8"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "6980e8d7511241f8acf4aebddbb1ff938df5eebe98691418c4468d0b72a96a67"
+dependencies = [
+ "lock_api",
+]
 
 [[package]]
 name = "spki"
@@ -8566,6 +8703,17 @@ version = "0.9.4"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "49874b5167b65d7193b8aba1567f5c7d93d001cafc34600cee003eda787e483f"
 
+[[package]]
+name = "visibility"
+version = "0.1.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b3fd98999db9227cf28e59d83e1f120f42bc233d4b152e8fab9bc87d5bb1e0f8"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn 2.0.52",
+]
+
 [[package]]
 name = "wait-timeout"
 version = "0.2.0"
diff --git a/Cargo.toml b/Cargo.toml
index 0d1f58a577..e35094a7f4 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -124,9 +124,9 @@ ledger-transport-hid = "0.10.0"
 libc = "0.2.97"
 libloading = "0.7.2"
 linkme = "0.3.24"
-# branch = "main"
-masp_primitives = { git = "https://github.com/anoma/masp", rev = "3aacc707c5948e7423589ac617305448bead9842" }
-masp_proofs = { git = "https://github.com/anoma/masp", rev = "3aacc707c5948e7423589ac617305448bead9842", default-features = false, features = ["local-prover"] }
+# branch = "murisi/namada-integration"
+masp_primitives = { git = "https://github.com/anoma/masp", rev = "3689da58610bb65aaa303912750f1fa89f8c4a34" }
+masp_proofs = { git = "https://github.com/anoma/masp", rev = "3689da58610bb65aaa303912750f1fa89f8c4a34", default-features = false, features = ["local-prover"] }
 num256 = "0.3.5"
 num_cpus = "1.13.0"
 num-derive = "0.4"
@@ -196,6 +196,10 @@ winapi = "0.3.9"
 yansi = "0.5.1"
 zeroize = { version = "1.5.5", features = ["zeroize_derive"] }
 
+[patch.crates-io]
+# Patch to the fork containing the correct personalization and basepoints for masp
+reddsa = { git = "https://github.com/heliaxdev/reddsa", rev = "46d363b929e1b940688fa0c53d637e304a755185" }
+
 [profile.release]
 lto = true
 opt-level = 3
diff --git a/crates/benches/native_vps.rs b/crates/benches/native_vps.rs
index 965a3dfe65..28dae6a57c 100644
--- a/crates/benches/native_vps.rs
+++ b/crates/benches/native_vps.rs
@@ -457,7 +457,8 @@ fn ibc(c: &mut Criterion) {
                         ibc.ctx.keys_changed,
                         ibc.ctx.verifiers,
                     )
-                    .is_ok())
+                    .is_ok()
+                )
             })
         });
     }
@@ -634,7 +635,8 @@ fn masp(c: &mut Criterion) {
                         masp.ctx.keys_changed,
                         masp.ctx.verifiers,
                     )
-                    .is_ok());
+                    .is_ok()
+                );
             })
         });
     }
@@ -918,7 +920,7 @@ fn customize_masp_tx_data(
     TransactionData<masp_primitives::transaction::Authorized>,
     Transaction,
 )> {
-    let (_, tx) = setup_storage_for_masp_verification("unshielding");
+    let (_, _, tx) = setup_storage_for_masp_verification("unshielding");
     let transaction = tx
         .sections
         .into_iter()
@@ -1087,8 +1089,11 @@ fn masp_batch_validate(c: &mut Criterion) {
                         ctx
                     },
                     |ctx| {
-                        assert!(ctx
-                            .validate(spend_vk, convert_vk, output_vk, OsRng))
+                        assert!(
+                            ctx.validate(
+                                spend_vk, convert_vk, output_vk, OsRng
+                            )
+                        )
                     },
                     BatchSize::SmallInput,
                 )
@@ -1177,7 +1182,8 @@ fn pgf(c: &mut Criterion) {
                         pgf.ctx.keys_changed,
                         pgf.ctx.verifiers,
                     )
-                    .is_ok())
+                    .is_ok()
+                )
             })
         });
     }
@@ -1253,7 +1259,8 @@ fn eth_bridge_nut(c: &mut Criterion) {
                     nut.ctx.keys_changed,
                     nut.ctx.verifiers,
                 )
-                .is_ok())
+                .is_ok()
+            )
         })
     });
 }
@@ -1569,7 +1576,8 @@ fn pos(c: &mut Criterion) {
                         pos.ctx.keys_changed,
                         pos.ctx.verifiers,
                     )
-                    .is_ok())
+                    .is_ok()
+                )
             })
         });
     }
diff --git a/wasm/Cargo.lock b/wasm/Cargo.lock
index 2a00208e8a..4a07422fff 100644
--- a/wasm/Cargo.lock
+++ b/wasm/Cargo.lock
@@ -3384,7 +3384,7 @@ dependencies = [
 [[package]]
 name = "masp_note_encryption"
 version = "1.0.0"
-source = "git+https://github.com/anoma/masp?rev=3aacc707c5948e7423589ac617305448bead9842#3aacc707c5948e7423589ac617305448bead9842"
+source = "git+https://github.com/anoma/masp?rev=3689da58610bb65aaa303912750f1fa89f8c4a34#3689da58610bb65aaa303912750f1fa89f8c4a34"
 dependencies = [
  "borsh 1.4.0",
  "chacha20",
@@ -3397,7 +3397,7 @@ dependencies = [
 [[package]]
 name = "masp_primitives"
 version = "1.0.0"
-source = "git+https://github.com/anoma/masp?rev=3aacc707c5948e7423589ac617305448bead9842#3aacc707c5948e7423589ac617305448bead9842"
+source = "git+https://github.com/anoma/masp?rev=3689da58610bb65aaa303912750f1fa89f8c4a34#3689da58610bb65aaa303912750f1fa89f8c4a34"
 dependencies = [
  "aes",
  "bip0039",
@@ -3429,7 +3429,7 @@ dependencies = [
 [[package]]
 name = "masp_proofs"
 version = "1.0.0"
-source = "git+https://github.com/anoma/masp?rev=3aacc707c5948e7423589ac617305448bead9842#3aacc707c5948e7423589ac617305448bead9842"
+source = "git+https://github.com/anoma/masp?rev=3689da58610bb65aaa303912750f1fa89f8c4a34#3689da58610bb65aaa303912750f1fa89f8c4a34"
 dependencies = [
  "bellman",
  "blake2b_simd",
diff --git a/wasm_for_tests/Cargo.lock b/wasm_for_tests/Cargo.lock
index 59aa30e730..5afaaed8f5 100644
--- a/wasm_for_tests/Cargo.lock
+++ b/wasm_for_tests/Cargo.lock
@@ -3364,7 +3364,7 @@ dependencies = [
 [[package]]
 name = "masp_note_encryption"
 version = "1.0.0"
-source = "git+https://github.com/anoma/masp?rev=3aacc707c5948e7423589ac617305448bead9842#3aacc707c5948e7423589ac617305448bead9842"
+source = "git+https://github.com/anoma/masp?rev=3689da58610bb65aaa303912750f1fa89f8c4a34#3689da58610bb65aaa303912750f1fa89f8c4a34"
 dependencies = [
  "borsh 1.2.1",
  "chacha20",
@@ -3377,7 +3377,7 @@ dependencies = [
 [[package]]
 name = "masp_primitives"
 version = "1.0.0"
-source = "git+https://github.com/anoma/masp?rev=3aacc707c5948e7423589ac617305448bead9842#3aacc707c5948e7423589ac617305448bead9842"
+source = "git+https://github.com/anoma/masp?rev=3689da58610bb65aaa303912750f1fa89f8c4a34#3689da58610bb65aaa303912750f1fa89f8c4a34"
 dependencies = [
  "aes",
  "bip0039",
@@ -3409,7 +3409,7 @@ dependencies = [
 [[package]]
 name = "masp_proofs"
 version = "1.0.0"
-source = "git+https://github.com/anoma/masp?rev=3aacc707c5948e7423589ac617305448bead9842#3aacc707c5948e7423589ac617305448bead9842"
+source = "git+https://github.com/anoma/masp?rev=3689da58610bb65aaa303912750f1fa89f8c4a34#3689da58610bb65aaa303912750f1fa89f8c4a34"
 dependencies = [
  "bellman",
  "blake2b_simd",

From ee2fd4600a88160608cc7e9fbfc8fe1de6f69bce Mon Sep 17 00:00:00 2001
From: Marco Granelli <marco@heliax.dev>
Date: Thu, 16 May 2024 17:08:19 +0200
Subject: [PATCH 4/7] Fixes masp benchmarks

---
 crates/benches/Cargo.toml    |  26 +--
 crates/benches/native_vps.rs | 324 +++++++++++++++++++++--------------
 2 files changed, 210 insertions(+), 140 deletions(-)

diff --git a/crates/benches/Cargo.toml b/crates/benches/Cargo.toml
index 01d25fa777..74bf5960e0 100644
--- a/crates/benches/Cargo.toml
+++ b/crates/benches/Cargo.toml
@@ -17,20 +17,20 @@ name = "native_vps"
 harness = false
 path = "native_vps.rs"
 
-[[bench]]
-name = "process_wrapper"
-harness = false
-path = "process_wrapper.rs"
+# [[bench]]
+# name = "process_wrapper"
+# harness = false
+# path = "process_wrapper.rs"
 
-[[bench]]
-name = "host_env"
-harness = false
-path = "host_env.rs"
+# [[bench]]
+# name = "host_env"
+# harness = false
+# path = "host_env.rs"
 
-[[bench]]
-name = "wasm_opcodes"
-harness = false
-path = "wasm_opcodes.rs"
+# [[bench]]
+# name = "wasm_opcodes"
+# harness = false
+# path = "wasm_opcodes.rs"
 
 [features]
 namada-eth-bridge = [
@@ -46,7 +46,7 @@ namada = { path = "../namada", features = ["rand", "benches"] }
 namada_apps_lib = { path = "../apps_lib" }
 namada_node = { path = "../node", features = ["benches"] }
 masp_primitives.workspace = true
-masp_proofs = { workspace = true, features = ["multicore"] }
+masp_proofs = { workspace = true, features = ["benchmarks"] }
 borsh.workspace = true
 borsh-ext.workspace = true
 criterion = { version = "0.5", features = ["html_reports"] }
diff --git a/crates/benches/native_vps.rs b/crates/benches/native_vps.rs
index 28dae6a57c..5f2b4455e4 100644
--- a/crates/benches/native_vps.rs
+++ b/crates/benches/native_vps.rs
@@ -904,24 +904,23 @@ fn masp_final_check(c: &mut Criterion) {
 }
 
 #[derive(Debug)]
-enum RequestedItem {
-    Signature,
+enum BenchNote {
     Spend,
     Convert,
     Output,
 }
 
-// Removes the unneeded notes from the generated transaction and replicates the
-// remaining ones if needed
+// Tweaks the transaction to match the desired benchmark
 fn customize_masp_tx_data(
     multi: bool,
-    request: &RequestedItem,
-) -> Option<(
+    request: &BenchNote,
+) -> (
     TransactionData<masp_primitives::transaction::Authorized>,
     Transaction,
-)> {
+) {
     let (_, _, tx) = setup_storage_for_masp_verification("unshielding");
     let transaction = tx
+        .tx
         .sections
         .into_iter()
         .filter_map(|section| match section {
@@ -935,84 +934,47 @@ fn customize_masp_tx_data(
     let mut sapling_bundle = transaction.sapling_bundle().unwrap().to_owned();
 
     match request {
-        RequestedItem::Signature => {
+        BenchNote::Spend => {
             if multi {
-                // No multisig benchmark
-                return None;
-            }
-            // ensure we only have one signature to verify (the binding one) and
-            // no proofs
-            sapling_bundle.shielded_spends.clear();
-            sapling_bundle.shielded_converts.clear();
-            sapling_bundle.shielded_outputs.clear();
-            assert_eq!(sapling_bundle.shielded_spends.len(), 0);
-            assert_eq!(sapling_bundle.shielded_converts.len(), 0);
-            assert_eq!(sapling_bundle.shielded_outputs.len(), 0);
-        }
-        RequestedItem::Spend => {
-            if multi {
-                // ensure we only have two spend proofs
+                // ensure we have two spend proofs
                 sapling_bundle.shielded_spends = [
                     sapling_bundle.shielded_spends.clone(),
                     sapling_bundle.shielded_spends,
                 ]
                 .concat();
-                sapling_bundle.shielded_outputs.clear();
-                sapling_bundle.shielded_converts.clear();
                 assert_eq!(sapling_bundle.shielded_spends.len(), 2);
             } else {
-                // ensure we only have one spend proof
-                sapling_bundle.shielded_outputs.clear();
-                sapling_bundle.shielded_converts.clear();
+                // ensure we have one spend proof
                 assert_eq!(sapling_bundle.shielded_spends.len(), 1);
             }
-            assert_eq!(sapling_bundle.shielded_converts.len(), 0);
-            assert_eq!(sapling_bundle.shielded_outputs.len(), 0);
         }
-        RequestedItem::Convert => {
+        BenchNote::Convert => {
             if multi {
-                // ensure we only have two convert proofs
+                // ensure we have two convert proofs
                 sapling_bundle.shielded_converts = [
                     sapling_bundle.shielded_converts.clone(),
                     sapling_bundle.shielded_converts,
                 ]
                 .concat();
-                sapling_bundle.shielded_spends.clear();
-                sapling_bundle.shielded_outputs.clear();
                 assert_eq!(sapling_bundle.shielded_converts.len(), 2);
             } else {
-                // ensure we only have one convert proof
-                sapling_bundle.shielded_spends.clear();
-                sapling_bundle.shielded_outputs.clear();
+                // ensure we have one convert proof
                 assert_eq!(sapling_bundle.shielded_converts.len(), 1);
             }
-            assert_eq!(sapling_bundle.shielded_spends.len(), 0);
-            assert_eq!(sapling_bundle.shielded_outputs.len(), 0);
         }
-        RequestedItem::Output => {
-            // From the cost remove the cost of signature(s) validation
+        BenchNote::Output => {
             if multi {
-                // ensure we only have two output proofs
-                sapling_bundle.shielded_outputs = [
-                    sapling_bundle.shielded_outputs.clone(),
-                    sapling_bundle.shielded_outputs,
-                ]
-                .concat();
-                sapling_bundle.shielded_spends.clear();
-                sapling_bundle.shielded_converts.clear();
+                // ensure we have two output proofs
                 assert_eq!(sapling_bundle.shielded_outputs.len(), 2);
             } else {
-                // ensure we only have one output proof
-                sapling_bundle.shielded_spends.clear();
-                sapling_bundle.shielded_converts.clear();
+                // ensure we have one output proof
+                sapling_bundle.shielded_outputs.truncate(1);
                 assert_eq!(sapling_bundle.shielded_outputs.len(), 1);
             }
-            assert_eq!(sapling_bundle.shielded_spends.len(), 0);
-            assert_eq!(sapling_bundle.shielded_converts.len(), 0);
         }
     };
 
-    Some((
+    (
         TransactionData::from_parts(
             transaction.version(),
             transaction.consensus_branch_id(),
@@ -1022,84 +984,189 @@ fn customize_masp_tx_data(
             Some(sapling_bundle),
         ),
         transaction,
-    ))
+    )
 }
 
-// For signatures: it's impossible to benchmark more than one signature without
-// pulling in the cost of a proof, so we just benchmark the cost of one binding
-// signature and then use its cost. We don't discount by the numebr of cores
-// because there might be the need for a call to the fallback single
-// verification.. For proofs: benchmarks both one and two proofs and take the
-// difference as the variable cost for every proofs/sigs. Compute the cost of
-// the non parallel parts (with the diff) and charge that if at least one note
-// is present, then charge the variable cost multiplied by the number of notes
-// and divided by the number of cores
-fn masp_batch_validate(c: &mut Criterion) {
-    let mut group = c.benchmark_group("masp_batch_validate");
-    let PVKs {
-        spend_vk,
-        convert_vk,
-        output_vk,
-    } = preload_verifying_keys();
-
-    for multi in [true, false] {
-        for bench in [
-            RequestedItem::Signature,
-            RequestedItem::Spend,
-            RequestedItem::Convert,
-            RequestedItem::Output,
-        ] {
-            // From the cost of proofs remove the cost of signature(s)
-            // validation
-            let (tx_data, transaction) =
-                if let Some(data) = customize_masp_tx_data(multi, &bench) {
-                    data
-                } else {
-                    continue;
-                };
+// benchmark the cost of validating two signatures in a batch.
+fn masp_batch_signature_verification(c: &mut Criterion) {
+    let (_, _, tx) = setup_storage_for_masp_verification("unshielding");
+    let transaction = tx
+        .tx
+        .sections
+        .into_iter()
+        .filter_map(|section| match section {
+            Section::MaspTx(transaction) => Some(transaction),
+            _ => None,
+        })
+        .collect::<Vec<Transaction>>()
+        .first()
+        .unwrap()
+        .to_owned();
+    let sapling_bundle = transaction.sapling_bundle().unwrap();
+    // ensure we have two signatures to verify (the binding and one spending)
+    assert_eq!(sapling_bundle.shielded_spends.len(), 1);
 
-            // Partially deauthorize the transparent bundle
-            let unauth_tx_data =
-                partial_deauthorize(transaction.deref()).unwrap();
-            let txid_parts = unauth_tx_data.digest(TxIdDigester);
-            let sighash = signature_hash(
-                &unauth_tx_data,
-                &SignableInput::Shielded,
-                &txid_parts,
-            )
+    // Partially deauthorize the transparent bundle
+    let unauth_tx_data = partial_deauthorize(transaction.deref()).unwrap();
+    let txid_parts = unauth_tx_data.digest(TxIdDigester);
+    let sighash =
+        signature_hash(&unauth_tx_data, &SignableInput::Shielded, &txid_parts)
             .as_ref()
             .to_owned();
-            let sapling_bundle = tx_data.sapling_bundle().unwrap();
 
-            let bench_name = if multi {
-                format!("{:#?}_multi", bench)
-            } else {
-                format!("{:#?}_single", bench)
-            };
-            group.bench_function(bench_name, |b| {
-                b.iter_batched(
-                    || {
-                        let mut ctx = BatchValidator::new();
-                        // Check bundle first
-                        if !ctx.check_bundle(sapling_bundle.to_owned(), sighash)
-                        {
-                            panic!("Failed check bundle");
-                        }
-
-                        ctx
-                    },
-                    |ctx| {
-                        assert!(
-                            ctx.validate(
-                                spend_vk, convert_vk, output_vk, OsRng
-                            )
-                        )
-                    },
-                    BatchSize::SmallInput,
-                )
-            });
-        }
+    c.bench_function("masp_batch_signature_verification", |b| {
+        b.iter_batched(
+            || {
+                let mut ctx = BatchValidator::new();
+                // Check bundle first
+                if !ctx.check_bundle(sapling_bundle.to_owned(), sighash) {
+                    panic!("Failed check bundle");
+                }
+
+                ctx
+            },
+            |ctx| assert!(ctx.verify_signatures(OsRng).is_ok()),
+            BatchSize::SmallInput,
+        )
+    });
+}
+
+// Benchmark both one and two proofs and take the difference as the variable
+// cost for every proofs. Charge the full cost for the first note and then
+// charge the variable cost multiplied by the number of remaining notes and
+// divided by the number of cores
+fn masp_batch_spend_proofs_validate(c: &mut Criterion) {
+    let mut group = c.benchmark_group("masp_batch_spend_proofs_validate");
+    let PVKs { spend_vk, .. } = preload_verifying_keys();
+
+    for double in [true, false] {
+        let (tx_data, transaction) =
+            customize_masp_tx_data(double, &BenchNote::Spend);
+
+        // Partially deauthorize the transparent bundle
+        let unauth_tx_data = partial_deauthorize(transaction.deref()).unwrap();
+        let txid_parts = unauth_tx_data.digest(TxIdDigester);
+        // Compute the sighash from the original, unmodified transaction
+        let sighash = signature_hash(
+            &unauth_tx_data,
+            &SignableInput::Shielded,
+            &txid_parts,
+        )
+        .as_ref()
+        .to_owned();
+        let sapling_bundle = tx_data.sapling_bundle().unwrap();
+
+        let bench_name = if double { "double" } else { "single" };
+        group.bench_function(bench_name, |b| {
+            b.iter_batched(
+                || {
+                    let mut ctx = BatchValidator::new();
+                    // Check bundle first
+                    if !ctx.check_bundle(sapling_bundle.to_owned(), sighash) {
+                        panic!("Failed check bundle");
+                    }
+
+                    ctx
+                },
+                |ctx| assert!(ctx.verify_spend_proofs(spend_vk).is_ok()),
+                BatchSize::SmallInput,
+            )
+        });
+    }
+
+    group.finish();
+}
+
+// Benchmark both one and two proofs and take the difference as the variable
+// cost for every proofs. Charge the full cost for the first note and then
+// charge the variable cost multiplied by the number of remaining notes and
+// divided by the number of cores
+fn masp_batch_convert_proofs_validate(c: &mut Criterion) {
+    let mut group = c.benchmark_group("masp_batch_convert_proofs_validate");
+    let PVKs { convert_vk, .. } = preload_verifying_keys();
+
+    for double in [true, false] {
+        let (tx_data, transaction) =
+            customize_masp_tx_data(double, &BenchNote::Convert);
+
+        // Partially deauthorize the transparent bundle
+        let unauth_tx_data = partial_deauthorize(transaction.deref()).unwrap();
+        let txid_parts = unauth_tx_data.digest(TxIdDigester);
+        // Compute the sighash from the original, unmodified transaction
+        let sighash = signature_hash(
+            &unauth_tx_data,
+            &SignableInput::Shielded,
+            &txid_parts,
+        )
+        .as_ref()
+        .to_owned();
+        let sapling_bundle = tx_data.sapling_bundle().unwrap();
+
+        let bench_name = if double { "double" } else { "single" };
+        group.bench_function(bench_name, |b| {
+            b.iter_batched(
+                || {
+                    let mut ctx = BatchValidator::new();
+                    // Check bundle first
+                    if !ctx.check_bundle(sapling_bundle.to_owned(), sighash) {
+                        panic!("Failed check bundle");
+                    }
+
+                    ctx
+                },
+                |ctx| assert!(ctx.verify_convert_proofs(convert_vk).is_ok()),
+                BatchSize::SmallInput,
+            )
+        });
+    }
+
+    group.finish();
+}
+
+// Benchmark both one and two proofs and take the difference as the variable
+// cost for every proofs. Charge the full cost for the first note and then
+// charge the variable cost multiplied by the number of remaining notes and
+// divided by the number of cores
+fn masp_batch_output_proofs_validate(c: &mut Criterion) {
+    let mut group = c.benchmark_group("masp_batch_output_proofs_validate");
+    let PVKs { output_vk, .. } = preload_verifying_keys();
+
+    for double in [true, false] {
+        let (tx_data, transaction) =
+            customize_masp_tx_data(double, &BenchNote::Output);
+
+        // Partially deauthorize the transparent bundle
+        let unauth_tx_data = partial_deauthorize(transaction.deref()).unwrap();
+        let txid_parts = unauth_tx_data.digest(TxIdDigester);
+        // Compute the sighash from the original, unmodified transaction
+        let sighash = signature_hash(
+            &unauth_tx_data,
+            &SignableInput::Shielded,
+            &txid_parts,
+        )
+        .as_ref()
+        .to_owned();
+        let sapling_bundle = tx_data.sapling_bundle().unwrap();
+
+        let bench_name = if double { "double" } else { "single" };
+        group.bench_function(bench_name, |b| {
+            b.iter_batched(
+                || {
+                    let mut ctx = BatchValidator::new();
+                    // Check bundle first
+                    if !ctx.check_bundle(sapling_bundle.to_owned(), sighash) {
+                        panic!("Failed check bundle");
+                    }
+
+                    ctx
+                },
+                |ctx| assert!(ctx.verify_output_proofs(output_vk).is_ok()),
+                BatchSize::SmallInput,
+            )
+        });
     }
+
+    group.finish();
 }
 
 fn pgf(c: &mut Criterion) {
@@ -1712,7 +1779,10 @@ criterion_group!(
     masp_check_convert,
     masp_check_output,
     masp_final_check,
-    masp_batch_validate,
+    masp_batch_signature_verification,
+    masp_batch_spend_proofs_validate,
+    masp_batch_convert_proofs_validate,
+    masp_batch_output_proofs_validate,
     vp_multitoken,
     pgf,
     eth_bridge_nut,

From c2849f883e4fa3a7956436250e5b5fd85abd57e9 Mon Sep 17 00:00:00 2001
From: Marco Granelli <marco@heliax.dev>
Date: Thu, 16 May 2024 19:11:18 +0200
Subject: [PATCH 5/7] Reworks gas metering for masp validation

---
 crates/sdk/src/masp.rs | 151 ++++++++++++++---------------------------
 1 file changed, 52 insertions(+), 99 deletions(-)

diff --git a/crates/sdk/src/masp.rs b/crates/sdk/src/masp.rs
index caa9a19d5a..b2934f341e 100644
--- a/crates/sdk/src/masp.rs
+++ b/crates/sdk/src/masp.rs
@@ -30,6 +30,9 @@ use masp_primitives::transaction::builder::{self, *};
 use masp_primitives::transaction::components::sapling::builder::{
     RngBuildParams, SaplingMetadata,
 };
+use masp_primitives::transaction::components::sapling::{
+    Authorized as SaplingAuthorized, Bundle as SaplingBundle,
+};
 use masp_primitives::transaction::components::transparent::builder::TransparentBuilder;
 use masp_primitives::transaction::components::{
     I128Sum, OutputDescription, TxOut, U64Sum, ValueSum,
@@ -70,6 +73,7 @@ use rand::rngs::StdRng;
 use rand_core::{CryptoRng, OsRng, RngCore, SeedableRng};
 use ripemd::Digest as RipemdDigest;
 use sha2::Digest;
+use smooth_operator::checked;
 use thiserror::Error;
 
 use crate::error::{Error, QueryError};
@@ -288,15 +292,7 @@ where
     let mut ctx = testing::MockBatchValidator::default();
 
     // Charge gas before check bundle
-    let spends_len = sapling_bundle.shielded_spends.len() as u64;
-    let converts_len = sapling_bundle.shielded_converts.len() as u64;
-    let outputs_len = sapling_bundle.shielded_outputs.len() as u64;
-    charge_masp_check_bundle_gas(
-        spends_len,
-        converts_len,
-        outputs_len,
-        &consume_verify_gas,
-    )?;
+    charge_masp_check_bundle_gas(sapling_bundle, &consume_verify_gas)?;
 
     if !ctx.check_bundle(sapling_bundle.to_owned(), sighash.as_ref().to_owned())
     {
@@ -306,12 +302,7 @@ where
     tracing::debug!("passed check bundle");
 
     // Charge gas before final validation
-    charge_masp_validate_gas(
-        spends_len,
-        converts_len,
-        outputs_len,
-        consume_verify_gas,
-    )?;
+    charge_masp_validate_gas(sapling_bundle, consume_verify_gas)?;
     if !ctx.validate(spend_vk, convert_vk, output_vk, OsRng) {
         return Err(StorageError::SimpleMessage(
             "Invalid proofs or signatures",
@@ -322,113 +313,75 @@ where
 
 // Charge gas for the check_bundle operation which does not leverage concurrency
 fn charge_masp_check_bundle_gas<F>(
-    spends_len: u64,
-    converts_len: u64,
-    outputs_len: u64,
+    sapling_bundle: &SaplingBundle<SaplingAuthorized>,
     consume_verify_gas: F,
 ) -> Result<(), namada_state::StorageError>
 where
     F: Fn(u64) -> std::result::Result<(), namada_state::StorageError>,
 {
-    consume_verify_gas(
-        spends_len
-            .checked_mul(namada_gas::MASP_SPEND_CHECK_GAS)
-            .ok_or(namada_state::StorageError::SimpleMessage(
-                "Overflow in gas calculation",
-            ))?,
-    )?;
-
-    consume_verify_gas(
-        converts_len
-            .checked_mul(namada_gas::MASP_CONVERT_CHECK_GAS)
-            .ok_or(namada_state::StorageError::SimpleMessage(
-                "Overflow in gas calculation",
-            ))?,
-    )?;
-
-    consume_verify_gas(
-        outputs_len
-            .checked_mul(namada_gas::MASP_OUTPUT_CHECK_GAS)
-            .ok_or(namada_state::StorageError::SimpleMessage(
-                "Overflow in gas calculation",
-            ))?,
-    )?;
-
-    Ok(())
+    consume_verify_gas(checked!(
+        (sapling_bundle.shielded_spends.len() as u64)
+            * namada_gas::MASP_SPEND_CHECK_GAS
+    )?)?;
+
+    consume_verify_gas(checked!(
+        (sapling_bundle.shielded_converts.len() as u64)
+            * namada_gas::MASP_CONVERT_CHECK_GAS
+    )?)?;
+
+    consume_verify_gas(checked!(
+        (sapling_bundle.shielded_outputs.len() as u64)
+            * namada_gas::MASP_OUTPUT_CHECK_GAS
+    )?)
 }
 
 // Charge gas for the final validation, taking advtange of concurrency for
-// proofs verification but not for signature (because of the fallback)
+// proofs verification but not for signatures
 fn charge_masp_validate_gas<F>(
-    spends_len: u64,
-    converts_len: u64,
-    outputs_len: u64,
+    sapling_bundle: &SaplingBundle<SaplingAuthorized>,
     consume_verify_gas: F,
 ) -> Result<(), namada_state::StorageError>
 where
     F: Fn(u64) -> std::result::Result<(), namada_state::StorageError>,
 {
-    // Signatures pay the entire cost (no discount because of the possible
-    // fallback)
-    consume_verify_gas(
-        spends_len
-            // Add one for the binding signature
-            .checked_add(1)
-            .ok_or(namada_state::StorageError::SimpleMessage(
-                "Overflow in gas calculation",
-            ))?
-            .checked_mul(namada_gas::MASP_VERIFY_SIG_GAS)
-            .ok_or(namada_state::StorageError::SimpleMessage(
-                "Overflow in gas calculation",
-            ))?,
-    )?;
+    // Signatures gas
+    consume_verify_gas(checked!(
+        // Add one for the binding signature
+        ((sapling_bundle.shielded_spends.len() as u64) + 1)
+            * namada_gas::MASP_VERIFY_SIG_GAS
+    )?)?;
 
     // If at least one note is present charge the fixed costs. Then charge the
-    // variable cost for every note, amortized on the fixed expected number of
-    // cores
-    if spends_len != 0 {
+    // variable cost for every other note, amortized on the fixed expected
+    // number of cores
+    if let Some(remaining_notes) =
+        sapling_bundle.shielded_spends.len().checked_sub(1)
+    {
         consume_verify_gas(namada_gas::MASP_FIXED_SPEND_GAS)?;
-        consume_verify_gas(
-            namada_gas::MASP_VARIABLE_SPEND_GAS
-                .checked_mul(spends_len)
-                .ok_or(namada_state::StorageError::SimpleMessage(
-                    "Overflow in gas calculation",
-                ))?
-                .checked_div(namada_gas::MASP_PARALLEL_GAS_DIVIDER)
-                .ok_or(namada_state::StorageError::SimpleMessage(
-                    "Overflow in gas calculation",
-                ))?,
-        )?;
+        consume_verify_gas(checked!(
+            namada_gas::MASP_VARIABLE_SPEND_GAS * remaining_notes as u64
+                / namada_gas::MASP_PARALLEL_GAS_DIVIDER
+        )?)?;
     }
 
-    if converts_len != 0 {
+    if let Some(remaining_notes) =
+        sapling_bundle.shielded_converts.len().checked_sub(1)
+    {
         consume_verify_gas(namada_gas::MASP_FIXED_CONVERT_GAS)?;
-        consume_verify_gas(
-            namada_gas::MASP_VARIABLE_CONVERT_GAS
-                .checked_mul(converts_len)
-                .ok_or(namada_state::StorageError::SimpleMessage(
-                    "Overflow in gas calculation",
-                ))?
-                .checked_div(namada_gas::MASP_PARALLEL_GAS_DIVIDER)
-                .ok_or(namada_state::StorageError::SimpleMessage(
-                    "Overflow in gas calculation",
-                ))?,
-        )?;
+        consume_verify_gas(checked!(
+            namada_gas::MASP_VARIABLE_CONVERT_GAS * remaining_notes as u64
+                / namada_gas::MASP_PARALLEL_GAS_DIVIDER
+        )?)?;
     }
 
-    if outputs_len != 0 {
+    if let Some(remaining_notes) =
+        sapling_bundle.shielded_outputs.len().checked_sub(1)
+    {
         consume_verify_gas(namada_gas::MASP_FIXED_OUTPUT_GAS)?;
-        consume_verify_gas(
-            namada_gas::MASP_VARIABLE_OUTPUT_GAS
-                .checked_mul(outputs_len)
-                .ok_or(namada_state::StorageError::SimpleMessage(
-                    "Overflow in gas calculation",
-                ))?
-                .checked_div(namada_gas::MASP_PARALLEL_GAS_DIVIDER)
-                .ok_or(namada_state::StorageError::SimpleMessage(
-                    "Overflow in gas calculation",
-                ))?,
-        )?;
+        consume_verify_gas(checked!(
+            namada_gas::MASP_VARIABLE_OUTPUT_GAS * remaining_notes as u64
+                / namada_gas::MASP_PARALLEL_GAS_DIVIDER
+        )?)?;
     }
 
     Ok(())

From b8b25f2461c7953c071adf21c6122d05aac6d939 Mon Sep 17 00:00:00 2001
From: Marco Granelli <marco@heliax.dev>
Date: Fri, 17 May 2024 19:25:40 +0200
Subject: [PATCH 6/7] Updates masp gas costs

---
 crates/gas/src/lib.rs | 23 +++++++++++------------
 1 file changed, 11 insertions(+), 12 deletions(-)

diff --git a/crates/gas/src/lib.rs b/crates/gas/src/lib.rs
index 88c73b7d8d..89e24209ff 100644
--- a/crates/gas/src/lib.rs
+++ b/crates/gas/src/lib.rs
@@ -84,29 +84,28 @@ pub const WASM_MEMORY_PAGE_GAS: u32 =
 pub const IBC_ACTION_VALIDATE_GAS: u64 = 1_472_023;
 /// The cost to execute an Ibc action
 pub const IBC_ACTION_EXECUTE_GAS: u64 = 3_678_745;
-// FIXME: update costs
 /// The cost of masp sig verification
-pub const MASP_VERIFY_SIG_GAS: u64 = 1;
+pub const MASP_VERIFY_SIG_GAS: u64 = 5_443_000;
 /// The fixed cost of spend note verification
-pub const MASP_FIXED_SPEND_GAS: u64 = 1;
+pub const MASP_FIXED_SPEND_GAS: u64 = 87_866_000;
 /// The variable cost of spend note verification
-pub const MASP_VARIABLE_SPEND_GAS: u64 = 1;
+pub const MASP_VARIABLE_SPEND_GAS: u64 = 14_384_000;
 /// The fixed cost of convert note verification
-pub const MASP_FIXED_CONVERT_GAS: u64 = 1;
+pub const MASP_FIXED_CONVERT_GAS: u64 = 70_308_000;
 /// The variable cost of convert note verification
-pub const MASP_VARIABLE_CONVERT_GAS: u64 = 1;
+pub const MASP_VARIABLE_CONVERT_GAS: u64 = 12_664_000;
 /// The fixed cost of output note verification
-pub const MASP_FIXED_OUTPUT_GAS: u64 = 1;
+pub const MASP_FIXED_OUTPUT_GAS: u64 = 78_203_000;
 /// The variable cost of output note verification
-pub const MASP_VARIABLE_OUTPUT_GAS: u64 = 1;
+pub const MASP_VARIABLE_OUTPUT_GAS: u64 = 14_586_000;
 /// The cost to process a masp spend note in the bundle
-pub const MASP_SPEND_CHECK_GAS: u64 = 1;
+pub const MASP_SPEND_CHECK_GAS: u64 = 479_730;
 /// The cost to process a masp convert note in the bundle
-pub const MASP_CONVERT_CHECK_GAS: u64 = 1;
+pub const MASP_CONVERT_CHECK_GAS: u64 = 173_570;
 /// The cost to process a masp output note in the bundle
-pub const MASP_OUTPUT_CHECK_GAS: u64 = 1;
+pub const MASP_OUTPUT_CHECK_GAS: u64 = 310_260;
 /// The cost to run the final masp check in the bundle
-pub const MASP_FINAL_CHECK_GAS: u64 = 1;
+pub const MASP_FINAL_CHECK_GAS: u64 = 44;
 /// Gas divider specific for the masp vp. Only allocates half of the cores to
 /// the masp vp since we can expect the other half to be busy with other vps
 pub const MASP_PARALLEL_GAS_DIVIDER: u64 = PARALLEL_GAS_DIVIDER / 2;

From 15affec2d5e8f14ec420a1cfd9891868369ab710 Mon Sep 17 00:00:00 2001
From: Marco Granelli <marco@heliax.dev>
Date: Wed, 22 May 2024 14:50:01 +0200
Subject: [PATCH 7/7] Changelog #2972

---
 .../unreleased/improvements/2972-masp-parallel-verification.md  | 2 ++
 1 file changed, 2 insertions(+)
 create mode 100644 .changelog/unreleased/improvements/2972-masp-parallel-verification.md

diff --git a/.changelog/unreleased/improvements/2972-masp-parallel-verification.md b/.changelog/unreleased/improvements/2972-masp-parallel-verification.md
new file mode 100644
index 0000000000..4aa10651bb
--- /dev/null
+++ b/.changelog/unreleased/improvements/2972-masp-parallel-verification.md
@@ -0,0 +1,2 @@
+- Improved masp vp verification to run in parallel.
+  ([\#2972](https://github.com/anoma/namada/pull/2972))
\ No newline at end of file