Wallet restore API to support restore on batch (mimblewimble#18)

* wallet restore API to support restore on batch

* rustfmt

impls/src/lmdb_wallet.rs

@@ -30,7 +30,7 @@ use crate::store::{self, option_to_not_found, to_key, to_key_u64};
 
 use crate::core::core::Transaction;
 use crate::core::{global, ser};
-use crate::libwallet::{check_repair, restore};
+use crate::libwallet::{check_repair, restore, restore_batch};
 use crate::libwallet::{
 	AcctPathMapping, Context, Error, ErrorKind, NodeClient, OutputData, PaymentCommits,
 	PaymentData, TxLogEntry, WalletBackend, WalletOutputBatch,
@@ -374,6 +374,11 @@ where
 		Ok(())
 	}
 
+	fn restore_batch(&mut self, start_index: u64, batch_size: u64) -> Result<(u64, u64), Error> {
+		let res = restore_batch(self, start_index, batch_size).context(ErrorKind::Restore)?;
+		Ok(res)
+	}
+
 	fn check_repair(&mut self, delete_unconfirmed: bool) -> Result<(), Error> {
 		check_repair(self, delete_unconfirmed).context(ErrorKind::Restore)?;
 		Ok(())

libwallet/src/api_impl/owner.rs

@@ -326,6 +326,20 @@ where
 	w.restore()
 }
 
+/// Attempt to restore contents of wallet on batch
+pub fn restore_batch<T: ?Sized, C, K>(
+	w: &mut T,
+	start_index: u64,
+	batch_size: u64,
+) -> Result<(u64, u64), Error>
+where
+	T: WalletBackend<C, K>,
+	C: NodeClient,
+	K: Keychain,
+{
+	w.restore_batch(start_index, batch_size)
+}
+
 /// check repair
 pub fn check_repair<T: ?Sized, C, K>(w: &mut T, delete_unconfirmed: bool) -> Result<(), Error>
 where

libwallet/src/internal/restore.rs

@@ -389,34 +389,22 @@ where
 	Ok(())
 }
 
-/// Restore a wallet
-pub fn restore<T, C, K>(wallet: &mut T) -> Result<(), Error>
+fn restore_from_outputs<T, C, K>(wallet: &mut T, outputs: Vec<OutputResult>) -> Result<(), Error>
 where
 	T: WalletBackend<C, K>,
 	C: NodeClient,
 	K: Keychain,
 {
-	// Don't proceed if wallet_data has anything in it
-	let is_empty = wallet.iter().next().is_none();
-	if !is_empty {
-		error!("Not restoring. Please back up and remove existing db directory first.");
-		return Ok(());
-	}
-
-	warn!("Starting restore.");
-
-	let result_vec = collect_chain_outputs(wallet)?;
-
 	warn!(
 		"Identified {} wallet_outputs as belonging to this wallet",
-		result_vec.len(),
+		outputs.len(),
 	);
 
 	let mut found_parents: HashMap<Identifier, u32> = HashMap::new();
 	let mut restore_stats = HashMap::new();
 
 	// Now save what we have
-	for output in result_vec {
+	for output in outputs {
 		restore_missing_output(
 			wallet,
 			output,
@@ -453,3 +441,46 @@ where
 	}
 	Ok(())
 }
+
+/// Restore a wallet
+pub fn restore<T, C, K>(wallet: &mut T) -> Result<(), Error>
+where
+	T: WalletBackend<C, K>,
+	C: NodeClient,
+	K: Keychain,
+{
+	// Don't proceed if wallet_data has anything in it
+	let is_empty = wallet.iter().next().is_none();
+	if !is_empty {
+		error!("Not restoring. Please back up and remove existing db directory first.");
+		return Ok(());
+	}
+
+	warn!("Starting restore.");
+
+	let result_vec = collect_chain_outputs(wallet)?;
+
+	restore_from_outputs(wallet, result_vec)
+}
+
+/// Restore outputs by index on batch
+pub fn restore_batch<T, C, K>(
+	wallet: &mut T,
+	start_index: u64,
+	batch_size: u64,
+) -> Result<(u64, u64), Error>
+where
+	T: WalletBackend<C, K>,
+	C: NodeClient,
+	K: Keychain,
+{
+	let mut result_vec: Vec<OutputResult> = vec![];
+	let (highest_index, last_retrieved_index, outputs) = wallet
+		.w2n_client()
+		.get_outputs_by_pmmr_index(start_index, batch_size)?;
+
+	result_vec.append(&mut identify_utxo_outputs(wallet, outputs.clone())?);
+
+	restore_from_outputs(wallet, result_vec)?;
+	Ok((highest_index, last_retrieved_index))
+}

libwallet/src/lib.rs

@@ -53,7 +53,7 @@ pub use api_impl::types::{
 	BlockFees, CbData, InitTxArgs, InitTxSendArgs, NodeHeightResult, OutputCommitMapping,
 	PaymentCommitMapping, SendTXArgs, VersionInfo,
 };
-pub use internal::restore::{check_repair, restore};
+pub use internal::restore::{check_repair, restore, restore_batch};
 pub use types::{
 	AcctPathMapping, BlockIdentifier, Context, NodeClient, OutputData, OutputStatus,
 	PaymentCommits, PaymentData, TxLogEntry, TxLogEntryType, TxWrapper, WalletBackend, WalletInfo,

libwallet/src/types.rs

@@ -133,6 +133,9 @@ where
 	/// Attempt to restore the contents of a wallet from seed
 	fn restore(&mut self) -> Result<(), Error>;
 
+	/// Attempt to restore the contents of a wallet from seed on batch outputs
+	fn restore_batch(&mut self, start_index: u64, batch_size: u64) -> Result<(u64, u64), Error>;
+
 	/// Attempt to check and fix wallet state
 	fn check_repair(&mut self, delete_unconfirmed: bool) -> Result<(), Error>;
 }