From 7c22067f6934abe254fba9b51d9cccf00b657173 Mon Sep 17 00:00:00 2001
From: Daniel Cousens <github@dcousens.com>
Date: Wed, 13 Aug 2014 13:56:17 +1000
Subject: [PATCH] Wallet: clarify getPrivateKeyForAddress method structure

This does repeat the O(n) lookup several times, but that can be fixed by
using an O(1) lookup instead (and will be later).

Clarity first.
---
 src/wallet.js  | 15 ++++++++++-----
 test/wallet.js |  2 +-
 2 files changed, 11 insertions(+), 6 deletions(-)

diff --git a/src/wallet.js b/src/wallet.js
index b2d9a05..f377377 100644
--- a/src/wallet.js
+++ b/src/wallet.js
@@ -208,13 +208,18 @@ function Wallet(seed, network) {
   }
 
   this.getPrivateKeyForAddress = function(address) {
-    var index
-    if((index = this.addresses.indexOf(address)) > -1) {
+    assert(isMyAddress(address), 'Unknown address. Make sure the address is from the keychain and has been generated')
+
+    if (isReceiveAddress(address)) {
+      var index = this.addresses.indexOf(address)
+
       return this.getPrivateKey(index)
-    } else if((index = this.changeAddresses.indexOf(address)) > -1) {
+    }
+
+    if (isChangeAddress(address)) {
+      var index = this.changeAddresses.indexOf(address)
+
       return this.getInternalPrivateKey(index)
-    } else {
-      throw new Error('Unknown address. Make sure the address is from the keychain and has been generated.')
     }
   }
 
diff --git a/test/wallet.js b/test/wallet.js
index 80fd7a8..456775d 100644
--- a/test/wallet.js
+++ b/test/wallet.js
@@ -164,7 +164,7 @@ describe('Wallet', function() {
       var wallet = new Wallet(seed, networks.testnet)
       assert.throws(function() {
         wallet.getPrivateKeyForAddress("n2fiWrHqD6GM5GiEqkbWAc6aaZQp3ba93X")
-      }, /Unknown address. Make sure the address is from the keychain and has been generated./)
+      }, /Unknown address. Make sure the address is from the keychain and has been generated/)
     })
   })