From 7ce4360c0052bcdc17270688610ec5dc9c08cc6e Mon Sep 17 00:00:00 2001
From: Mathias Buus <mathiasbuus@gmail.com>
Date: Tue, 7 May 2024 08:45:18 +0200
Subject: [PATCH] cache the merkle batch in batches (#516)

---
 lib/batch.js       | 33 +++++++++++++++++++++++++--------
 lib/merkle-tree.js | 12 ++++++++++--
 2 files changed, 35 insertions(+), 10 deletions(-)

diff --git a/lib/batch.js b/lib/batch.js
index 2dea959e..8a2c36f5 100644
--- a/lib/batch.js
+++ b/lib/batch.js
@@ -25,6 +25,7 @@ module.exports = class HypercoreBatch extends EventEmitter {
     this._sessionLength = 0
     this._sessionByteLength = 0
     this._sessionBatch = null
+    this._cachedBatch = null
     this._flushing = null
     this._clear = clear
 
@@ -207,20 +208,32 @@ module.exports = class HypercoreBatch extends EventEmitter {
     return this.session.core.tree.restoreBatch(length)
   }
 
-  createTreeBatch (length, blocks = []) {
+  _catchupBatch (clone) {
+    if (this._cachedBatch === null) this._cachedBatch = this._sessionBatch.clone()
+
+    if (this.length > this._cachedBatch.length) {
+      const offset = this._cachedBatch.length - this._sessionBatch.length
+
+      for (let i = offset; i < this._appendsActual.length; i++) {
+        this._cachedBatch.append(this._appendsActual[i])
+      }
+    }
+
+    return clone ? this._cachedBatch.clone() : this._cachedBatch
+  }
+
+  createTreeBatch (length, opts = {}) {
+    if (Array.isArray(opts)) opts = { blocks: opts }
+
+    const { blocks = [], clone = true } = opts
     if (!length && length !== 0) length = this.length + blocks.length
 
     const maxLength = this.length + blocks.length
-    const b = this._sessionBatch.clone()
+    const b = this._catchupBatch(clone || (blocks.length > 0 || length !== this.length))
     const len = Math.min(length, this.length)
 
     if (len < this._sessionLength || length > maxLength) return null
-
-    for (let i = 0; i < len - this._sessionLength; i++) {
-      b.append(this._appendsActual[i])
-    }
-
-    if (len < this.length) return b
+    if (len < b.length) b.checkout(len, this._sessionBatch.roots)
 
     for (let i = 0; i < length - len; i++) {
       b.append(this._appendsActual === this._appends ? blocks[i] : this._encrypt(b.length, blocks[i]))
@@ -239,6 +252,8 @@ module.exports = class HypercoreBatch extends EventEmitter {
     if (typeof opts === 'number') opts = { fork: opts }
     const { fork = this.fork + 1, force = false } = opts
 
+    this._cachedBatch = null
+
     const length = this._sessionLength
     if (newLength < length) {
       if (!force) throw new Error('Cannot truncate committed blocks')
@@ -380,6 +395,8 @@ module.exports = class HypercoreBatch extends EventEmitter {
     this._sessionByteLength = info.byteLength
     this._sessionBatch = newBatch
 
+    if (this._cachedBatch !== null) this._cachedBatch.prune(info.length)
+
     const same = this._appends === this._appendsActual
 
     this._appends = this._appends.slice(flushingLength)
diff --git a/lib/merkle-tree.js b/lib/merkle-tree.js
index 7851a013..e8ebf63b 100644
--- a/lib/merkle-tree.js
+++ b/lib/merkle-tree.js
@@ -57,13 +57,19 @@ class MerkleTreeBatch {
     this.upgraded = false
   }
 
-  checkout (length) {
+  checkout (length, additionalRoots) {
     const roots = []
     let r = 0
 
     const head = 2 * length - 2
     const gaps = new Set()
     const all = new Map()
+
+    // additional roots is so the original roots can be passed (we mutate the array in appendRoot)
+    if (additionalRoots) {
+      for (const node of additionalRoots) all.set(node.index, node)
+    }
+
     for (const node of this.nodes) all.set(node.index, node)
 
     for (const index of flat.fullRoots(head + 2)) {
@@ -74,7 +80,9 @@ class MerkleTreeBatch {
         roots.push(this.roots[r++])
         continue
       }
-      roots.push(all.get(index))
+      const node = all.get(index)
+      if (!node) throw new BAD_ARGUMENT('root missing for given length')
+      roots.push(node)
     }
 
     this.roots = roots