From a478e395246d0f6453adfa9bf71384a5b7d4baad Mon Sep 17 00:00:00 2001
From: William Wolf <throughnothing@gmail.com>
Date: Mon, 23 Feb 2015 08:46:33 -0800
Subject: [PATCH] Review Fixes

---
 lib/block/merkleblock.js | 108 ++++++++++++++++++++-------------------
 test/merkleblock.js      |   2 +-
 2 files changed, 57 insertions(+), 53 deletions(-)

diff --git a/lib/block/merkleblock.js b/lib/block/merkleblock.js
index bebbf17..74d425c 100644
--- a/lib/block/merkleblock.js
+++ b/lib/block/merkleblock.js
@@ -57,10 +57,12 @@ function MerkleBlock(arg) {
       flags: arg.flags
     };
   } else {
-    throw new TypeError('Unrecognized argument for Block');
+    throw new TypeError('Unrecognized argument for MerkleBlock');
   }
   _.extend(this,info);
   this._validMerkleTree = null;
+  this._flagBitsUsed = 0;
+  this._hashesUsed = 0;
   return this;
 }
 
@@ -140,13 +142,10 @@ MerkleBlock.prototype.toJSON = function toJSON() {
  * @returns {Bool} - True/False whether this MerkleBlock is Valid
  */
 MerkleBlock.prototype.validMerkleTree = function validMerkleTree() {
-  $.checkState(this.flags instanceof Array, 'MerkleBlock flags is not an array');
-  $.checkState(this.hashes instanceof Array, 'MerkleBlock flags is not an array');
-  var self = this;
-  if(this._validMerkleTree === true) {
-    return true;
-  } else if (this._validMerkleTree === false) {
-    return false;
+  $.checkState(_.isArray(this.flags), 'MerkleBlock flags is not an array');
+  $.checkState(_.isArray(this.hashes), 'MerkleBlock hashes is not an array');
+  if(_.isBoolean(this._validMerkleTree)) {
+    return this._validMerkleTree;
   }
 
   // Can't have more hashes than numTransactions
@@ -159,66 +158,71 @@ MerkleBlock.prototype.validMerkleTree = function validMerkleTree() {
     return this._setValidMerkleTree(false);
   }
 
-  // Calculate height of tree
-  // From Bitcoin Core merkleblock.h CalcTreeWidth() + CPartialMerkleTree
   var height = 0;
-  while (calcTreeWidth(height) > 1) {
+  while (this._calcTreeWidth(height) > 1) {
     height++;
   }
 
-  var flagBitsUsed = 0;
-  var hashesUsed = 0;
-
-  // Modeled after Bitcoin Core merkleblock.h CalcTreeWidth()
-  function calcTreeWidth(height) {
-    return (self.numTransactions + (1 << height) - 1) >> height;
-  }
-
-  // Modeled after Bitcoin Core merkleblock.cpp TraverseAndExtract()
-  function traverse(depth, pos) {
-    if(flagBitsUsed > self.flags.length * 8) {
-      return null;
-    }
-    var isParentOfMatch = (self.flags[flagBitsUsed >> 3] >>> (flagBitsUsed++ & 7)) & 1;
-    if(depth === 0 || !isParentOfMatch) {
-      if(hashesUsed >= self.hashes.length) {
-        return null;
-      }
-      var hash = self.hashes[hashesUsed++];
-      return new Buffer(hash, 'hex');
-    } else {
-      var left = traverse(depth-1, pos*2);
-      var right;
-      if(pos*2+1 < calcTreeWidth(depth-1)) {
-        right = traverse(depth-1, pos*2+1);
-      } else {
-        right = left;
-      }
-      return Hash.sha256sha256(new Buffer.concat([left, right]));
-    }
-  }
-
-  var root = traverse(height, 0);
-  if(hashesUsed !== this.hashes.length) {
+  this._flagBitsUsed = 0;
+  this._hashesUsed = 0;
+  var root = this._traverseMerkleTree(height, 0);
+  if(this._hashesUsed !== this.hashes.length) {
     return this._setValidMerkleTree(false);
   }
   return this._setValidMerkleTree(BufferUtil.equals(root, this.header.merkleRoot));
 }
 
+/** Traverse a the tree in this MerkleBlock, validating it along the way
+ *  Modeled after Bitcoin Core merkleblock.cpp TraverseAndExtract()
+ * @param {Number} - Current height
+ * @param {Number} - Current position in the tree
+ * @returns {Buffer|null} - Buffer containing the Merkle Hash far that height
+ * @private
+ */
+MerkleBlock.prototype._traverseMerkleTree = function traverseMerkleTree(depth, pos) {
+  if(this._flagBitsUsed > this.flags.length * 8) {
+    return null;
+  }
+  var isParentOfMatch = (this.flags[this._flagBitsUsed >> 3] >>> (this._flagBitsUsed++ & 7)) & 1;
+  if(depth === 0 || !isParentOfMatch) {
+    if(this._hashesUsed >= this.hashes.length) {
+      return null;
+    }
+    var hash = this.hashes[this._hashesUsed++];
+    return new Buffer(hash, 'hex');
+  } else {
+    var left = this._traverseMerkleTree(depth-1, pos*2);
+    var right;
+    if(pos*2+1 < this._calcTreeWidth(depth-1)) {
+      right = this._traverseMerkleTree(depth-1, pos*2+1);
+    } else {
+      right = left;
+    }
+    return Hash.sha256sha256(new Buffer.concat([left, right]));
+  }
+}
+
+/** Calculates the width of a merkle tree at a given height.
+ *  Modeled after Bitcoin Core merkleblock.h CalcTreeWidth()
+ * @param {Number} - Height at which we want the tree width
+ * @returns {Number} - Width of the tree at a given height
+ * @private
+ */
+MerkleBlock.prototype._calcTreeWidth = function calcTreeWidth(height) {
+  return (this.numTransactions + (1 << height) - 1) >> height;
+}
+
 /**
  * @param {Transaction|String} - Transaction or Transaction ID Hash
  * @returns {Bool} - return true/false if this MerkleBlock has the TX or not
  * @private
  */
 MerkleBlock.prototype.hasTransaction = function hasTransaction(tx) {
-  $.checkArgument(!_.isUndefined(tx), 'No transaction given');
-  $.checkArgument(tx instanceof Transaction
-    || typeof tx === 'string', 'No transaction given');
+  $.checkArgument(!_.isUndefined(tx), 'tx cannot be undefined');
+  $.checkArgument(tx instanceof Transaction || typeof tx === 'string',
+      'Invalid tx given, tx must be a "string" or "Transaction"');
 
-  var hash = tx;
-  if(tx instanceof Transaction) {
-    hash = tx.id;
-  }
+  var hash = tx.id || tx;
   var revHash = BufferUtil.reverse(new Buffer(hash,'hex')).toString('hex');
 
   return (this.hashes.indexOf(hash) !== -1
diff --git a/test/merkleblock.js b/test/merkleblock.js
index 4c34a0c..61e5201 100644
--- a/test/merkleblock.js
+++ b/test/merkleblock.js
@@ -34,7 +34,7 @@ describe('MerkleBlock', function() {
     it('should not make an empty block', function() {
       (function() {
         return new MerkleBlock();
-      }).should.throw('Unrecognized argument for Block');
+      }).should.throw('Unrecognized argument for MerkleBlock');
     });
   });