HDKeyGeneration: add generate() method, use in DeterministicKeyChain

core/src/main/java/org/bitcoinj/crypto/HDKeyDerivation.java

@@ -23,6 +23,8 @@ import java.math.*;
 import java.nio.*;
 import java.security.*;
 import java.util.*;
+import java.util.function.Supplier;
+import java.util.stream.Stream;
 
 import static com.google.common.base.Preconditions.*;
 
@@ -122,6 +124,18 @@ public final class HDKeyDerivation {
 
     }
 
+    /**
+     * Generate an infinite stream of {@link DeterministicKey}s from the given parent and index.
+     * <p>
+     * Note: Use {@link Stream#limit(long)} to get the desired number of keys.
+     * @param parent The parent key
+     * @param childNumber The index of the first child to supply/generate
+     * @return A stream of {@code DeterministicKey}s
+     */
+    public static Stream<DeterministicKey> generate(DeterministicKey parent, int childNumber) {
+        return Stream.generate(new KeySupplier(parent, childNumber));
+    }
+
     /**
      * @throws HDDerivationException if private derivation is attempted for a public-only parent key, or
      * if the resulting derived key is invalid (eg. private key == 0).
@@ -228,6 +242,33 @@ public final class HDKeyDerivation {
             throw new HDDerivationException(errorMessage);
     }
 
+    /**
+     * A supplier of a sequence of {@code DeterministicKey}s generated from a starting
+     * parent and child index.
+     * <p><b>Not threadsafe: Should be used on a single thread</b>
+     */
+    private static class KeySupplier implements Supplier<DeterministicKey> {
+        private final DeterministicKey parent;
+        private int nextChild;
+
+        /**
+         * Construct a key supplier with the specified starting point
+         * @param parent The parent key
+         * @param nextChild The index of the next child to supply/generate
+         */
+        public KeySupplier(DeterministicKey parent, int nextChild) {
+            this.parent = parent;
+            this.nextChild = nextChild;
+        }
+
+        @Override
+        public DeterministicKey get() {
+            DeterministicKey key = HDKeyDerivation.deriveThisOrNextChildKey(parent, nextChild);
+            nextChild = key.getChildNumber().num() + 1;
+            return key;
+        }
+    }
+
     public static class RawKeyBytes {
         public final byte[] keyBytes, chainCode;
 

core/src/main/java/org/bitcoinj/wallet/DeterministicKeyChain.java

@@ -41,6 +41,7 @@ import java.util.*;
 import java.util.concurrent.Executor;
 import java.util.concurrent.locks.ReentrantLock;
 import java.util.function.Predicate;
+import java.util.stream.Collectors;
 
 import static com.google.common.base.Preconditions.*;
 import static java.util.stream.Collectors.collectingAndThen;
@@ -1202,23 +1203,19 @@ public class DeterministicKeyChain implements EncryptableKeyChain {
         checkState(lock.isHeldByCurrentThread());
         final int numChildren = hierarchy.getNumChildren(parent.getPath());
         final int needed = issued + lookaheadSize + lookaheadThreshold - numChildren;
-
-        if (needed <= lookaheadThreshold)
-            return new ArrayList<>();
+        final int limit = (needed > lookaheadThreshold) ? needed : 0;
 
         log.info("{} keys needed for {} = {} issued + {} lookahead size + {} lookahead threshold - {} num children",
-                needed, parent.getPathAsString(), issued, lookaheadSize, lookaheadThreshold, numChildren);
+                limit, parent.getPathAsString(), issued, lookaheadSize, lookaheadThreshold, numChildren);
 
-        List<DeterministicKey> result  = new ArrayList<>(needed);
         final Stopwatch watch = Stopwatch.createStarted();
-        int nextChild = numChildren;
-        for (int i = 0; i < needed; i++) {
-            DeterministicKey key = HDKeyDerivation.deriveThisOrNextChildKey(parent, nextChild);
-            key = key.dropPrivateBytes();
+        List<DeterministicKey> result = HDKeyDerivation.generate(parent, numChildren)
+                .limit(limit)
+                .map(DeterministicKey::dropPrivateBytes)
+                .collect(Collectors.toList());
+        result.forEach(key -> {
             hierarchy.putKey(key);
-            result.add(key);
-            nextChild = key.getChildNumber().num() + 1;
-        }
+        });
         watch.stop();
         log.info("Took {}", watch);
         return result;

core/src/test/java/org/bitcoinj/crypto/HDKeyDerivationTest.java

@@ -22,6 +22,9 @@ import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
 
 import java.math.BigInteger;
+import java.util.List;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
 
 import org.bitcoinj.core.Utils;
 import org.bitcoinj.crypto.HDKeyDerivation.PublicDeriveMode;
@@ -136,4 +139,30 @@ public class HDKeyDerivationTest {
         assertEquals(EXPECTED_CHILD_PRIVATE_KEY, fromPublicWithInversion.getPrivateKeyAsHex());
         assertEquals(EXPECTED_CHILD_PUBLIC_KEY, fromPublicWithInversion.getPublicKeyAsHex());
     }
+
+    @Test
+    public void testGenerate() {
+        DeterministicKey parent = new DeterministicKey(HDPath.m(), new byte[32], BigInteger.TEN,
+                null);
+        assertFalse(parent.isPubKeyOnly());
+        assertFalse(parent.isEncrypted());
+
+        List<DeterministicKey> keys0 = HDKeyDerivation.generate(parent, CHILD_NUMBER.num())
+                                                        .limit(0)
+                                                        .collect(Collectors.toList());
+        assertEquals(0, keys0.size());
+
+        List<DeterministicKey> keys1 = HDKeyDerivation.generate(parent, CHILD_NUMBER.num())
+                                                        .limit(1)
+                                                        .collect(Collectors.toList());
+        assertEquals(1, keys1.size());
+        assertEquals(HDPath.m(CHILD_NUMBER), keys1.get(0).getPath());
+
+        List<DeterministicKey> keys2 = HDKeyDerivation.generate(parent, CHILD_NUMBER.num())
+                                                        .limit(2)
+                                                        .collect(Collectors.toList());
+        assertEquals(2, keys2.size());
+        assertEquals(HDPath.parsePath("m/1"), keys2.get(0).getPath());
+        assertEquals(HDPath.parsePath("m/2"), keys2.get(1).getPath());
+    }
 }