From 33171fd0698a8cbd308e74b0d078d412b198290b Mon Sep 17 00:00:00 2001
From: Chris Stewart <stewart.chris1234@gmail.com>
Date: Sat, 26 Mar 2022 06:22:09 -0500
Subject: [PATCH] Use the set of addresses in the wallet when rescanning rather
 than generating more (#4211)

---
 .../bitcoins/wallet/RescanHandlingTest.scala  |  2 +-
 .../wallet/internal/RescanHandling.scala      | 77 +++++++++++++------
 2 files changed, 55 insertions(+), 24 deletions(-)

diff --git a/wallet-test/src/test/scala/org/bitcoins/wallet/RescanHandlingTest.scala b/wallet-test/src/test/scala/org/bitcoins/wallet/RescanHandlingTest.scala
index 3faa91a68f..7ba0e549cb 100644
--- a/wallet-test/src/test/scala/org/bitcoins/wallet/RescanHandlingTest.scala
+++ b/wallet-test/src/test/scala/org/bitcoins/wallet/RescanHandlingTest.scala
@@ -310,7 +310,7 @@ class RescanHandlingTest extends BitcoinSWalletTestCachedBitcoindNewest {
 
       //slight delay to make sure other rescan is started
       val alreadyStartedF =
-        AsyncUtil.nonBlockingSleep(500.millis).flatMap { _ =>
+        AsyncUtil.nonBlockingSleep(100.millis).flatMap { _ =>
           wallet.rescanNeutrinoWallet(startOpt = None,
                                       endOpt = None,
                                       addressBatchSize =
diff --git a/wallet/src/main/scala/org/bitcoins/wallet/internal/RescanHandling.scala b/wallet/src/main/scala/org/bitcoins/wallet/internal/RescanHandling.scala
index acc06f9935..e5cd8e05da 100644
--- a/wallet/src/main/scala/org/bitcoins/wallet/internal/RescanHandling.scala
+++ b/wallet/src/main/scala/org/bitcoins/wallet/internal/RescanHandling.scala
@@ -147,13 +147,19 @@ private[wallet] trait RescanHandling extends WalletLogger {
       addressBatchSize: Int): Future[Unit] = {
     for {
       scriptPubKeys <- generateScriptPubKeys(account, addressBatchSize)
+      addressCount <- addressDAO.count()
       _ <- matchBlocks(scriptPubKeys = scriptPubKeys,
                        endOpt = endOpt,
                        startOpt = startOpt)
       externalGap <- calcAddressGap(HDChainType.External, account)
       changeGap <- calcAddressGap(HDChainType.Change, account)
-      res <-
-        if (
+      res <- {
+        logger.info(s"addressCount=$addressCount externalGap=$externalGap")
+        if (addressCount != 0) {
+          logger.info(
+            s"We have a small number of addresses preloaded into the wallet")
+          Future.unit
+        } else if (
           externalGap >= walletConfig.addressGapLimit && changeGap >= walletConfig.addressGapLimit
         ) {
           logger.info(
@@ -165,6 +171,7 @@ private[wallet] trait RescanHandling extends WalletLogger {
               s"match within our address gap limit of ${walletConfig.addressGapLimit}")
           doNeutrinoRescan(account, startOpt, endOpt, addressBatchSize)
         }
+      }
     } yield res
   }
 
@@ -225,34 +232,58 @@ private[wallet] trait RescanHandling extends WalletLogger {
     blocksF
   }
 
+  /** Use to generate a list of addresses to search when restoring our wallet
+    *  from our mneomnic seed
+    */
+  private def generateAddressesForRescan(
+      account: HDAccount,
+      count: Int): Future[Vector[BitcoinAddress]] = {
+    val receiveAddressesF = 1
+      .to(count)
+      .foldLeft(Future.successful(Vector.empty[BitcoinAddress])) {
+        (prevFuture, _) =>
+          for {
+            prev <- prevFuture
+            address <- getNewAddress(account)
+          } yield prev :+ address
+      }
+
+    val changeAddressesF = 1
+      .to(count)
+      .foldLeft(Future.successful(Vector.empty[BitcoinAddress])) {
+        (prevFuture, _) =>
+          for {
+            prev <- prevFuture
+            address <- getNewChangeAddress(account)
+          } yield prev :+ address
+      }
+    for {
+      receiveAddresses <- receiveAddressesF
+      changeAddresses <- changeAddressesF
+    } yield receiveAddresses ++ changeAddresses
+
+  }
+
   private def generateScriptPubKeys(
       account: HDAccount,
       count: Int): Future[Vector[ScriptPubKey]] = {
+    val addressCountF = addressDAO.count()
     for {
-      addresses <-
-        1
-          .to(count)
-          .foldLeft(Future.successful(Vector.empty[BitcoinAddress])) {
-            (prevFuture, _) =>
-              for {
-                prev <- prevFuture
-                address <- getNewAddress(account)
-              } yield prev :+ address
-          }
-      changeAddresses <-
-        1
-          .to(count)
-          .foldLeft(Future.successful(Vector.empty[BitcoinAddress])) {
-            (prevFuture, _) =>
-              for {
-                prev <- prevFuture
-                address <- getNewChangeAddress(account)
-              } yield prev :+ address
-          }
+      addressCount <- addressCountF
+      addresses <- {
+        if (addressCount == 0) {
+          generateAddressesForRescan(account, count)
+        } else {
+          //we don't want to continously generate addresses
+          //if our wallet already has them, so just use what is in the
+          //database already
+          addressDAO.findAllAddresses().map(_.map(_.address))
+        }
+      }
       spksDb <- scriptPubKeyDAO.findAll()
     } yield {
       val addrSpks =
-        addresses.map(_.scriptPubKey) ++ changeAddresses.map(_.scriptPubKey)
+        addresses.map(_.scriptPubKey)
       val otherSpks = spksDb.map(_.scriptPubKey)
 
       (addrSpks ++ otherSpks).distinct