Merge branch 'tomas/pregen-masp-proofs' (#1768)

* origin/tomas/pregen-masp-proofs:
  changelog: add #1768
  make: add recipes for integration tests with saved MASP proofs
  test: add masp_proofs test fixtures
  shared/masp: allow to save and load proofs for tests

.changelog/unreleased/testing/1768-pregen-masp-proofs.md

@@ -0,0 +1,2 @@
+- Added pre-built MASP proofs for integration tests.
+  ([\#1768](https://github.com/anoma/namada/pull/1768))

Cargo.toml

@@ -20,6 +20,7 @@ members = [
 exclude = [
   "wasm",
   "wasm_for_tests",
+  "test_fixtures",
 ]
 
 [workspace.package]

Makefile

@@ -4,6 +4,7 @@ package = namada
 NAMADA_E2E_USE_PREBUILT_BINARIES ?= true
 NAMADA_E2E_DEBUG ?= true
 RUST_BACKTRACE ?= 1
+NAMADA_MASP_TEST_SEED ?= 0
 
 cargo := $(env) cargo
 rustup := $(env) rustup
@@ -116,12 +117,14 @@ audit:
 
 test: test-unit test-e2e test-wasm
 
-# Unit tests with coverage report
 test-coverage:
+	# Run integration tests with pre-built MASP proofs
+	NAMADA_MASP_TEST_SEED=$(NAMADA_MASP_TEST_SEED) \
+	NAMADA_MASP_TEST_PROOFS=load \
 	$(cargo) +$(nightly) llvm-cov --output-dir target \
 		--features namada/testing \
 		--html \
-		-- --skip e2e --skip integration -Z unstable-options --report-time
+		-- --skip e2e -Z unstable-options --report-time
 
 # NOTE: `TEST_FILTER` is prepended with `e2e::`. Since filters in `cargo test`
 # work with a substring search, TEST_FILTER only works if it contains a string
@@ -137,7 +140,23 @@ test-e2e:
 	--test-threads=1 \
 	-Z unstable-options --report-time
 
+# Run integration tests with pre-built MASP proofs
 test-integration:
+	NAMADA_MASP_TEST_SEED=$(NAMADA_MASP_TEST_SEED) \
+	NAMADA_MASP_TEST_PROOFS=load \
+	make test-integration-slow
+
+# Clear pre-built proofs, run integration tests and save the new proofs
+test-integration-save-proofs:
+    # Clear old proofs first
+	rm --force test_fixtures/masp_proofs/*.bin || true
+	NAMADA_MASP_TEST_SEED=$(NAMADA_MASP_TEST_SEED) \
+	NAMADA_MASP_TEST_PROOFS=save \
+	TEST_FILTER=masp \
+	make test-integration-slow
+
+# Run integration tests without specifiying any pre-built MASP proofs option
+test-integration-slow:
 	RUST_BACKTRACE=$(RUST_BACKTRACE) \
 	$(cargo) +$(nightly) test integration::$(TEST_FILTER) \
 	-Z unstable-options \

shared/src/ledger/masp.rs

@@ -1,7 +1,6 @@
 //! MASP verification wrappers.
 
-use std::collections::hash_map::Entry;
-use std::collections::{BTreeMap, HashMap, HashSet};
+use std::collections::{btree_map, BTreeMap, BTreeSet, HashMap, HashSet};
 use std::env;
 use std::fmt::Debug;
 #[cfg(feature = "masp-tx-gen")]
@@ -79,6 +78,17 @@ use crate::types::transaction::{EllipticCurve, PairingEngine, WrapperTx};
 /// the default OS specific path is used.
 pub const ENV_VAR_MASP_PARAMS_DIR: &str = "NAMADA_MASP_PARAMS_DIR";
 
+/// Env var to either "save" proofs into files or to "load" them from
+/// files.
+pub const ENV_VAR_MASP_TEST_PROOFS: &str = "NAMADA_MASP_TEST_PROOFS";
+
+/// Randomness seed for MASP integration tests to build proofs with
+/// deterministic rng.
+pub const ENV_VAR_MASP_TEST_SEED: &str = "NAMADA_MASP_TEST_SEED";
+
+/// A directory to save serialized proofs for tests.
+pub const MASP_TEST_PROOFS_DIR: &str = "test_fixtures/masp_proofs";
+
 /// The network to use for MASP
 #[cfg(feature = "mainnet")]
 const NETWORK: MainNetwork = MainNetwork;
@@ -93,6 +103,26 @@ pub const OUTPUT_NAME: &str = "masp-output.params";
 /// Convert circuit name
 pub const CONVERT_NAME: &str = "masp-convert.params";
 
+/// Shielded transfer
+#[derive(Clone, Debug, BorshSerialize, BorshDeserialize)]
+pub struct ShieldedTransfer {
+    /// Shielded transfer builder
+    pub builder: Builder<(), (), ExtendedFullViewingKey, ()>,
+    /// MASP transaction
+    pub masp_tx: Transaction,
+    /// Metadata
+    pub metadata: SaplingMetadata,
+    /// Epoch in which the transaction was created
+    pub epoch: Epoch,
+}
+
+#[derive(Clone, Copy, Debug)]
+enum LoadOrSaveProofs {
+    Load,
+    Save,
+    Neither,
+}
+
 fn load_pvks() -> (
     PreparedVerifyingKey<Bls12>,
     PreparedVerifyingKey<Bls12>,
@@ -511,7 +541,7 @@ impl From<MaspAmount> for Amount {
 
 /// Represents the amount used of different conversions
 pub type Conversions =
-    HashMap<AssetType, (AllowedConversion, MerklePath<Node>, i128)>;
+    BTreeMap<AssetType, (AllowedConversion, MerklePath<Node>, i128)>;
 
 /// Represents the changes that were made to a list of transparent accounts
 pub type TransferDelta = HashMap<Address, MaspChange>;
@@ -531,7 +561,7 @@ pub struct ShieldedContext<U: ShieldedUtils> {
     /// The commitment tree produced by scanning all transactions up to tx_pos
     pub tree: CommitmentTree<Node>,
     /// Maps viewing keys to applicable note positions
-    pub pos_map: HashMap<ViewingKey, HashSet<usize>>,
+    pub pos_map: HashMap<ViewingKey, BTreeSet<usize>>,
     /// Maps a nullifier to the note position to which it applies
     pub nf_map: HashMap<Nullifier, usize>,
     /// Maps note positions to their corresponding notes
@@ -657,7 +687,7 @@ impl<U: ShieldedUtils> ShieldedContext<U> {
                 ..Default::default()
             };
             for vk in unknown_keys {
-                tx_ctx.pos_map.entry(vk).or_insert_with(HashSet::new);
+                tx_ctx.pos_map.entry(vk).or_insert_with(BTreeSet::new);
             }
             // Update this unknown shielded context until it is level with self
             while tx_ctx.last_txidx != self.last_txidx {
@@ -931,7 +961,9 @@ impl<U: ShieldedUtils> ShieldedContext<U> {
         asset_type: AssetType,
         conversions: &'a mut Conversions,
     ) {
-        if let Entry::Vacant(conv_entry) = conversions.entry(asset_type) {
+        if let btree_map::Entry::Vacant(conv_entry) =
+            conversions.entry(asset_type)
+        {
             // Query for the ID of the last accepted transaction
             if let Some((addr, denom, ep, conv, path)) =
                 query_conversion(client, asset_type).await
@@ -962,7 +994,7 @@ impl<U: ShieldedUtils> ShieldedContext<U> {
                     client,
                     balance,
                     target_epoch,
-                    HashMap::new(),
+                    BTreeMap::new(),
                 )
                 .await
                 .0;
@@ -1142,7 +1174,7 @@ impl<U: ShieldedUtils> ShieldedContext<U> {
         Conversions,
     ) {
         // Establish connection with which to do exchange rate queries
-        let mut conversions = HashMap::new();
+        let mut conversions = BTreeMap::new();
         let mut val_acc = Amount::zero();
         let mut notes = Vec::new();
         // Retrieve the notes that can be spent by this key
@@ -1286,7 +1318,7 @@ impl<U: ShieldedUtils> ShieldedContext<U> {
         println!("Decoded pinned balance: {:?}", amount);
         // Finally, exchange the balance to the transaction's epoch
         let computed_amount = self
-            .compute_exchanged_amount(client, amount, ep, HashMap::new())
+            .compute_exchanged_amount(client, amount, ep, BTreeMap::new())
             .await
             .0;
         println!("Exchanged amount: {:?}", computed_amount);
@@ -1357,16 +1389,17 @@ impl<U: ShieldedUtils> ShieldedContext<U> {
         args: &args::TxTransfer,
         shielded_gas: bool,
     ) -> Result<
-        Option<(
-            Builder<(), (), ExtendedFullViewingKey, ()>,
-            Transaction,
-            SaplingMetadata,
-            Epoch,
-        )>,
+        Option<ShieldedTransfer>,
         builder::Error<std::convert::Infallible>,
     > {
         // No shielded components are needed when neither source nor destination
         // are shielded
+
+        use std::str::FromStr;
+
+        use rand::rngs::StdRng;
+        use rand_core::SeedableRng;
+
         let spending_key = args.source.spending_key();
         let payment_address = args.target.payment_address();
         // No shielded components are needed when neither source nor
@@ -1388,8 +1421,27 @@ impl<U: ShieldedUtils> ShieldedContext<U> {
         // possesion
         let memo = MemoBytes::empty();
 
+        // Try to get a seed from env var, if any.
+        let rng = if let Ok(seed) =
+            env::var(ENV_VAR_MASP_TEST_SEED).map(|seed| {
+                let exp_str =
+                    format!("Env var {ENV_VAR_MASP_TEST_SEED} must be a u64.");
+                let parsed_seed: u64 =
+                    FromStr::from_str(&seed).expect(&exp_str);
+                parsed_seed
+            }) {
+            tracing::warn!(
+                "UNSAFE: Using a seed from {ENV_VAR_MASP_TEST_SEED} env var \
+                 to build proofs."
+            );
+            StdRng::seed_from_u64(seed)
+        } else {
+            StdRng::from_rng(OsRng).unwrap()
+        };
+
         // Now we build up the transaction within this object
-        let mut builder = Builder::<TestNetwork, OsRng>::new(NETWORK, 1.into());
+        let mut builder =
+            Builder::<TestNetwork, _>::new_with_rng(NETWORK, 1.into(), rng);
 
         // break up a transfer into a number of transfers with suitable
         // denominations
@@ -1548,16 +1600,81 @@ impl<U: ShieldedUtils> ShieldedContext<U> {
             }
         }
 
-        // Build and return the constructed transaction
-        builder
-            .clone()
-            .build(
+        // To speed up integration tests, we can save and load proofs
+        let load_or_save = if let Ok(masp_proofs) =
+            env::var(ENV_VAR_MASP_TEST_PROOFS)
+        {
+            let parsed = match masp_proofs.to_ascii_lowercase().as_str() {
+                "load" => LoadOrSaveProofs::Load,
+                "save" => LoadOrSaveProofs::Save,
+                env_var => panic!(
+                    "Unexpected value for {ENV_VAR_MASP_TEST_PROOFS} env var. \
+                     Expecting \"save\" or \"load\", but got \"{env_var}\"."
+                ),
+            };
+            if env::var(ENV_VAR_MASP_TEST_SEED).is_err() {
+                panic!(
+                    "Ensure to set a seed with {ENV_VAR_MASP_TEST_SEED} env \
+                     var when using {ENV_VAR_MASP_TEST_PROOFS} for \
+                     deterministic proofs."
+                );
+            }
+            parsed
+        } else {
+            LoadOrSaveProofs::Neither
+        };
+
+        let builder_clone = builder.clone().map_builder(WalletMap);
+        let builder_bytes = BorshSerialize::try_to_vec(&builder_clone).unwrap();
+        let builder_hash =
+            namada_core::types::hash::Hash::sha256(&builder_bytes);
+        let saved_filepath = env::current_dir()
+            .unwrap()
+            // One up from "tests" dir to the root dir
+            .parent()
+            .unwrap()
+            .join(MASP_TEST_PROOFS_DIR)
+            .join(format!("{builder_hash}.bin"));
+
+        if let LoadOrSaveProofs::Load = load_or_save {
+            let recommendation = format!(
+                "Re-run the tests with {ENV_VAR_MASP_TEST_PROOFS}=save to \
+                 re-generate proofs."
+            );
+            let exp_str = format!(
+                "Read saved MASP proofs from {}. {recommendation}",
+                saved_filepath.to_string_lossy()
+            );
+            let loaded_bytes =
+                tokio::fs::read(&saved_filepath).await.expect(&exp_str);
+            let exp_str = format!(
+                "Valid `ShieldedTransfer` bytes in {}. {recommendation}",
+                saved_filepath.to_string_lossy()
+            );
+            let loaded: ShieldedTransfer =
+                BorshDeserialize::try_from_slice(&loaded_bytes)
+                    .expect(&exp_str);
+            Ok(Some(loaded))
+        } else {
+            // Build and return the constructed transaction
+            let (masp_tx, metadata) = builder.build(
                 &self.utils.local_tx_prover(),
                 &FeeRule::non_standard(tx_fee),
-            )
-            .map(|(tx, metadata)| {
-                Some((builder.map_builder(WalletMap), tx, metadata, epoch))
-            })
+            )?;
+            let built = ShieldedTransfer {
+                builder: builder_clone,
+                masp_tx,
+                metadata,
+                epoch,
+            };
+            if let LoadOrSaveProofs::Save = load_or_save {
+                let built_bytes = BorshSerialize::try_to_vec(&built).unwrap();
+                tokio::fs::write(&saved_filepath, built_bytes)
+                    .await
+                    .unwrap();
+            }
+            Ok(Some(built))
+        }
     }
 
     /// Obtain the known effects of all accepted shielded and transparent

shared/src/ledger/tx.rs

@@ -42,7 +42,7 @@ use crate::ibc::tx_msg::Msg;
 use crate::ibc::Height as IbcHeight;
 use crate::ibc_proto::cosmos::base::v1beta1::Coin;
 use crate::ledger::args::{self, InputAmount};
-use crate::ledger::masp::{ShieldedContext, ShieldedUtils};
+use crate::ledger::masp::{ShieldedContext, ShieldedTransfer, ShieldedUtils};
 use crate::ledger::rpc::{
     self, format_denominated_amount, validate_amount, TxBroadcastData,
     TxResponse,
@@ -1607,30 +1607,34 @@ pub async fn build_transfer<
     let mut tx = Tx::new(chain_id, args.tx.expiration);
 
     // Add the MASP Transaction and its Builder to facilitate validation
-    let (masp_hash, shielded_tx_epoch) = if let Some(shielded_parts) =
-        shielded_parts
+    let (masp_hash, shielded_tx_epoch) = if let Some(ShieldedTransfer {
+        builder,
+        masp_tx,
+        metadata,
+        epoch,
+    }) = shielded_parts
     {
         // Add a MASP Transaction section to the Tx and get the tx hash
-        let masp_tx_hash = tx.add_masp_tx_section(shielded_parts.1).1;
+        let masp_tx_hash = tx.add_masp_tx_section(masp_tx).1;
 
         // Get the decoded asset types used in the transaction to give
         // offline wallet users more information
-        let asset_types = used_asset_types(shielded, client, &shielded_parts.0)
+        let asset_types = used_asset_types(shielded, client, &builder)
             .await
             .unwrap_or_default();
 
         tx.add_masp_builder(MaspBuilder {
             asset_types,
             // Store how the Info objects map to Descriptors/Outputs
-            metadata: shielded_parts.2,
+            metadata,
             // Store the data that was used to construct the Transaction
-            builder: shielded_parts.0,
+            builder,
             // Link the Builder to the Transaction by hash code
             target: masp_tx_hash,
         });
 
         // The MASP Transaction section hash will be used in Transfer
-        (Some(masp_tx_hash), Some(shielded_parts.3))
+        (Some(masp_tx_hash), Some(epoch))
     } else {
         (None, None)
     };

test_fixtures/masp_proofs/README.md

@@ -0,0 +1,11 @@
+# MASP proofs for tests
+
+This directory contains pre-built MASP transaction proofs used to speed-up integration tests.
+
+```shell
+# Run the tests with the saved proofs from here.
+make test-integration
+
+# Delete old proofs, run the tests and save the new proofs.
+make test-integration-save-proofs
+```