fix: reduce cost of signer key-search algorithm with pre-compute table

CHANGELOG.md

@@ -17,6 +17,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 ### Removed
 
 ### Fixed
+- [\#129](https://github.com/Manta-Network/manta-rs/pull/129) Reduce cost of signer key-search algorithm by adding dynamic pre-computation table
 
 ### Security
 

manta-accounting/src/key.rs

@@ -21,8 +21,6 @@
 //!
 //! [`BIP-0044`]: https://github.com/bitcoin/bips/blob/master/bip-0044.mediawiki
 
-// TODO: Build custom iterator types for [`keypairs`] and [`generate_keys`].
-
 use alloc::vec::Vec;
 use core::{
     cmp,
@@ -35,6 +33,7 @@ use manta_crypto::{
     key::KeyDerivationFunction,
     rand::{RngCore, Sample},
 };
+use manta_util::collections::btree_map::{self, BTreeMap};
 
 #[cfg(feature = "serde")]
 use manta_util::serde::{Deserialize, Serialize};
@@ -524,6 +523,70 @@ where
         })
     }
 
+    /// Returns a new [`ViewKeyTable`] for `self`.
+    #[inline]
+    pub fn view_key_table(self) -> ViewKeyTable<'h, H> {
+        ViewKeyTable::new(self)
+    }
+}
+
+/// View Key Table
+pub struct ViewKeyTable<'h, H>
+where
+    H: HierarchicalKeyDerivationScheme + ?Sized,
+{
+    /// Account Keys
+    keys: AccountKeysMut<'h, H>,
+
+    /// Pre-computed View Keys
+    view_keys: BTreeMap<KeyIndex, H::SecretKey>,
+}
+
+impl<'h, H> ViewKeyTable<'h, H>
+where
+    H: HierarchicalKeyDerivationScheme + ?Sized,
+{
+    /// View Key Buffer Maximum Size Limit
+    pub const VIEW_KEY_BUFFER_LIMIT: usize = 16 * (H::GAP_LIMIT as usize);
+
+    /// Builds a new [`ViewKeyTable`] over the account `keys`.
+    #[inline]
+    pub fn new(keys: AccountKeysMut<'h, H>) -> Self {
+        Self {
+            keys,
+            view_keys: Default::default(),
+        }
+    }
+
+    /// Returns the account keys associated to `self`.
+    #[inline]
+    pub fn into_keys(self) -> AccountKeysMut<'h, H> {
+        self.keys
+    }
+
+    /// Returns the view key for this account at `index`, if it does not exceed the maximum index.
+    ///
+    /// # Limits
+    ///
+    /// This function uses a view key buffer that stores the computed keys to reduce the number of
+    /// times a re-compute of the view keys is needed while searching. The buffer only grows past
+    /// the current key bounds with a call to [`find_index_with_gap`](Self::find_index_with_gap)
+    /// which extends the buffer by at most [`GAP_LIMIT`]-many keys per round. To prevent allocating
+    /// too much memory, the internal buffer is capped at [`VIEW_KEY_BUFFER_LIMIT`]-many elements.
+    ///
+    /// [`GAP_LIMIT`]: HierarchicalKeyDerivationScheme::GAP_LIMIT
+    /// [`VIEW_KEY_BUFFER_LIMIT`]: Self::VIEW_KEY_BUFFER_LIMIT
+    #[inline]
+    pub fn view_key(&mut self, index: KeyIndex) -> Option<&H::SecretKey> {
+        btree_map::get_or_mutate(&mut self.view_keys, &index, |map| {
+            let next_key = self.keys.view_key(index)?;
+            if map.len() == Self::VIEW_KEY_BUFFER_LIMIT {
+                btree_map::pop_last(map);
+            }
+            Some(btree_map::insert_then_get(map, index, next_key))
+        })
+    }
+
     /// Applies `f` to the view keys generated by `self` returning the first non-`None` result with
     /// it's key index and key attached, or returns `None` if every application of `f` returned
     /// `None`.
