Add WalletWrite::insert_address_with_diversifier_index function

This enables use cases where the diversifier index is prescribed instead of sequentially assigned.

zcash_client_backend/CHANGELOG.md

@@ -21,6 +21,7 @@ and this library adheres to Rust's notion of
 - `zcash_client_backend::fees::ChangeValue::orchard`
 - `zcash_client_backend::wallet`:
   - `Note::Orchard`
+- `WalletWrite::insert_address_with_diversifier_index`
 
 ### Changed
 - `zcash_client_backend::data_api`:

zcash_client_backend/src/data_api.rs

@@ -22,6 +22,7 @@ use zcash_primitives::{
     },
     zip32::{AccountId, Scope},
 };
+use zip32::DiversifierIndex;
 
 use crate::{
     address::{AddressMetadata, UnifiedAddress},
@@ -1019,6 +1020,27 @@ pub trait WalletWrite: WalletRead {
         request: UnifiedAddressRequest,
     ) -> Result<Option<UnifiedAddress>, Self::Error>;
 
+    /// Generates and persists a new address for the specified account, with the specified
+    /// diversifier index.
+    ///
+    /// Returns the new address, or an error if the account identifier does not correspond to a
+    /// known account.
+    /// A conflict with an existing row in the database is considered acceptable and no error is returned.
+    /// If the diversifier index cannot produce a valid Sapling address, no sapling receiver will
+    /// be included in the returned address.
+    /// If the diversifier is outside the range for transparent addresses, no transparent receiver
+    /// will be included in the returned address.
+    ///
+    /// This supports a more advanced use case than `get_next_available_address` where the caller
+    /// simply gets the next diversified address sequentially. Mixing use of the two functions
+    /// is not recommended because `get_next_available_address` will return the next address
+    /// after the highest diversifier index in the database, which may leave gaps in the sequence.
+    fn insert_address_with_diversifier_index(
+        &mut self,
+        account: AccountId,
+        diversifier_index: DiversifierIndex,
+    ) -> Result<UnifiedAddress, Self::Error>;
+
     /// Updates the state of the wallet database by persisting the provided block information,
     /// along with the note commitments that were detected when scanning the block for transactions
     /// pertaining to this wallet.
@@ -1105,6 +1127,7 @@ pub mod testing {
     use secrecy::{ExposeSecret, SecretVec};
     use shardtree::{error::ShardTreeError, store::memory::MemoryShardStore, ShardTree};
     use std::{collections::HashMap, convert::Infallible, num::NonZeroU32};
+    use zip32::DiversifierIndex;
 
     use zcash_primitives::{
         block::BlockHash,
@@ -1323,6 +1346,14 @@ pub mod testing {
             Ok(None)
         }
 
+        fn insert_address_with_diversifier_index(
+            &mut self,
+            _account: AccountId,
+            _diversifier_index: DiversifierIndex,
+        ) -> Result<UnifiedAddress, Self::Error> {
+            todo!()
+        }
+
         #[allow(clippy::type_complexity)]
         fn put_blocks(
             &mut self,

zcash_client_sqlite/Cargo.toml

@@ -56,6 +56,7 @@ shardtree = { workspace = true, features = ["legacy-api"] }
 # CocoaPods, due to being bound to React Native. We need to ensure that the SQLite
 # version required for `rusqlite` is a version that is available through CocoaPods.
 rusqlite = { version = "0.29.0", features = ["bundled", "time", "array"] }
+libsqlite3-sys = { version = "0.26.0", features = ["bundled"] }
 schemer = "0.2"
 schemer-rusqlite = "0.2.2"
 time = "0.3.22"

zcash_client_sqlite/src/lib.rs

@@ -42,6 +42,7 @@ use std::{
     borrow::Borrow, collections::HashMap, convert::AsRef, fmt, num::NonZeroU32, ops::Range,
     path::Path,
 };
+use zcash_keys::keys::AddressGenerationError;
 
 use incrementalmerkletree::Position;
 use shardtree::{error::ShardTreeError, ShardTree};
@@ -445,6 +446,63 @@ impl<P: consensus::Parameters> WalletWrite for WalletDb<rusqlite::Connection, P>
         )
     }
 
+    fn insert_address_with_diversifier_index(
+        &mut self,
+        account: AccountId,
+        diversifier_index: DiversifierIndex,
+    ) -> Result<UnifiedAddress, SqliteClientError> {
+        self.transactionally(|wdb| {
+            let keys = wdb.get_unified_full_viewing_keys()?;
+            let ufvk = keys
+                .get(&account)
+                .ok_or(SqliteClientError::AccountUnknown(account))?;
+
+            let has_orchard = true;
+            let mut has_sapling = true;
+            let mut has_transparent = true;
+
+            // Get the most comprehensive UA available for the given diversifier index.
+            // We may have to drop the sapling and/or the transparent receiver if the diversifier index is invalid or out of range.
+            let addr = loop {
+                if let Some(addr) = match ufvk.address(
+                    diversifier_index,
+                    UnifiedAddressRequest::unsafe_new(has_orchard, has_sapling, has_transparent),
+                ) {
+                    Ok(addr) => Some(addr),
+                    Err(AddressGenerationError::InvalidSaplingDiversifierIndex(_)) => {
+                        has_sapling = false;
+                        None
+                    }
+                    Err(AddressGenerationError::InvalidTransparentChildIndex(_)) => {
+                        has_transparent = false;
+                        None
+                    }
+                    Err(_) => return Err(SqliteClientError::DiversifierIndexOutOfRange),
+                } {
+                    break addr;
+                }
+            };
+
+            return match wallet::insert_address(
+                wdb.conn.0,
+                &wdb.params,
+                account,
+                diversifier_index,
+                &addr,
+            ) {
+                Ok(_) => Ok(addr),
+                Err(rusqlite::Error::SqliteFailure(
+                    libsqlite3_sys::Error {
+                        code: libsqlite3_sys::ErrorCode::ConstraintViolation,
+                        ..
+                    },
+                    _,
+                )) => Ok(addr), // conflicts are ignorable
+                Err(e) => Err(e.into()),
+            };
+        })
+    }
+
     #[tracing::instrument(skip_all, fields(height = blocks.first().map(|b| u32::from(b.height()))))]
     #[allow(clippy::type_complexity)]
     fn put_blocks(