From 282a8f7bb47c99a26a964703f55aef30010b5ac6 Mon Sep 17 00:00:00 2001
From: Christopher Jeffrey <chjjeffrey@gmail.com>
Date: Sat, 22 Oct 2016 05:12:27 -0700
Subject: [PATCH] walletdb: improve block connection.

---
 lib/chain/chaindb.js   |   6 ++
 lib/wallet/txdb.js     | 147 ++++++++++++++++++++++++++++++++++-------
 lib/wallet/wallet.js   |  18 ++---
 lib/wallet/walletdb.js | 145 ++++++++++++++++------------------------
 4 files changed, 194 insertions(+), 122 deletions(-)

diff --git a/lib/chain/chaindb.js b/lib/chain/chaindb.js
index 96c1e549..18fcf299 100644
--- a/lib/chain/chaindb.js
+++ b/lib/chain/chaindb.js
@@ -1062,6 +1062,12 @@ ChainDB.prototype.scan = co(function* scan(start, filter, iter) {
 
   entry = yield this.getEntry(start);
 
+  if (!entry)
+    return;
+
+  if (!(yield entry.isMainChain()))
+    throw new Error('Cannot rescan an alternate chain.');
+
   while (entry) {
     block = yield this.getFullBlock(entry.hash);
     total++;
diff --git a/lib/wallet/txdb.js b/lib/wallet/txdb.js
index 42f97e53..2aec72a1 100644
--- a/lib/wallet/txdb.js
+++ b/lib/wallet/txdb.js
@@ -865,13 +865,13 @@ TXDB.prototype.isSpent = co(function* isSpent(hash, index) {
  * @returns {Promise}
  */
 
-TXDB.prototype.add = co(function* add(tx) {
+TXDB.prototype.add = co(function* add(tx, block) {
   var result;
 
   this.start();
 
   try {
-    result = yield this._add(tx);
+    result = yield this._add(tx, block);
   } catch (e) {
     this.drop();
     throw e;
@@ -889,7 +889,7 @@ TXDB.prototype.add = co(function* add(tx) {
  * @returns {Promise}
  */
 
-TXDB.prototype._add = co(function* add(tx) {
+TXDB.prototype._add = co(function* add(tx, block) {
   var hash = tx.hash('hex');
   var existing = yield this.getTX(hash);
 
@@ -909,7 +909,7 @@ TXDB.prototype._add = co(function* add(tx) {
     tx.ps = existing.ps;
 
     // Confirm transaction.
-    return yield this.confirm(tx);
+    return yield this.confirm(tx, block);
   }
 
   if (tx.height === -1) {
@@ -935,7 +935,7 @@ TXDB.prototype._add = co(function* add(tx) {
   }
 
   // Finally we can do a regular insertion.
-  return yield this.insert(tx);
+  return yield this.insert(tx, block);
 });
 
 /**
@@ -945,9 +945,10 @@ TXDB.prototype._add = co(function* add(tx) {
  * @returns {Promise}
  */
 
-TXDB.prototype.insert = co(function* insert(tx) {
+TXDB.prototype.insert = co(function* insert(tx, block) {
   var hash = tx.hash('hex');
-  var details = new Details(this, tx);
+  var height = block ? block.height : this.walletdb.height;
+  var details = new Details(this, tx, height);
   var updated = false;
   var i, input, output, coin;
   var prevout, credit, path, account;
@@ -1082,6 +1083,11 @@ TXDB.prototype.insert = co(function* insert(tx) {
       this.put(layout.H(account, tx.height, hash), DUMMY);
   }
 
+  yield this.addTXRecord(tx);
+
+  if (block)
+    yield this.addBlockRecord(tx, block);
+
   // Update the transaction counter and
   // commit the new state. This state will
   // only overwrite the best state once
@@ -1104,6 +1110,88 @@ TXDB.prototype.insert = co(function* insert(tx) {
   return details;
 });
 
+TXDB.prototype.addTXRecord = co(function* addTXRecord(tx) {
+  var hash = tx.hash('hex');
+  var wids = yield this.walletdb.getWalletsByTX(hash);
+  var result;
+
+  if (!wids)
+    wids = [];
+
+  result = utils.binaryInsert(wids, this.wallet.wid, cmp, true);
+
+  if (result === -1)
+    return;
+
+  this.walletdb.writeTX(this.wallet, hash, wids);
+});
+
+TXDB.prototype.addBlockRecord = co(function* addBlockRecord(tx, record) {
+  var hash = tx.hash('hex');
+  var block = yield this.walletdb.getBlock(record.hash);
+
+  if (!block)
+    block = record.clone();
+
+  if (block.hashes.indexOf(hash) !== -1)
+    return;
+
+  block.hashes.push(hash);
+
+  this.walletdb.writeBlock(this.wallet, block);
+});
+
+TXDB.prototype.removeTXRecord = co(function* removeTXRecord(tx) {
+  var hash = tx.hash('hex');
+  var wids = yield this.walletdb.getWalletsByTX(hash);
+  var result;
+
+  if (!wids)
+    return true;
+
+  result = utils.binaryRemove(wids, this.wallet.wid, cmp);
+
+  if (!result)
+    return false;
+
+  if (wids.length > 0) {
+    this.walletdb.writeTX(this.wallet, hash, wids);
+    return false;
+  }
+
+  this.walletdb.removeTX(this.wallet, hash);
+
+  return true;
+});
+
+TXDB.prototype.removeBlockRecord = co(function* removeBlockRecord(tx, blk) {
+  var hash = tx.hash('hex');
+  var block = yield this.walletdb.getBlock(blk);
+  var index;
+
+  if (!block)
+    return;
+
+  index = block.hashes.indexOf(hash);
+
+  if (index === -1)
+    return;
+
+  block.hashes.splice(index, 1);
+
+  if (block.hashes.length > 0) {
+    this.walletdb.writeBlock(this.wallet, block);
+    return;
+  }
+
+  this.walletdb.removeBlock(this.wallet, block);
+});
+
+
+function cmp(a, b) {
+  return a - b;
+}
+
 /**
  * Attempt to confirm a transaction.
  * @private
@@ -1111,9 +1199,10 @@ TXDB.prototype.insert = co(function* insert(tx) {
  * @returns {Promise}
  */
 
-TXDB.prototype.confirm = co(function* confirm(tx) {
+TXDB.prototype.confirm = co(function* confirm(tx, block) {
   var hash = tx.hash('hex');
-  var details = new Details(this, tx);
+  var height = block ? block.height : this.walletdb.height;
+  var details = new Details(this, tx, height);
   var i, account, output, coin, input, prevout;
   var path, credit, credits;
 
@@ -1209,6 +1298,11 @@ TXDB.prototype.confirm = co(function* confirm(tx) {
     this.put(layout.H(account, tx.height, hash), DUMMY);
   }
 
+  yield this.addTXRecord(tx);
+
+  if (block)
+    yield this.addBlockRecord(tx, block);
+
   // Commit the new state. The balance has updated.
   this.put(layout.R, this.pending.commit());
 
@@ -1246,7 +1340,7 @@ TXDB.prototype.remove = co(function* remove(hash) {
 
 TXDB.prototype.erase = co(function* erase(tx) {
   var hash = tx.hash('hex');
-  var details = new Details(this, tx);
+  var details = new Details(this, tx, this.walletdb.height);
   var i, path, account, credits;
   var input, output, coin, credit;
 
@@ -1333,6 +1427,11 @@ TXDB.prototype.erase = co(function* erase(tx) {
       this.del(layout.H(account, tx.height, hash));
   }
 
+  if (yield this.removeTXRecord(tx)) {
+    if (tx.block)
+      yield this.removeBlockRecord(tx, tx.block);
+  }
+
   // Update the transaction counter
   // and commit new state due to
   // balance change.
@@ -1389,13 +1488,13 @@ TXDB.prototype.removeRecursive = co(function* removeRecursive(tx) {
  * @returns {Promise}
  */
 
-TXDB.prototype.unconfirm = co(function* unconfirm(hash) {
+TXDB.prototype.unconfirm = co(function* unconfirm(hash, block) {
   var details;
 
   this.start();
 
   try {
-    details = yield this._unconfirm(hash);
+    details = yield this._unconfirm(hash, block);
   } catch (e) {
     this.drop();
     throw e;
@@ -1413,13 +1512,13 @@ TXDB.prototype.unconfirm = co(function* unconfirm(hash) {
  * @returns {Promise}
  */
 
-TXDB.prototype._unconfirm = co(function* unconfirm(hash) {
+TXDB.prototype._unconfirm = co(function* unconfirm(hash, block) {
   var tx = yield this.getTX(hash);
 
   if (!tx)
     return;
 
-  return yield this.disconnect(tx);
+  return yield this.disconnect(tx, block);
 });
 
 /**
@@ -1428,9 +1527,9 @@ TXDB.prototype._unconfirm = co(function* unconfirm(hash) {
  * @returns {Promise}
  */
 
-TXDB.prototype.disconnect = co(function* disconnect(tx) {
+TXDB.prototype.disconnect = co(function* disconnect(tx, block) {
   var hash = tx.hash('hex');
-  var details = new Details(this, tx);
+  var details = new Details(this, tx, block.height - 1);
   var height = tx.height;
   var i, account, output, coin, credits;
   var input, path, credit;
@@ -2338,7 +2437,7 @@ TXDB.prototype.toDetails = co(function* toDetails(txs) {
  */
 
 TXDB.prototype._toDetails = co(function* _toDetails(tx) {
-  var details = new Details(this, tx);
+  var details = new Details(this, tx, this.walletdb.height);
   var coins = yield this.fillHistory(tx);
   var i, coin, path, output;
 
@@ -2887,14 +2986,14 @@ Credit.fromTX = function fromTX(tx, index) {
  * @param {TX} tx
  */
 
-function Details(txdb, tx) {
+function Details(txdb, tx, height) {
   if (!(this instanceof Details))
-    return new Details(txdb, tx);
+    return new Details(txdb, tx, height);
 
-  this.db = txdb.walletdb;
-  this.network = this.db.network;
-  this.wid = txdb.wallet.wid;
-  this.id = txdb.wallet.id;
+  this.wallet = txdb.wallet;
+  this.network = this.wallet.network;
+  this.wid = this.wallet.wid;
+  this.id = this.wallet.id;
 
   this.hash = tx.hash('hex');
   this.tx = tx;
@@ -2903,7 +3002,7 @@ function Details(txdb, tx) {
   this.height = tx.height;
   this.ts = tx.ts;
   this.index = tx.index;
-  this.confirmations = tx.getConfirmations(this.db.height);
+  this.confirmations = tx.getConfirmations(height);
 
   this.ps = tx.ps;
 
diff --git a/lib/wallet/wallet.js b/lib/wallet/wallet.js
index 26b463e1..8b64734b 100644
--- a/lib/wallet/wallet.js
+++ b/lib/wallet/wallet.js
@@ -1813,10 +1813,10 @@ Wallet.prototype.getTX = function getTX(hash) {
  * @returns {Promise}
  */
 
-Wallet.prototype.add = co(function* add(tx) {
+Wallet.prototype.add = co(function* add(tx, block) {
   var unlock = yield this.writeLock.lock();
   try {
-    return yield this._add(tx);
+    return yield this._add(tx, block);
   } finally {
     unlock();
   }
@@ -1830,8 +1830,8 @@ Wallet.prototype.add = co(function* add(tx) {
  * @returns {Promise}
  */
 
-Wallet.prototype._add = co(function* add(tx) {
-  var resolved = yield this.txdb.resolve(tx);
+Wallet.prototype._add = co(function* add(tx, block) {
+  var resolved = yield this.txdb.resolve(tx, block);
   var result = false;
   var i;
 
@@ -1839,7 +1839,7 @@ Wallet.prototype._add = co(function* add(tx) {
     return true;
 
   for (i = 0; i < resolved.length; i++) {
-    if (yield this._insert(resolved[i]))
+    if (yield this._insert(resolved[i], block))
       result = true;
   }
 
@@ -1853,13 +1853,13 @@ Wallet.prototype._add = co(function* add(tx) {
  * @returns {Promise}
  */
 
-Wallet.prototype._insert = co(function* insert(tx) {
+Wallet.prototype._insert = co(function* insert(tx, block) {
   var details, derived;
 
   this.txdb.start();
 
   try {
-    details = yield this.txdb._add(tx);
+    details = yield this.txdb._add(tx, block);
     derived = yield this.syncOutputDepth(details);
   } catch (e) {
     this.txdb.drop();
@@ -1882,10 +1882,10 @@ Wallet.prototype._insert = co(function* insert(tx) {
  * @returns {Promise}
  */
 
-Wallet.prototype.unconfirm = co(function* unconfirm(hash) {
+Wallet.prototype.unconfirm = co(function* unconfirm(hash, block) {
   var unlock = yield this.writeLock.lock();
   try {
-    return yield this.txdb.unconfirm(hash);
+    return yield this.txdb.unconfirm(hash, block);
   } finally {
     unlock();
   }
diff --git a/lib/wallet/walletdb.js b/lib/wallet/walletdb.js
index 4762ae10..8e5bd0c0 100644
--- a/lib/wallet/walletdb.js
+++ b/lib/wallet/walletdb.js
@@ -1343,40 +1343,43 @@ WalletDB.prototype.setTip = co(function* setTip(hash, height) {
  * @returns {Promise}
  */
 
-WalletDB.prototype.writeBlock = function writeBlock(block, matches) {
-  var batch = this.db.batch();
-  var i, hash, wids;
-
-  batch.put(layout.R, block.toTip());
-
-  if (block.hashes.length === 0)
-    return batch.write();
-
+WalletDB.prototype.writeBlock = function writeBlock(wallet, block) {
+  var batch = this.batch(wallet);
   batch.put(layout.b(block.hash), block.toRaw());
-
-  for (i = 0; i < block.hashes.length; i++) {
-    hash = block.hashes[i];
-    wids = matches[i];
-    batch.put(layout.e(hash), serializeWallets(wids));
-  }
-
-  return batch.write();
 };
 
 /**
- * Disconnect a block.
+ * Connect a transaction.
  * @param {WalletBlock} block
  * @returns {Promise}
  */
 
-WalletDB.prototype.unwriteBlock = function unwriteBlock(block) {
-  var batch = this.db.batch();
-  var prev = new WalletBlock(block.prevBlock, block.height - 1);
+WalletDB.prototype.writeTX = function writeTX(wallet, hash, wids) {
+  var batch = this.batch(wallet);
+  this.addFilter(hash);
+  batch.put(layout.e(hash), serializeWallets(wids));
+};
 
-  batch.put(layout.R, prev.toTip());
+/**
+ * Connect a block.
+ * @param {WalletBlock} block
+ * @returns {Promise}
+ */
+
+WalletDB.prototype.removeBlock = function removeBlock(wallet, block) {
+  var batch = this.batch(wallet);
   batch.del(layout.b(block.hash));
+};
 
-  return batch.write();
+/**
+ * Connect a transaction.
+ * @param {WalletBlock} block
+ * @returns {Promise}
+ */
+
+WalletDB.prototype.removeTX = function removeTX(wallet, hash) {
+  var batch = this.batch(wallet);
+  batch.del(layout.e(hash));
 };
 
 /**
@@ -1432,7 +1435,8 @@ WalletDB.prototype.addBlock = co(function* addBlock(entry, txs) {
  */
 
 WalletDB.prototype._addBlock = co(function* addBlock(entry, txs) {
-  var i, block, matches, hash, tx, wids;
+  var total = 0;
+  var i, block, tx;
 
   if (this.options.useCheckpoints) {
     if (entry.height <= this.network.checkpoints.lastHeight) {
@@ -1442,37 +1446,22 @@ WalletDB.prototype._addBlock = co(function* addBlock(entry, txs) {
   }
 
   block = WalletBlock.fromEntry(entry);
-  matches = [];
-
-  // Update these early so transactions
-  // get correct confirmation calculations.
-  this.tip = block.hash;
-  this.height = block.height;
 
   // Atomicity doesn't matter here. If we crash
   // during this loop, the automatic rescan will get
   // the database back into the correct state.
   for (i = 0; i < txs.length; i++) {
     tx = txs[i];
-    wids = yield this._insertTX(tx);
-
-    if (!wids)
-      continue;
-
-    hash = tx.hash('hex');
-
-    this.addFilter(hash);
-
-    block.hashes.push(hash);
-    matches.push(wids);
+    if (yield this._insertTX(tx, block))
+      total++;
   }
 
-  if (block.hashes.length > 0) {
+  if (total > 0) {
     this.logger.info('Connecting block %s (%d txs).',
-      utils.revHex(block.hash), block.hashes.length);
+      utils.revHex(block.hash), total);
   }
 
-  yield this.writeBlock(block, matches);
+  yield this.setTip(entry.hash, entry.height);
 });
 
 /**
@@ -1500,15 +1489,8 @@ WalletDB.prototype.removeBlock = co(function* removeBlock(entry) {
 
 WalletDB.prototype._removeBlock = co(function* removeBlock(entry) {
   var block = WalletBlock.fromEntry(entry);
-  var i, data, hash;
-
-  // If we crash during a reorg, there's not much to do.
-  // Reorgs cannot be rescanned. The database will be
-  // in an odd state, with some txs being confirmed
-  // when they shouldn't be. That being said, this
-  // should eventually resolve itself when a new block
-  // comes in.
-  data = yield this.getBlock(block.hash);
+  var data = yield this.getBlock(block.hash);
+  var i, hash;
 
   if (data)
     block.hashes = data.hashes;
@@ -1518,16 +1500,12 @@ WalletDB.prototype._removeBlock = co(function* removeBlock(entry) {
       utils.revHex(block.hash), block.hashes.length);
   }
 
-  // Unwrite the tip as fast as we can.
-  yield this.unwriteBlock(block);
-
   for (i = block.hashes.length - 1; i >= 0; i--) {
     hash = block.hashes[i];
-    yield this._unconfirmTX(hash);
+    yield this._unconfirmTX(hash, block);
   }
 
-  this.tip = block.hash;
-  this.height = block.height;
+  yield this.setTip(block.prevBlock, block.height - 1);
 });
 
 /**
@@ -1541,7 +1519,7 @@ WalletDB.prototype._removeBlock = co(function* removeBlock(entry) {
 WalletDB.prototype.addTX = co(function* addTX(tx) {
   var unlock = yield this.txLock.lock();
   try {
-    return yield this._addTX(tx);
+    return yield this._insertTX(tx);
   } finally {
     unlock();
   }
@@ -1554,30 +1532,7 @@ WalletDB.prototype.addTX = co(function* addTX(tx) {
  * @returns {Promise}
  */
 
-WalletDB.prototype._addTX = co(function* addTX(tx) {
-  var wids = yield this._insertTX(tx);
-  var hash;
-
-  if (!wids)
-    return;
-
-  hash = tx.hash('hex');
-
-  yield this.db.put(layout.e(hash), serializeWallets(wids));
-
-  this.addFilter(hash);
-
-  return wids;
-});
-
-/**
- * Add a transaction to the database without a lock.
- * @private
- * @param {TX} tx
- * @returns {Promise}
- */
-
-WalletDB.prototype._insertTX = co(function* insertTX(tx) {
+WalletDB.prototype._insertTX = co(function* insertTX(tx, block) {
   var result = false;
   var i, wids, wid, wallet;
 
@@ -1598,7 +1553,7 @@ WalletDB.prototype._insertTX = co(function* insertTX(tx) {
 
     assert(wallet);
 
-    if (yield wallet.add(tx)) {
+    if (yield wallet.add(tx, block)) {
       this.logger.debug(
         'Added transaction to wallet: %s (%d).',
         wallet.id, wid);
@@ -1618,10 +1573,10 @@ WalletDB.prototype._insertTX = co(function* insertTX(tx) {
  * @returns {Promise}
  */
 
-WalletDB.prototype.unconfirmTX = co(function* unconfirmTX(hash) {
+WalletDB.prototype.unconfirmTX = co(function* unconfirmTX(hash, block) {
   var unlock = yield this.txLock.lock();
   try {
-    return yield this._unconfirmTX(hash);
+    return yield this._unconfirmTX(hash, block);
   } finally {
     unlock();
   }
@@ -1635,7 +1590,7 @@ WalletDB.prototype.unconfirmTX = co(function* unconfirmTX(hash) {
  * @returns {Promise}
  */
 
-WalletDB.prototype._unconfirmTX = co(function* unconfirmTX(hash) {
+WalletDB.prototype._unconfirmTX = co(function* unconfirmTX(hash, block) {
   var wids = yield this.getWalletsByTX(hash);
   var i, wid, wallet;
 
@@ -1646,7 +1601,7 @@ WalletDB.prototype._unconfirmTX = co(function* unconfirmTX(hash) {
     wid = wids[i];
     wallet = yield this.get(wid);
     assert(wallet);
-    yield wallet.unconfirm(hash);
+    yield wallet.unconfirm(hash, block);
   }
 });
 
@@ -1701,6 +1656,18 @@ function WalletBlock(hash, height) {
   this.hashes = [];
 }
 
+/**
+ * Clone the block.
+ * @returns {WalletBlock}
+ */
+
+WalletBlock.prototype.clone = function clone() {
+  var block = new WalletBlock(this.hash, this.height);
+  block.prevBlock = this.prevBlock;
+  block.hashes = this.hashes.slice();
+  return block;
+};
+
 /**
  * Instantiate wallet block from chain entry.
  * @private