From 3bb3d82aaca05c78b1d389f1d03fb8b79a0e32dd Mon Sep 17 00:00:00 2001
From: Matias Alejo Garcia <ematiu@gmail.com>
Date: Wed, 10 Feb 2016 15:03:34 -0500
Subject: [PATCH] add counter for address mempool index

---
 lib/services/address/index.js       | 31 +++++++++++++++++++----------
 test/services/address/index.unit.js |  5 ++++-
 2 files changed, 25 insertions(+), 11 deletions(-)

diff --git a/lib/services/address/index.js b/lib/services/address/index.js
index b5a98e18..35829758 100644
--- a/lib/services/address/index.js
+++ b/lib/services/address/index.js
@@ -275,6 +275,25 @@ AddressService.prototype.transactionHandler = function(txInfo, callback) {
 
 };
 
+AddressService.prototype._updateAddressIndex = function(key, add) {
+
+  var currentValue = this.mempoolAddressIndex[key] || 0;
+
+  if(add) {
+    if (currentValue > 0)
+      this.mempoolAddressIndex[key] = currentValue + 1;
+    else 
+      this.mempoolAddressIndex[key] = 1;
+  } else {
+    if (currentValue < 1) {
+      delete this.mempoolAddressIndex[key];
+    } else {
+      this.mempoolAddressIndex[key]--;
+    }
+  }
+};
+
+
 /**
  * This function will update the mempool address index with the necessary
  * information for further lookups.
@@ -309,11 +328,7 @@ AddressService.prototype.updateMempoolIndex = function(tx, add, callback) {
 
     var addressIndexKey = encoding.encodeMempoolAddressIndexKey(addressInfo.hashBuffer, addressInfo.hashTypeBuffer);
 
-    if(add) {
-      this.mempoolAddressIndex[addressIndexKey] = true;
-    } else {
-      delete this.mempoolAddressIndex[addressIndexKey];
-    }
+    this._updateAddressIndex(addressIndexKey, add);
 
     // Update output index
     var outputIndexBuffer = new Buffer(4);
@@ -409,11 +424,7 @@ AddressService.prototype.updateMempoolIndex = function(tx, add, callback) {
 
     var addressIndexKey = encoding.encodeMempoolAddressIndexKey(inputHashBuffer, inputHashType);
 
-    if(add) {
-      this.mempoolAddressIndex[addressIndexKey] = true;
-    } else {
-      delete this.mempoolAddressIndex[addressIndexKey];
-    }
+    this._updateAddressIndex(addressIndexKey, add);
   }
 
   if (!callback) {
diff --git a/test/services/address/index.unit.js b/test/services/address/index.unit.js
index 62745395..d5632c8f 100644
--- a/test/services/address/index.unit.js
+++ b/test/services/address/index.unit.js
@@ -9,6 +9,7 @@ var bitcorenode = require('../../../');
 var AddressService = bitcorenode.services.Address;
 var blockData = require('../../data/livenet-345003.json');
 var bitcore = require('bitcore-lib');
+var _ = bitcore.deps._;
 var memdown = require('memdown');
 var leveldown = require('leveldown');
 var Networks = bitcore.Networks;
@@ -1978,6 +1979,7 @@ describe('Address Service', function() {
         callback.should.be.a('function');
         Object.keys(am.mempoolSpentIndex).length.should.equal(14);
         Object.keys(am.mempoolAddressIndex).length.should.equal(5);
+        _.values(am.mempoolAddressIndex).should.deep.equal([1,1,12,1,1]);
         for (var i = 0; i < operations.length; i++) {
           operations[i].type.should.equal('put');
         }
@@ -2013,6 +2015,7 @@ describe('Address Service', function() {
         for (var i = 0; i < operations.length; i++) {
           operations[i].type.should.equal('del');
         }
+        Object.keys(am.mempoolAddressIndex).length.should.equal(0);
       };
       am.updateMempoolIndex(tx, false);
     });
@@ -2523,7 +2526,7 @@ describe('Address Service', function() {
 
       var hashTypeBuffer = constants.HASH_TYPES_MAP[address.type];
       var addressIndex = encoding.encodeMempoolAddressIndexKey(address.hashBuffer, hashTypeBuffer);
-      as.mempoolAddressIndex[addressIndex] = true;
+      as.mempoolAddressIndex[addressIndex] = 1;
 
       as._getInputsMempool = sinon.stub().callsArgWith(3, null, mempoolInputs);
       as._getOutputsMempool = sinon.stub().callsArgWith(3, null, mempoolOutputs);