From 53900f3196362b6705ab62e1822ac5879ac73f5c Mon Sep 17 00:00:00 2001
From: Esteban Ordano <eordano@gmail.com>
Date: Fri, 28 Nov 2014 12:24:32 -0300
Subject: [PATCH] Add cache to derivation

---
 lib/hdkeycache.js   | 16 ++++++++++++++++
 lib/hdprivatekey.js | 12 +++++++++++-
 lib/hdpublickey.js  |  9 ++++++++-
 test/hdpublickey.js |  7 +++++++
 4 files changed, 42 insertions(+), 2 deletions(-)
 create mode 100644 lib/hdkeycache.js

diff --git a/lib/hdkeycache.js b/lib/hdkeycache.js
new file mode 100644
index 0000000..d2e566b
--- /dev/null
+++ b/lib/hdkeycache.js
@@ -0,0 +1,16 @@
+'use strict';
+
+var cache = {};
+
+module.exports = {
+  get: function(xkey, number, hardened) {
+    var key = xkey + '/' + number + '/' + hardened;
+    if (cache[key]) {
+      return cache[key];
+    }
+  },
+  set: function(xkey, number, hardened, derived) {
+    var key = xkey + '/' + number + '/' + hardened;
+    cache[key] = derived;
+  }
+};
diff --git a/lib/hdprivatekey.js b/lib/hdprivatekey.js
index f7c1955..78d5e53 100644
--- a/lib/hdprivatekey.js
+++ b/lib/hdprivatekey.js
@@ -6,6 +6,7 @@ var Base58 = require('./encoding/base58');
 var Base58Check = require('./encoding/base58check');
 var Hash = require('./crypto/hash');
 var Network = require('./networks');
+var HDKeyCache = require('./hdkeycache');
 var Point = require('./crypto/point');
 var PrivateKey = require('./privatekey');
 var Random = require('./crypto/random');
@@ -59,12 +60,18 @@ HDPrivateKey.prototype.derive = function(arg, hardened) {
 };
 
 HDPrivateKey.prototype._deriveWithNumber = function(index, hardened) {
+  /* jshint maxstatements: 20 */
+  /* jshint maxcomplexity: 10 */
   if (index >= HDPrivateKey.Hardened) {
     hardened = true;
   }
   if (index < HDPrivateKey.Hardened && hardened) {
     index += HDPrivateKey.Hardened;
   }
+  var cached = HDKeyCache.get(this.xprivkey, index, hardened);
+  if (cached) {
+    return cached;
+  }
 
   var indexBuffer = util.integerAsBuffer(index);
   var data;
@@ -79,7 +86,7 @@ HDPrivateKey.prototype._deriveWithNumber = function(index, hardened) {
 
   var privateKey = leftPart.add(this.privateKey.toBigNumber()).mod(Point.getN()).toBuffer({size: 32});
 
-  return new HDPrivateKey({
+  var derived = new HDPrivateKey({
     network: this.network,
     depth: this.depth + 1,
     parentFingerPrint: this.fingerPrint,
@@ -87,6 +94,8 @@ HDPrivateKey.prototype._deriveWithNumber = function(index, hardened) {
     chainCode: chainCode,
     privateKey: privateKey
   });
+  HDKeyCache.set(this.xprivkey, index, hardened, derived);
+  return derived;
 };
 
 HDPrivateKey.prototype._deriveFromString = function(path) {
@@ -174,6 +183,7 @@ HDPrivateKey.prototype._buildFromJson = function(arg) {
 };
 
 HDPrivateKey.prototype._buildFromObject = function(arg) {
+  /* jshint maxcomplexity: 12 */
   // TODO: Type validation
   var buffers = {
     version: arg.network ? util.integerAsBuffer(Network.get(arg.network).xprivkey) : arg.version,
diff --git a/lib/hdpublickey.js b/lib/hdpublickey.js
index e947016..faf5fbd 100644
--- a/lib/hdpublickey.js
+++ b/lib/hdpublickey.js
@@ -6,6 +6,7 @@ var Base58 = require('./encoding/base58');
 var Base58Check = require('./encoding/base58check');
 var Hash = require('./crypto/hash');
 var HDPrivateKey = require('./hdprivatekey');
+var HDKeyCache = require('./hdkeycache');
 var Network = require('./networks');
 var Point = require('./crypto/point');
 var PublicKey = require('./publickey');
@@ -67,6 +68,10 @@ HDPublicKey.prototype._deriveWithNumber = function (index, hardened) {
   if (hardened || index >= HDPublicKey.Hardened) {
     throw new Error(HDPublicKey.Errors.InvalidIndexCantDeriveHardened);
   }
+  var cached = HDKeyCache.get(this.xpubkey, index, hardened);
+  if (cached) {
+    return cached;
+  }
 
   var indexBuffer = util.integerAsBuffer(index);
   var data = buffer.Buffer.concat([this.publicKey.toBuffer(), indexBuffer]);
@@ -76,7 +81,7 @@ HDPublicKey.prototype._deriveWithNumber = function (index, hardened) {
 
   var publicKey = PublicKey.fromPoint(Point.getG().mul(leftPart).add(this.publicKey.point));
 
-  return new HDPublicKey({
+  var derived = new HDPublicKey({
     network: this.network,
     depth: this.depth + 1,
     parentFingerPrint: this.fingerPrint,
@@ -84,6 +89,8 @@ HDPublicKey.prototype._deriveWithNumber = function (index, hardened) {
     chainCode: chainCode,
     publicKey: publicKey
   });
+  HDKeyCache.set(this.xpubkey, index, hardened, derived);
+  return derived;
 };
 
 HDPublicKey.prototype._deriveFromString = function (path) {
diff --git a/test/hdpublickey.js b/test/hdpublickey.js
index 2c208a6..219eaee 100644
--- a/test/hdpublickey.js
+++ b/test/hdpublickey.js
@@ -149,5 +149,12 @@ describe('HDPublicKey interface', function() {
       expect(function() { return new HDPublicKey(xpubkey).derive(HDPublicKey.Hardened + 1); })
         .to.throw(HDPublicKey.Errors.InvalidIndexCantDeriveHardened);
     });
+
+    it('should use the cache', function() {
+      var pubkey = new HDPublicKey(xpubkey);
+      var derived1 = pubkey.derive(0);
+      var derived2 = pubkey.derive(0);
+      derived1.xpubkey.should.equal(derived2.xpubkey);
+    });
   });
 });