@@ -534,12 +597,12 @@ where
     {
         let mut index = KeyIndex::default();
         loop {
-            let view_key = self.view_key(index)?;
-            if let Some(item) = f(&view_key) {
-                self.account.last_used_index = cmp::max(self.account.last_used_index, index);
+            if let Some(item) = f(self.view_key(index)?) {
+                self.keys.account.last_used_index =
+                    cmp::max(self.keys.account.last_used_index, index);
                 return Some(ViewKeySelection {
                     index,
-                    keypair: SecretKeyPair::new(self.derive_spend(index), view_key),
+                    keypair: self.keys.derive_pair(index),
                     item,
                 });
             }
@@ -563,15 +626,15 @@ where
     where
         F: FnMut(&H::SecretKey) -> Option<T>,
     {
-        let previous_maximum = self.account.maximum_index;
-        self.account.maximum_index.index += H::GAP_LIMIT;
+        let previous_maximum = self.keys.account.maximum_index;
+        self.keys.account.maximum_index.index += H::GAP_LIMIT;
         match self.find_index(f) {
             Some(result) => {
-                self.account.maximum_index = cmp::max(previous_maximum, result.index);
+                self.keys.account.maximum_index = cmp::max(previous_maximum, result.index);
                 Some(result)
             }
             _ => {
-                self.account.maximum_index = previous_maximum;
+                self.keys.account.maximum_index = previous_maximum;
                 None
             }
         }
@@ -612,6 +675,25 @@ where
     pub item: T,
 }
 
+impl<H, T> ViewKeySelection<H, T>
+where
+    H: HierarchicalKeyDerivationScheme + ?Sized,
+{
+    /// Computes `f` on `self.item` returning a new [`ViewKeySelection`] with the same `index` and
+    /// `keypair`.
+    #[inline]
+    pub fn map<U, F>(self, f: F) -> ViewKeySelection<H, U>
+    where
+        F: FnOnce(T) -> U,
+    {
+        ViewKeySelection {
+            index: self.index,
+            keypair: self.keypair,
+            item: f(self.item),
+        }
+    }
+}
+
 /// Account
 #[cfg_attr(
     feature = "serde",

manta-accounting/src/wallet/signer.rs

@@ -28,7 +28,10 @@
 
 use crate::{
     asset::{Asset, AssetId, AssetMap, AssetMetadata, AssetValue},
-    key::{self, HierarchicalKeyDerivationScheme, KeyIndex, SecretKeyPair, ViewKeySelection},
+    key::{
+        self, HierarchicalKeyDerivationScheme, KeyIndex, SecretKeyPair, ViewKeySelection,
+        ViewKeyTable,
+    },
     transfer::{
         self,
         batch::Join,
@@ -630,63 +633,69 @@ where
         )
     }
 
-    /// Inserts the new `utxo`-`encrypted_note` pair if a known key can decrypt the note and
-    /// validate the utxo.
+    /// Finds the next viewing key that can decrypt the `encrypted_note` from the `view_key_table`.
     #[inline]
-    fn insert_next_item(
-        &mut self,
+    fn find_next_key<'h>(
+        view_key_table: &mut ViewKeyTable<'h, C::HierarchicalKeyDerivationScheme>,
         parameters: &Parameters<C>,
         with_recovery: bool,
-        utxo: Utxo<C>,
         encrypted_note: EncryptedNote<C>,
+    ) -> Option<ViewKeySelection<C::HierarchicalKeyDerivationScheme, Note<C>>> {
+        let mut finder = DecryptedMessage::find(encrypted_note);
+        view_key_table
+            .find_index_with_maybe_gap(with_recovery, move |k| {
+                finder.decrypt(&parameters.note_encryption_scheme, k)
+            })
+            .map(|selection| selection.map(|item| item.plaintext))
+    }
+
+    /// Inserts the new `utxo`-`note` pair into the `utxo_accumulator` adding the spendable amount
+    /// to `assets` if there is no void number to match it.
+    #[inline]
+    fn insert_next_item(
+        utxo_accumulator: &mut C::UtxoAccumulator,
+        assets: &mut C::AssetMap,
+        parameters: &Parameters<C>,
+        utxo: Utxo<C>,
+        selection: ViewKeySelection<C::HierarchicalKeyDerivationScheme, Note<C>>,
         void_numbers: &mut Vec<VoidNumber<C>>,
         deposit: &mut Vec<Asset>,
-    ) -> Result<(), SyncError<C::Checkpoint>> {
-        let mut finder = DecryptedMessage::find(encrypted_note);
-        if let Some(ViewKeySelection {
+    ) {
+        let ViewKeySelection {
             index,
             keypair,
-            item,
-        }) = self
-            .accounts
-            .get_mut_default()
-            .find_index_with_maybe_gap(with_recovery, |k| {
-                finder.decrypt(&parameters.note_encryption_scheme, k)
-            })
-        {
-            let Note {
+            item: Note {
                 ephemeral_secret_key,
                 asset,
-            } = item.plaintext;
-            if let Some(void_number) =
-                parameters.check_full_asset(&keypair.spend, &ephemeral_secret_key, &asset, &utxo)
-            {
-                if let Some(index) = void_numbers.iter().position(move |v| v == &void_number) {
-                    void_numbers.remove(index);
-                } else {
-                    self.utxo_accumulator.insert(&utxo);
-                    self.assets.insert((index, ephemeral_secret_key), asset);
-                    if !asset.is_zero() {
-                        deposit.push(asset);
-                    }
-                    return Ok(());
+            },
+        } = selection;
+        if let Some(void_number) =
+            parameters.check_full_asset(&keypair.spend, &ephemeral_secret_key, &asset, &utxo)
+        {
+            if let Some(index) = void_numbers.iter().position(move |v| v == &void_number) {
+                void_numbers.remove(index);
+            } else {
+                utxo_accumulator.insert(&utxo);
+                assets.insert((index, ephemeral_secret_key), asset);
+                if !asset.is_zero() {
+                    deposit.push(asset);
                 }
+                return;
             }
         }
-        self.utxo_accumulator.insert_nonprovable(&utxo);
-        Ok(())
+        utxo_accumulator.insert_nonprovable(&utxo);
     }
 
     /// Checks if `asset` matches with `void_number`, removing it from the `utxo_accumulator` and
     /// inserting it into the `withdraw` set if this is the case.
     #[inline]
     fn is_asset_unspent(
+        utxo_accumulator: &mut C::UtxoAccumulator,
         parameters: &Parameters<C>,
         secret_spend_key: &SecretKey<C>,
         ephemeral_secret_key: &SecretKey<C>,
         asset: Asset,
         void_numbers: &mut Vec<VoidNumber<C>>,
-        utxo_accumulator: &mut C::UtxoAccumulator,
         withdraw: &mut Vec<Asset>,
     ) -> bool {
         let utxo = parameters.utxo(
@@ -716,33 +725,44 @@ where
         inserts: I,
         mut void_numbers: Vec<VoidNumber<C>>,
         is_partial: bool,
-    ) -> Result<SyncResponse<C::Checkpoint>, SyncError<C::Checkpoint>>
+    ) -> SyncResponse<C::Checkpoint>
     where
         I: Iterator<Item = (Utxo<C>, EncryptedNote<C>)>,
     {
         let void_number_count = void_numbers.len();
         let mut deposit = Vec::new();
         let mut withdraw = Vec::new();
+        let mut view_key_table = self.accounts.get_mut_default().view_key_table();
         for (utxo, encrypted_note) in inserts {
-            self.insert_next_item(
+            if let Some(selection) = Self::find_next_key(
+                &mut view_key_table,
                 parameters,
                 with_recovery,
-                utxo,
                 encrypted_note,
-                &mut void_numbers,
-                &mut deposit,
-            )?;
+            ) {
+                Self::insert_next_item(
+                    &mut self.utxo_accumulator,
+                    &mut self.assets,
+                    parameters,
+                    utxo,
+                    selection,
+                    &mut void_numbers,
+                    &mut deposit,
+                );
+            } else {
+                self.utxo_accumulator.insert_nonprovable(&utxo);
+            }
         }
         self.assets.retain(|(index, ephemeral_secret_key), assets| {
             assets.retain(
                 |asset| match self.accounts.get_default().spend_key(*index) {
                     Some(secret_spend_key) => Self::is_asset_unspent(
+                        &mut self.utxo_accumulator,
                         parameters,
                         &secret_spend_key,
                         ephemeral_secret_key,
                         *asset,
                         &mut void_numbers,
-                        &mut self.utxo_accumulator,
                         &mut withdraw,
                     ),
                     _ => true,
@@ -753,7 +773,7 @@ where
         self.checkpoint.update_from_void_numbers(void_number_count);
         self.checkpoint
             .update_from_utxo_accumulator(&self.utxo_accumulator);
-        Ok(SyncResponse {
+        SyncResponse {
             checkpoint: self.checkpoint.clone(),
             balance_update: if is_partial {
                 // TODO: Whenever we are doing a full update, don't even build the `deposit` and
@@ -764,7 +784,7 @@ where
                     assets: self.assets.assets().into(),
                 }
             },
-        })
+        }
     }
 
     /// Builds the pre-sender associated to `key` and `asset`.
@@ -1144,15 +1164,15 @@ where
         } else {
             let has_pruned = request.prune(checkpoint);
             let SyncData { receivers, senders } = request.data;
-            let result = self.state.sync_with(
+            let response = self.state.sync_with(
                 &self.parameters.parameters,
                 request.with_recovery,
                 receivers.into_iter(),
                 senders,
                 !has_pruned,
             );
             self.state.utxo_accumulator.commit();
-            result
+            Ok(response)
         }
     }
 

manta-util/src/array.rs

@@ -30,7 +30,7 @@ use alloc::{boxed::Box, vec::Vec};
 #[cfg(feature = "serde-array")]
 use crate::serde::{Deserialize, Serialize};
 
-/// Error Message for the [`into_array_unchecked`] and [`into_boxed_array_unchecked`] messages.
+/// Error Message for the [`into_array_unchecked`] and [`into_boxed_array_unchecked`] Functions
 const INTO_UNCHECKED_ERROR_MESSAGE: &str =
     "Input did not have the correct length to match the output array of length";