From 1554033a9ae30a81841132b680e72787d22064cb Mon Sep 17 00:00:00 2001
From: emjshrx <emjshrx@gmail.com>
Date: Sun, 18 Feb 2024 16:01:14 +0530
Subject: [PATCH] feat: add lookahead

---
 src/wallet/db/db.interface.ts |  3 ++-
 src/wallet/db/level/db.ts     | 11 +++++++++--
 src/wallet/db/level/layout.ts |  3 ++-
 src/wallet/wallet.ts          | 20 ++++++++++++++++----
 4 files changed, 29 insertions(+), 8 deletions(-)

diff --git a/src/wallet/db/db.interface.ts b/src/wallet/db/db.interface.ts
index d5800b2..b106dee 100644
--- a/src/wallet/db/db.interface.ts
+++ b/src/wallet/db/db.interface.ts
@@ -15,7 +15,8 @@ export type DbInterface = {
         encryptedChainCode: string,
     ): Promise<void>;
     saveAddress(address: string, path: string): Promise<void>;
-    getAddress(address: string): Promise<string>;
+    getPathFromAddress(address: string): Promise<string>;
+    getAddressFromPath(path: string): Promise<string>;
     hasAddress(address: string): Promise<boolean>;
     getReceiveDepth(): Promise<number>;
     setReceiveDepth(depth: number): Promise<void>;
diff --git a/src/wallet/db/level/db.ts b/src/wallet/db/level/db.ts
index a62a8ab..01d83fc 100644
--- a/src/wallet/db/level/db.ts
+++ b/src/wallet/db/level/db.ts
@@ -56,13 +56,20 @@ export class WalletDB implements DbInterface {
     }
 
     public async saveAddress(address: string, path: string): Promise<void> {
-        await this.db.sublevel(wdb.A).put(address, path);
+        await Promise.all([
+            this.db.sublevel(wdb.A).put(address, path),
+            this.db.sublevel(wdb.P).put(path, address),
+        ]);
     }
 
-    public async getAddress(address: string): Promise<string> {
+    public async getPathFromAddress(address: string): Promise<string> {
         return await this.db.sublevel(wdb.A).get(address);
     }
 
+    public async getAddressFromPath(path: string): Promise<string> {
+        return await this.db.sublevel(wdb.P).get(path);
+    }
+
     async hasAddress(address: string): Promise<boolean> {
         return (await this.db.sublevel(wdb.A).get(address)) !== undefined;
     }
diff --git a/src/wallet/db/level/layout.ts b/src/wallet/db/level/layout.ts
index 3e6cd59..3f2ab1a 100644
--- a/src/wallet/db/level/layout.ts
+++ b/src/wallet/db/level/layout.ts
@@ -2,7 +2,8 @@
 export const wdb = {
     V: 'V', // Version
     M: 'M', // Master key
-    A: 'A', // Address
+    A: 'A', // Address:Path
+    P: 'P', // Path:Address
     D: 'D', // Address depth
     C: 'C', // Coins
     SP: 'SP', // Silent payment address
diff --git a/src/wallet/wallet.ts b/src/wallet/wallet.ts
index 1506ae3..320f7e4 100644
--- a/src/wallet/wallet.ts
+++ b/src/wallet/wallet.ts
@@ -20,9 +20,11 @@ const bip32 = BIP32Factory(ecc);
 export type WalletConfigOptions = {
     db: DbInterface;
     networkClient: NetworkInterface;
+    lookahead?: number;
 };
 
 const DEFAULT_ENCRYPTION_PASSWORD = '12345678';
+const DEFAULT_LOOKAHEAD = 10;
 
 export class Wallet {
     private readonly db: DbInterface;
@@ -30,10 +32,12 @@ export class Wallet {
     private masterKey: BIP32Interface;
     private receiveDepth: number = 0;
     private changeDepth: number = 0;
+    private lookahead: number;
 
     constructor(config: WalletConfigOptions) {
         this.db = config.db;
         this.network = config.networkClient;
+        this.lookahead = config.lookahead ?? DEFAULT_LOOKAHEAD;
     }
 
     async init(params?: { mnemonic?: string; password?: string }) {
@@ -44,6 +48,9 @@ export class Wallet {
             const seed = mnemonicToSeedSync(mnemonic).toString('hex');
             this.masterKey = bip32.fromSeed(Buffer.from(seed, 'hex'));
             this.setPassword(password ?? DEFAULT_ENCRYPTION_PASSWORD);
+            for (let i = 0; i < this.lookahead; i++) {
+                await this.deriveAddress(`m/84'/0'/0'/0/${i}`);
+            }
         } else {
             const { encryptedPrivateKey, encryptedChainCode } =
                 await this.db.getMasterKey();
@@ -102,8 +109,11 @@ export class Wallet {
     }
 
     async deriveReceiveAddress(): Promise<string> {
-        const path = `m/84'/0'/0'/0/${this.receiveDepth}`;
-        const address = await this.deriveAddress(path);
+        const nextPath = `m/84'/0'/0'/0/${this.receiveDepth + this.lookahead}`;
+        await this.deriveAddress(nextPath);
+        const address = await this.db.getAddressFromPath(
+            `m/84'/0'/0'/0/${this.receiveDepth}`,
+        );
         this.receiveDepth++;
         return address;
     }
@@ -133,7 +143,7 @@ export class Wallet {
 
     private async signTransaction(psbt: Psbt, coins: Coin[]): Promise<void> {
         for (let index = 0; index < coins.length; index++) {
-            const path = await this.db.getAddress(coins[index].address);
+            const path = await this.db.getPathFromAddress(coins[index].address);
             const privateKey = this.masterKey.derivePath(path);
             psbt.signInput(index, privateKey);
         }
@@ -237,7 +247,9 @@ export class Wallet {
 
         const privateKeys = (
             await Promise.all(
-                selectedCoins.map((coin) => this.db.getAddress(coin.address)),
+                selectedCoins.map((coin) =>
+                    this.db.getPathFromAddress(coin.address),
+                ),
             )
         ).map((path) => this.masterKey.derivePath(path));
         const outpoints = selectedCoins.map((coin) => ({