Iterator pattern for walking derived keys

2 files changed, 60 insertions, 44 deletions

diff --git a/packages/subproviders/src/subproviders/ledger.ts b/packages/subproviders/src/subproviders/ledger.ts
index 3c0442e4c..6c685c84d 100644
--- a/packages/subproviders/src/subproviders/ledger.ts
+++ b/packages/subproviders/src/subproviders/ledger.ts
@@ -88,9 +88,8 @@ export class LedgerSubprovider extends BaseWalletSubprovider {
     public async getAccountsAsync(
         numberOfAccounts: number = walletUtils.DEFAULT_NUM_ADDRESSES_TO_FETCH,
     ): Promise<string[]> {
-        const offset = 0;
         const initialHDerivedKey = await this._initialDerivedKeyAsync();
-        const derivedKeys = walletUtils.calculateDerivedHDKeys(initialHDerivedKey, numberOfAccounts, offset);
+        const derivedKeys = walletUtils.calculateDerivedHDKeys(initialHDerivedKey, numberOfAccounts);
         const accounts = _.map(derivedKeys, 'address');
         return accounts;
     }
diff --git a/packages/subproviders/src/utils/wallet_utils.ts b/packages/subproviders/src/utils/wallet_utils.ts
index 0c3e895b6..bd1851d9a 100644
--- a/packages/subproviders/src/utils/wallet_utils.ts
+++ b/packages/subproviders/src/utils/wallet_utils.ts
@@ -6,64 +6,81 @@ import { DerivedHDKey, WalletSubproviderErrors } from '../types';
 
 const DEFAULT_ADDRESS_SEARCH_OFFSET = 0;
 const BATCH_SIZE = 10;
+const DEFAULT_ADDRESS_SEARCH_LIMIT = 1000;
+
+class DerivedHDKeyIterator implements IterableIterator<DerivedHDKey> {
+    private _initialDerivedKey: DerivedHDKey;
+    private _searchLimit: number;
+    private _index: number;
+
+    constructor(initialDerivedKey: DerivedHDKey, searchLimit: number = DEFAULT_ADDRESS_SEARCH_OFFSET) {
+        this._searchLimit = searchLimit;
+        this._initialDerivedKey = initialDerivedKey;
+        this._index = 0;
+    }
+
+    public next(): IteratorResult<DerivedHDKey> {
+        const derivationPath = this._initialDerivedKey.derivationPath;
+        const derivationIndex = this._index;
+        // If the DerivedHDKey is a child then we walk relative, if not we walk the full derivation path
+        const path = this._initialDerivedKey.isChildKey
+            ? `m/${derivationIndex}`
+            : `m/${derivationPath}/${derivationIndex}`;
+        const hdKey = this._initialDerivedKey.hdKey.derive(path);
+        const address = walletUtils.addressOfHDKey(hdKey);
+        const derivedKey: DerivedHDKey = {
+            address,
+            hdKey,
+            derivationPath,
+            derivationIndex,
+            isChildKey: this._initialDerivedKey.isChildKey,
+        };
+        const done = this._index === this._searchLimit;
+        this._index++;
+        return {
+            done,
+            value: derivedKey,
+        };
+    }
+
+    public [Symbol.iterator](): IterableIterator<DerivedHDKey> {
+        return this;
+    }
+}
 
 export const walletUtils = {
+    DEFAULT_ADDRESS_SEARCH_LIMIT,
     DEFAULT_NUM_ADDRESSES_TO_FETCH: 10,
-    DEFAULT_ADDRESS_SEARCH_LIMIT: 1000,
-    calculateDerivedHDKeys(
-        initialDerivedKey: DerivedHDKey,
-        searchLimit: number,
-        offset: number = DEFAULT_ADDRESS_SEARCH_OFFSET,
-    ): DerivedHDKey[] {
+    calculateDerivedHDKeys(initialDerivedKey: DerivedHDKey, searchLimit: number): DerivedHDKey[] {
         const derivedKeys: DerivedHDKey[] = [];
-        _.times(searchLimit, i => {
-            const derivationPath = initialDerivedKey.derivationPath;
-            const derivationIndex = offset + i;
-            // If the DerivedHDKey is a child then we walk relative, if not we walk the full derivation path
-            const path = initialDerivedKey.isChildKey
-                ? `m/${derivationIndex}`
-                : `m/${derivationPath}/${derivationIndex}`;
-            const hdKey = initialDerivedKey.hdKey.derive(path);
-            const address = walletUtils.addressOfHDKey(hdKey);
-            const derivedKey: DerivedHDKey = {
-                address,
-                hdKey,
-                derivationPath,
-                derivationIndex,
-                isChildKey: initialDerivedKey.isChildKey,
-            };
-            derivedKeys.push(derivedKey);
-        });
+        const derivedKeyIterator = new DerivedHDKeyIterator(initialDerivedKey, searchLimit);
+        for (const key of derivedKeyIterator) {
+            derivedKeys.push(key);
+        }
         return derivedKeys;
     },
-    addressOfHDKey(hdKey: HDNode): string {
-        const shouldSanitizePublicKey = true;
-        const derivedPublicKey = hdKey.publicKey;
-        const ethereumAddressUnprefixed = ethUtil
-            .publicToAddress(derivedPublicKey, shouldSanitizePublicKey)
-            .toString('hex');
-        const address = ethUtil.addHexPrefix(ethereumAddressUnprefixed);
-        return address;
-    },
     findDerivedKeyByAddress(
         address: string,
         initialDerivedKey: DerivedHDKey,
         searchLimit: number,
     ): DerivedHDKey | undefined {
         let matchedKey: DerivedHDKey | undefined;
-        for (let index = 0; index < searchLimit; index = index + BATCH_SIZE) {
-            const derivedKeys = walletUtils.calculateDerivedHDKeys(initialDerivedKey, BATCH_SIZE, index);
-            matchedKey = _.find(derivedKeys, derivedKey => derivedKey.address === address);
-            if (matchedKey) {
+        const derivedKeyIterator = new DerivedHDKeyIterator(initialDerivedKey, searchLimit);
+        for (const key of derivedKeyIterator) {
+            if (key.address === address) {
+                matchedKey = key;
                 break;
             }
         }
         return matchedKey;
     },
-
-    _firstDerivedKey(initialDerivedKey: DerivedHDKey): DerivedHDKey {
-        const derivedKeys = walletUtils.calculateDerivedHDKeys(initialDerivedKey, 1, 0);
-        const firstDerivedKey = derivedKeys[0];
-        return firstDerivedKey;
+    addressOfHDKey(hdKey: HDNode): string {
+        const shouldSanitizePublicKey = true;
+        const derivedPublicKey = hdKey.publicKey;
+        const ethereumAddressUnprefixed = ethUtil
+            .publicToAddress(derivedPublicKey, shouldSanitizePublicKey)
+            .toString('hex');
+        const address = ethUtil.addHexPrefix(ethereumAddressUnprefixed);
+        return address;
     },
 };