hypercore 10.38.2 → 11.0.1

package/lib/merkle-tree.js

@@ -1,16 +1,10 @@
 const flat = require('flat-tree')
 const crypto = require('hypercore-crypto')
-const c = require('compact-encoding')
-const Xache = require('xache')
 const b4a = require('b4a')
 const unslab = require('unslab')
 const caps = require('./caps')
 const { INVALID_PROOF, INVALID_CHECKSUM, INVALID_OPERATION, BAD_ARGUMENT, ASSERTION } = require('hypercore-errors')
 
-const BLANK_HASH = b4a.alloc(32)
-const OLD_TREE = b4a.from([5, 2, 87, 2, 0, 0, 40, 7, 66, 76, 65, 75, 69, 50, 98])
-const TREE_CACHE = 128 // speeds up linear scans by A LOT
-
 class NodeQueue {
   constructor (nodes, extra = null) {
     this.i = 0
@@ -42,18 +36,22 @@ class NodeQueue {
 }
 
 class MerkleTreeBatch {
-  constructor (tree) {
-    this.fork = tree.fork
-    this.roots = [...tree.roots]
-    this.length = tree.length
-    this.ancestors = tree.length
-    this.byteLength = tree.byteLength
-    this.signature = null
+  constructor (tree, session) {
+    this.fork = session.fork
+    this.roots = [...session.roots]
+    this.length = session.length
+    this.signature = session.signature
+    this.ancestors = session.length
+    this.byteLength = session.byteLength
     this.hashCached = null
 
-    this.treeLength = tree.length
-    this.treeFork = tree.fork
+    this.committed = false
+    this.truncated = false
+    this.treeLength = session.length
+    this.treeFork = session.fork
     this.tree = tree
+    this.storage = tree.storage
+    this.session = session
     this.nodes = []
     this.upgraded = false
   }
@@ -121,7 +119,7 @@ class MerkleTreeBatch {
   }
 
   clone () {
-    const b = new MerkleTreeBatch(this.tree)
+    const b = new MerkleTreeBatch(this.tree, this.session)
 
     b.fork = this.fork
     b.roots = [...this.roots]
@@ -138,7 +136,7 @@ class MerkleTreeBatch {
   }
 
   hash () {
-    if (this.hashCached === null) this.hashCached = unslab(this.tree.crypto.tree(this.roots))
+    if (this.hashCached === null) this.hashCached = unslab(crypto.tree(this.roots))
     return this.hashCached
   }
 
@@ -162,22 +160,28 @@ class MerkleTreeBatch {
     return this.tree.get(index, error)
   }
 
-  proof ({ block, hash, seek, upgrade }) {
-    return generateProof(this, block, hash, seek, upgrade)
+  proof (batch, { block, hash, seek, upgrade }) {
+    return generateProof(batch, this, block, hash, seek, upgrade)
   }
 
   verifyUpgrade (proof) {
-    const unverified = verifyTree(proof, this.tree.crypto, this.nodes)
+    const unverified = verifyTree(proof, this.nodes)
 
     if (!proof.upgrade) throw INVALID_OPERATION('Expected upgrade proof')
 
     return verifyUpgrade(proof, unverified, this)
   }
 
+  addNodesUnsafe (nodes) {
+    for (let i = 0; i < nodes.length; i++) {
+      this.nodes.push(nodes[i])
+    }
+  }
+
   append (buf) {
     const head = this.length * 2
     const ite = flat.iterator(head)
-    const node = blockNode(this.tree.crypto, head, buf)
+    const node = blockNode(head, buf)
 
     this.appendRoot(node, ite)
   }
@@ -201,7 +205,7 @@ class MerkleTreeBatch {
         break
       }
 
-      const node = unslabNode(parentNode(this.tree.crypto, ite.parent(), a, b))
+      const node = unslabNode(parentNode(ite.parent(), a, b))
       this.nodes.push(node)
       this.roots.pop()
       this.roots.pop()
@@ -210,76 +214,95 @@ class MerkleTreeBatch {
   }
 
   commitable () {
-    return this.treeFork === this.tree.fork && (
+    return this.treeFork === this.session.fork && (
       this.upgraded
-        ? this.treeLength === this.tree.length
-        : this.treeLength <= this.tree.length
+        ? this.treeLength === this.session.length
+        : this.treeLength <= this.session.length
     )
   }
 
-  commit () {
+  commit (tx) {
+    if (tx === undefined) throw INVALID_OPERATION('No database batch was passed')
     if (!this.commitable()) throw INVALID_OPERATION('Tree was modified during batch, refusing to commit')
 
-    if (this.upgraded) this._commitUpgrade()
+    if (this.upgraded) this._commitUpgrade(tx)
 
     for (let i = 0; i < this.nodes.length; i++) {
       const node = this.nodes[i]
-      this.tree.unflushed.set(node.index, node)
+      tx.putTreeNode(node)
     }
+
+    this.committed = true
+
+    return this
   }
 
-  _commitUpgrade () {
+  _commitUpgrade (tx) {
     // TODO: If easy to detect, we should refuse an trunc+append here without a fork id
     // change. Will only happen on user error so mostly to prevent that.
 
     if (this.ancestors < this.treeLength) {
+      tx.deleteTreeNodeRange(this.ancestors * 2, this.treeLength * 2)
+
       if (this.ancestors > 0) {
-        const head = 2 * this.ancestors
+        const head = this.ancestors * 2
         const ite = flat.iterator(head - 2)
 
         while (true) {
           if (ite.contains(head) && ite.index < head) {
-            this.tree.unflushed.set(ite.index, blankNode(ite.index))
+            tx.deleteTreeNode(ite.index)
           }
           if (ite.offset === 0) break
           ite.parent()
         }
-      }
 
-      this.tree.truncateTo = this.tree.truncated
-        ? Math.min(this.tree.truncateTo, this.ancestors)
-        : this.ancestors
-
-      this.tree.truncated = true
-      this.tree.cache = new Xache({ maxSize: this.tree.cache.maxSize })
-      truncateMap(this.tree.unflushed, this.ancestors)
-      if (this.tree.flushing !== null) truncateMap(this.tree.flushing, this.ancestors)
+        this.truncated = true
+      }
     }
-
-    this.tree.roots = this.roots
-    this.tree.length = this.length
-    this.tree.byteLength = this.byteLength
-    this.tree.fork = this.fork
-    this.tree.signature = this.signature
   }
 
   seek (bytes, padding) {
-    return new ByteSeeker(this, bytes, padding)
+    return new ByteSeeker(this, this, bytes, padding)
   }
 
   byteRange (index) {
-    return getByteRange(this, index)
+    const readBatch = this.storage.read()
+    const range = getByteRange(this, index, readBatch)
+    readBatch.tryFlush()
+
+    return range
   }
 
   byteOffset (index) {
-    if (index === 2 * this.tree.length) return this.tree.byteLength
-    return getByteOffset(this, index)
+    if (index === 2 * this.length) return this.byteLength
+
+    const readBatch = this.storage.read()
+    const offset = getByteOffset(this, index, readBatch)
+    readBatch.tryFlush()
+
+    return offset
+  }
+
+  async restore (length) {
+    if (length === this.length) return this
+
+    const roots = unslabNodes(await this.tree.getRoots(length))
+
+    this.roots = roots
+    this.length = length
+    this.byteLength = totalSize(roots)
+    this.ancestors = length
+
+    for (const node of roots) this.byteLength += node.size
+
+    return this
   }
 }
 
 class ReorgBatch extends MerkleTreeBatch {
-  constructor (tree) {
-    super(tree)
+  constructor (tree, session) {
+    super(tree, session)
+
     this.roots = []
     this.length = 0
     this.byteLength = 0
@@ -303,7 +326,7 @@ class ReorgBatch extends MerkleTreeBatch {
     if (this.want === null) return true
 
     const nodes = []
-    const root = verifyTree(proof, this.tree.crypto, nodes)
+    const root = verifyTree(proof, nodes)
 
     if (root === null || !b4a.equals(root.hash, this.diff.hash)) return false
 
@@ -363,8 +386,9 @@ class ReorgBatch extends MerkleTreeBatch {
 }
 
 class ByteSeeker {
-  constructor (tree, bytes, padding = 0) {
+  constructor (tree, session, bytes, padding = 0) {
     this.tree = tree
+    this.session = session
     this.bytes = bytes
     this.padding = padding
 
@@ -377,7 +401,7 @@ class ByteSeeker {
   async _seek (bytes) {
     if (!bytes) return [0, 0]
 
-    for (const node of this.tree.roots) { // all async ticks happen once we find the root so safe
+    for (const node of this.session.roots) { // all async ticks happen once we find the root so safe
       const size = getUnpaddedSize(node, this.padding, null)
 
       if (bytes === size) return [flat.rightSpan(node.index) + 2, 0]
@@ -423,81 +447,102 @@ class ByteSeeker {
   }
 }
 
-module.exports = class MerkleTree {
-  constructor (storage, roots, fork, signature, prologue) {
-    this.crypto = crypto
-    this.fork = fork
-    this.roots = roots
-    this.length = roots.length ? totalSpan(roots) / 2 : 0
-    this.byteLength = totalSize(roots)
-    this.signature = signature
-    this.prologue = prologue
-
-    this.storage = storage
-    this.unflushed = new Map()
-    this.cache = new Xache({ maxSize: TREE_CACHE })
-    this.flushing = null
-    this.truncated = false
-    this.truncateTo = 0
+class TreeProof {
+  constructor (tree, block, hash, seek, upgrade) {
+    this.fork = tree.fork
+    this.signature = tree.signature
+
+    this.block = block
+    this.hash = hash
+    this.seek = seek
+    this.upgrade = upgrade
+
+    this.pending = {
+      node: null,
+      seek: null,
+      upgrade: null,
+      additionalUpgrade: null
+    }
   }
 
-  addNode (node) {
-    if (node.size === 0 && b4a.equals(node.hash, BLANK_HASH)) node = blankNode(node.index)
-    this.unflushed.set(node.index, node)
-  }
+  async settle () {
+    const result = { fork: this.fork, block: null, hash: null, seek: null, upgrade: null, manifest: null }
 
-  batch () {
-    return new MerkleTreeBatch(this)
-  }
+    const [pNode, pSeek, pUpgrade, pAdditional] = await settleProof(this.pending)
 
-  async restoreBatch (length) {
-    const batch = new MerkleTreeBatch(this)
-    if (length === this.length) return batch
+    if (this.block) {
+      if (pNode === null) throw INVALID_OPERATION('Invalid block request')
+      result.block = {
+        index: this.block.index,
+        value: null, // populated upstream, alloc it here for simplicity
+        nodes: pNode
+      }
+    } else if (this.hash) {
+      if (pNode === null) throw INVALID_OPERATION('Invalid block request')
+      result.hash = {
+        index: this.hash.index,
+        nodes: pNode
+      }
+    }
 
-    const roots = unslabNodes(await this.getRoots(length))
+    if (this.seek && pSeek !== null) {
+      result.seek = {
+        bytes: this.seek.bytes,
+        nodes: pSeek
+      }
+    }
 
-    batch.roots = roots
-    batch.length = length
-    batch.byteLength = 0
-    batch.ancestors = length
+    if (this.upgrade) {
+      result.upgrade = {
+        start: this.upgrade.start,
+        length: this.upgrade.length,
+        nodes: pUpgrade,
+        additionalNodes: pAdditional || [],
+        signature: this.signature
+      }
+    }
 
-    for (const node of roots) batch.byteLength += node.size
+    return result
+  }
+}
 
-    return batch
+class MerkleTree {
+  constructor (storage) {
+    this.storage = storage
   }
 
-  seek (bytes, padding) {
-    return new ByteSeeker(this, bytes, padding)
+  static hash (s) {
+    return unslab(crypto.tree(s.roots))
   }
 
-  hash () {
-    return unslab(this.crypto.tree(this.roots))
+  static signable (s, namespace) {
+    return caps.treeSignable(namespace, MerkleTree.hash(s), s.length, s.fork)
+  }
+
+  static size (roots) {
+    return totalSize(roots)
+  }
+
+  static span (roots) {
+    return totalSpan(roots)
   }
 
-  signable (namespace) {
-    return caps.treeSignable(namespace, this.hash(), this.length, this.fork)
+  clone (storage) {
+    return new MerkleTree(storage)
   }
 
   getRoots (length) {
     const indexes = flat.fullRoots(2 * length)
     const roots = new Array(indexes.length)
+    const readBatch = this.storage.read()
 
     for (let i = 0; i < indexes.length; i++) {
-      roots[i] = this.get(indexes[i], true)
+      roots[i] = readBatch.getTreeNode(indexes[i], true)
     }
 
-    return Promise.all(roots)
-  }
+    readBatch.tryFlush()
 
-  setPrologue ({ hash, length }) {
-    this.prologue = { hash, length }
-  }
-
-  addNodes (nodes) {
-    for (let i = 0; i < nodes.length; i++) {
-      const node = nodes[i]
-      this.unflushed.set(node.index, node)
-    }
+    return Promise.all(roots)
   }
 
   getNeededNodes (length, start, end) {
@@ -527,11 +572,14 @@ module.exports = class MerkleTree {
   async upgradeable (length) {
     const indexes = flat.fullRoots(2 * length)
     const roots = new Array(indexes.length)
+    const readBatch = this.storage.read()
 
     for (let i = 0; i < indexes.length; i++) {
-      roots[i] = this.get(indexes[i], false)
+      roots[i] = readBatch.getTreeNode(indexes[i], false)
     }
 
+    readBatch.tryFlush()
+
     for (const node of await Promise.all(roots)) {
       if (node === null) return false
     }
@@ -539,131 +587,25 @@ module.exports = class MerkleTree {
     return true
   }
 
-  blankNode (index) {
-    return blankNode(index)
-  }
-
-  get (index, error = true) {
-    const c = this.cache.get(index)
-    if (c) return c
-
-    let node = this.unflushed.get(index)
-
-    if (this.flushing !== null && node === undefined) {
-      node = this.flushing.get(index)
-    }
-
-    // TODO: test this
-    if (this.truncated && node !== undefined && node.index >= 2 * this.truncateTo) {
-      node = blankNode(index)
-    }
-
-    if (node !== undefined) {
-      if (node.hash === BLANK_HASH) {
-        if (error) throw INVALID_OPERATION('Could not load node: ' + index)
-        return Promise.resolve(null)
-      }
-      return Promise.resolve(node)
-    }
-
-    return getStoredNode(this.storage, index, this.cache, error)
-  }
-
-  async flush () {
-    this.flushing = this.unflushed
-    this.unflushed = new Map()
-
-    try {
-      if (this.truncated) await this._flushTruncation()
-      await this._flushNodes()
-    } catch (err) {
-      for (const node of this.flushing.values()) {
-        if (!this.unflushed.has(node.index)) this.unflushed.set(node.index, node)
-      }
-      throw err
-    } finally {
-      this.flushing = null
-    }
-  }
-
-  _flushTruncation () {
-    return new Promise((resolve, reject) => {
-      const t = this.truncateTo
-      const offset = t === 0 ? 0 : (t - 1) * 80 + 40
-
-      this.storage.truncate(offset, (err) => {
-        if (err) return reject(err)
-
-        if (this.truncateTo === t) {
-          this.truncateTo = 0
-          this.truncated = false
-        }
-
-        resolve()
-      })
-    })
+  seek (session, bytes, padding) {
+    return new ByteSeeker(this, session, bytes, padding)
   }
 
-  _flushNodes () {
-    // TODO: write neighbors together etc etc
-    // TODO: bench loading a full disk page and copy to that instead
-    return new Promise((resolve, reject) => {
-      const slab = b4a.allocUnsafe(40 * this.flushing.size)
+  get (index, error = true, readBatch = null) {
+    if (readBatch) return readBatch.getTreeNode(index, error)
 
-      let error = null
-      let missing = this.flushing.size + 1
-      let offset = 0
-
-      for (const node of this.flushing.values()) {
-        const state = {
-          start: 0,
-          end: 40,
-          buffer: slab.subarray(offset, offset += 40)
-        }
-
-        c.uint64.encode(state, node.size)
-        c.raw.encode(state, node.hash)
-
-        this.storage.write(node.index * 40, state.buffer, done)
-      }
-
-      done(null)
-
-      function done (err) {
-        if (err) error = err
-        if (--missing > 0) return
-        if (error) reject(error)
-        else resolve()
-      }
-    })
+    return getTreeNode(this.storage, index, error)
   }
 
-  clear () {
-    this.cache = new Xache({ maxSize: this.cache.maxSize })
+  clear (tx) {
     this.truncated = true
     this.truncateTo = 0
-    this.roots = []
-    this.length = 0
-    this.byteLength = 0
-    this.fork = 0
-    this.signature = null
-    if (this.flushing !== null) this.flushing.clear()
-    this.unflushed.clear()
-    return this.flush()
-  }
 
-  close () {
-    return new Promise((resolve, reject) => {
-      this.storage.close(err => {
-        if (err) reject(err)
-        else resolve()
-      })
-    })
+    return tx.deleteTreeNodeRange(0, -1)
   }
 
-  async truncate (length, fork = this.fork) {
+  async truncate (length, batch, fork = batch.fork) {
     const head = length * 2
-    const batch = new MerkleTreeBatch(this)
     const fullRoots = flat.fullRoots(head)
 
     for (let i = 0; i < fullRoots.length; i++) {
@@ -687,13 +629,11 @@ module.exports = class MerkleTree {
     return batch
   }
 
-  async reorg (proof) {
-    const batch = new ReorgBatch(this)
-
+  async reorg (proof, batch) {
     let unverified = null
 
     if (proof.block || proof.hash || proof.seek) {
-      unverified = verifyTree(proof, this.crypto, batch.nodes)
+      unverified = verifyTree(proof, batch.nodes)
     }
 
     if (!verifyUpgrade(proof, unverified, batch)) {
@@ -717,9 +657,9 @@ module.exports = class MerkleTree {
     return batch
   }
 
-  verifyFullyRemote (proof) {
+  verifyFullyRemote (proof, session) {
     // TODO: impl this less hackishly
-    const batch = new MerkleTreeBatch(this)
+    const batch = new MerkleTreeBatch(this, session)
 
     batch.fork = proof.fork
     batch.roots = []
@@ -727,7 +667,7 @@ module.exports = class MerkleTree {
     batch.ancestors = 0
     batch.byteLength = 0
 
-    let unverified = verifyTree(proof, this.crypto, batch.nodes)
+    let unverified = verifyTree(proof, batch.nodes)
 
     if (proof.upgrade) {
       if (verifyUpgrade(proof, unverified, batch)) {
@@ -738,10 +678,10 @@ module.exports = class MerkleTree {
     return batch
   }
 
-  async verify (proof) {
-    const batch = new MerkleTreeBatch(this)
+  async verify (proof, session) {
+    const batch = new MerkleTreeBatch(this, session)
 
-    let unverified = verifyTree(proof, this.crypto, batch.nodes)
+    let unverified = verifyTree(proof, batch.nodes)
 
     if (proof.upgrade) {
       if (verifyUpgrade(proof, unverified, batch)) {
@@ -759,13 +699,13 @@ module.exports = class MerkleTree {
     return batch
   }
 
-  proof ({ block, hash, seek, upgrade }) {
-    return generateProof(this, block, hash, seek, upgrade)
+  proof (rx, batch, { block, hash, seek, upgrade }) {
+    return generateProof(rx, batch, block, hash, seek, upgrade)
   }
 
   // Successor to .nodes()
-  async missingNodes (index) {
-    const head = 2 * this.length
+  async missingNodes (index, length) {
+    const head = 2 * length
     const ite = flat.iterator(index)
 
     // See iterator.rightSpan()
@@ -774,7 +714,8 @@ module.exports = class MerkleTree {
     if (iteRightSpan >= head) return 0
 
     let cnt = 0
-    while (!ite.contains(head) && (await this.get(ite.index, false)) === null) {
+    // TODO: we could prop use a read batch here and do this in blocks of X for perf
+    while (!ite.contains(head) && !(await hasTreeNode(this.storage, ite.index))) {
       cnt++
       ite.parent()
     }
@@ -796,56 +737,38 @@ module.exports = class MerkleTree {
     return cnt
   }
 
-  byteRange (index) {
-    return getByteRange(this, index)
-  }
-
-  byteOffset (index) {
-    return getByteOffset(this, index)
-  }
-
-  static async open (storage, opts = {}) {
-    await new Promise((resolve, reject) => {
-      storage.read(0, OLD_TREE.length, (err, buf) => {
-        if (err) return resolve()
-        if (b4a.equals(buf, OLD_TREE)) return reject(new Error('Storage contains an incompatible merkle tree'))
-        resolve()
-      })
-    })
-
-    const length = typeof opts.length === 'number'
-      ? opts.length
-      : await autoLength(storage)
-
+  static async open (storage, length, opts = {}) {
     const roots = []
     for (const index of flat.fullRoots(2 * length)) {
-      roots.push(unslabNode(await getStoredNode(storage, index, null, true)))
+      roots.push(unslabNode(await getTreeNode(storage, index, true)))
     }
 
     return new MerkleTree(storage, roots, opts.fork || 0, opts.signature || null, opts.prologue || null)
   }
 }
 
-async function getByteRange (tree, index) {
-  const head = 2 * tree.length
-  if (((index & 1) === 0 ? index : flat.rightSpan(index)) >= head) {
-    throw BAD_ARGUMENT('Index is out of bounds')
-  }
-  return [await tree.byteOffset(index), (await tree.get(index)).size]
+module.exports = {
+  MerkleTreeBatch,
+  ReorgBatch,
+  MerkleTree
 }
 
-async function getByteOffset (tree, index) {
-  if (index === 2 * tree.length) return tree.byteLength
+async function getNodeSize (index, readBatch) {
+  return (await readBatch.getTreeNode(index, true)).size
+}
+
+function getByteOffsetTree (tree, index, readBatch) {
   if ((index & 1) === 1) index = flat.leftSpan(index)
 
   let head = 0
-  let offset = 0
+
+  const promises = []
 
   for (const node of tree.roots) { // all async ticks happen once we find the root so safe
     head += 2 * ((node.index - head) + 1)
 
     if (index >= head) {
-      offset += node.size
+      promises.push(node.size)
       continue
     }
 
@@ -855,20 +778,43 @@ async function getByteOffset (tree, index) {
       if (index < ite.index) {
         ite.leftChild()
       } else {
-        offset += (await tree.get(ite.leftChild())).size
+        promises.push(tree.get(ite.leftChild(), true, readBatch))
         ite.sibling()
       }
     }
 
-    return offset
+    return Promise.all(promises)
   }
 
   throw ASSERTION('Failed to find offset')
 }
 
+async function getByteOffset (tree, index, readBatch) {
+  if (index === 2 * tree.length) return tree.byteLength
+
+  const treeNodes = await getByteOffsetTree(tree, index, readBatch)
+
+  let offset = 0
+  for (const node of treeNodes) offset += node.size
+
+  return offset
+}
+
+function getByteRange (tree, index, readBatch) {
+  const head = 2 * tree.length
+  if (((index & 1) === 0 ? index : flat.rightSpan(index)) >= head) {
+    throw BAD_ARGUMENT('Index is out of bounds')
+  }
+
+  const offset = getByteOffset(tree, index, readBatch)
+  const size = getNodeSize(index, readBatch)
+
+  return Promise.all([offset, size])
+}
+
 // All the methods needed for proof verification
 
-function verifyTree ({ block, hash, seek }, crypto, nodes) {
+function verifyTree ({ block, hash, seek }, nodes) {
   const untrustedNode = block
     ? { index: 2 * block.index, value: block.value, nodes: block.nodes }
     : hash
@@ -889,7 +835,7 @@ function verifyTree ({ block, hash, seek }, crypto, nodes) {
     while (q.length > 0) {
       const node = q.shift(ite.sibling())
 
-      root = parentNode(crypto, ite.parent(), root, node)
+      root = parentNode(ite.parent(), root, node)
       nodes.push(node)
       nodes.push(root)
     }
@@ -898,7 +844,7 @@ function verifyTree ({ block, hash, seek }, crypto, nodes) {
   if (untrustedNode === null) return root
 
   const ite = flat.iterator(untrustedNode.index)
-  const blockHash = untrustedNode.value && blockNode(crypto, ite.index, untrustedNode.value)
+  const blockHash = untrustedNode.value && blockNode(ite.index, untrustedNode.value)
 
   const q = new NodeQueue(untrustedNode.nodes, root)
 
@@ -908,7 +854,7 @@ function verifyTree ({ block, hash, seek }, crypto, nodes) {
   while (q.length > 0) {
     const node = q.shift(ite.sibling())
 
-    root = parentNode(crypto, ite.parent(), root, node)
+    root = parentNode(ite.parent(), root, node)
     nodes.push(node)
     nodes.push(root)
   }
@@ -917,7 +863,7 @@ function verifyTree ({ block, hash, seek }, crypto, nodes) {
 }
 
 function verifyUpgrade ({ fork, upgrade }, blockRoot, batch) {
-  const prologue = batch.tree.prologue
+  const prologue = batch.prologue
 
   if (prologue) {
     const { start, length } = upgrade
@@ -993,7 +939,7 @@ async function seekFromHead (tree, head, bytes, padding) {
 
   for (let i = 0; i < roots.length; i++) {
     const root = roots[i]
-    const node = await tree.get(root)
+    const node = await tree.get(root, true)
     const size = getUnpaddedSize(node, padding, null)
 
     if (bytes === size) return root
@@ -1017,6 +963,7 @@ async function seekTrustedTree (tree, root, bytes, padding) {
 
   while ((ite.index & 1) !== 0) {
     const l = await tree.get(ite.leftChild(), false)
+
     if (l) {
       const size = getUnpaddedSize(l, padding, ite)
       if (size === bytes) return ite.index
@@ -1042,7 +989,7 @@ async function seekUntrustedTree (tree, root, bytes, padding) {
 
   bytes -= offset
 
-  const node = await tree.get(root)
+  const node = await tree.get(root, true)
 
   if (getUnpaddedSize(node, padding, null) <= bytes) throw INVALID_OPERATION('Invalid seek')
 
@@ -1053,41 +1000,41 @@ async function seekUntrustedTree (tree, root, bytes, padding) {
 // Note, that all these methods are sync as we can statically infer which nodes
 // are needed for the remote to verify given they arguments they passed us
 
-function seekProof (tree, seekRoot, root, p) {
+function seekProof (tree, batch, seekRoot, root, p) {
   const ite = flat.iterator(seekRoot)
 
   p.seek = []
-  p.seek.push(tree.get(ite.index))
+  p.seek.push(tree.get(ite.index, true, batch))
 
   while (ite.index !== root) {
     ite.sibling()
-    p.seek.push(tree.get(ite.index))
+    p.seek.push(tree.get(ite.index, true, batch))
     ite.parent()
   }
 }
 
-function blockAndSeekProof (tree, node, seek, seekRoot, root, p) {
-  if (!node) return seekProof(tree, seekRoot, root, p)
+function blockAndSeekProof (tree, batch, node, seek, seekRoot, root, p) {
+  if (!node) return seekProof(tree, batch, seekRoot, root, p)
 
   const ite = flat.iterator(node.index)
 
   p.node = []
-  if (!node.value) p.node.push(tree.get(ite.index))
+  if (!node.value) p.node.push(tree.get(ite.index, true, batch))
 
   while (ite.index !== root) {
     ite.sibling()
 
     if (seek && ite.contains(seekRoot) && ite.index !== seekRoot) {
-      seekProof(tree, seekRoot, ite.index, p)
+      seekProof(tree, batch, seekRoot, ite.index, p)
     } else {
-      p.node.push(tree.get(ite.index))
+      p.node.push(tree.get(ite.index, true, batch))
     }
 
     ite.parent()
   }
 }
 
-function upgradeProof (tree, node, seek, from, to, subTree, p) {
+function upgradeProof (tree, batch, node, seek, from, to, subTree, p) {
   if (from === 0) p.upgrade = []
 
   for (const ite = flat.iterator(0); ite.fullRoot(to); ite.nextTree()) {
@@ -1107,9 +1054,9 @@ function upgradeProof (tree, node, seek, from, to, subTree, p) {
         ite.sibling()
         if (ite.index > target) {
           if (p.node === null && p.seek === null && ite.contains(subTree)) {
-            blockAndSeekProof(tree, node, seek, subTree, ite.index, p)
+            blockAndSeekProof(tree, batch, node, seek, subTree, ite.index, p)
           } else {
-            p.upgrade.push(tree.get(ite.index))
+            p.upgrade.push(batch.getTreeNode(ite.index, true))
           }
         }
         ite.parent()
@@ -1125,16 +1072,16 @@ function upgradeProof (tree, node, seek, from, to, subTree, p) {
     // if the subtree included is a child of this tree, include that one
     // instead of a dup node
     if (p.node === null && p.seek === null && ite.contains(subTree)) {
-      blockAndSeekProof(tree, node, seek, subTree, ite.index, p)
+      blockAndSeekProof(tree, batch, node, seek, subTree, ite.index, p)
       continue
     }
 
     // add root (can be optimised since the root might be in tree.roots)
-    p.upgrade.push(tree.get(ite.index))
+    p.upgrade.push(tree.get(ite.index, true, batch))
   }
 }
 
-function additionalUpgradeProof (tree, from, to, p) {
+function additionalUpgradeProof (tree, batch, from, to, p) {
   if (from === 0) p.additionalUpgrade = []
 
   for (const ite = flat.iterator(0); ite.fullRoot(to); ite.nextTree()) {
@@ -1153,7 +1100,7 @@ function additionalUpgradeProof (tree, from, to, p) {
       while (ite.index !== root) {
         ite.sibling()
         if (ite.index > target) {
-          p.additionalUpgrade.push(tree.get(ite.index))
+          p.additionalUpgrade.push(tree.get(ite.index, true, batch))
         }
         ite.parent()
       }
@@ -1166,7 +1113,7 @@ function additionalUpgradeProof (tree, from, to, p) {
     }
 
     // add root (can be optimised since the root is in tree.roots)
-    p.additionalUpgrade.push(tree.get(ite.index))
+    p.additionalUpgrade.push(tree.get(ite.index, true, batch))
   }
 }
 
@@ -1193,75 +1140,14 @@ function totalSpan (nodes) {
   return s
 }
 
-function blockNode (crypto, index, value) {
+function blockNode (index, value) {
   return { index, size: value.byteLength, hash: crypto.data(value) }
 }
 
-function parentNode (crypto, index, a, b) {
+function parentNode (index, a, b) {
   return { index, size: a.size + b.size, hash: crypto.parent(a, b) }
 }
 
-function blankNode (index) {
-  return { index, size: 0, hash: BLANK_HASH }
-}
-
-// Storage methods
-
-function getStoredNode (storage, index, cache, error) {
-  return new Promise((resolve, reject) => {
-    storage.read(40 * index, 40, (err, data) => {
-      if (err) {
-        if (error) return reject(err)
-        else resolve(null)
-        return
-      }
-
-      const hash = data.subarray(8)
-      const size = c.decode(c.uint64, data)
-
-      if (size === 0 && b4a.compare(hash, BLANK_HASH) === 0) {
-        if (error) reject(new Error('Could not load node: ' + index))
-        else resolve(null)
-        return
-      }
-
-      const node = { index, size, hash }
-
-      if (cache !== null) {
-        // Copy hash to a new buffer to avoid blocking gc of its original slab
-        node.hash = unslab(hash)
-        cache.set(index, node)
-      }
-
-      resolve(node)
-    })
-  })
-}
-
-function storedNodes (storage) {
-  return new Promise((resolve) => {
-    storage.stat((_, st) => {
-      if (!st) return resolve(0)
-      resolve((st.size - (st.size % 40)) / 40)
-    })
-  })
-}
-
-async function autoLength (storage) {
-  const nodes = await storedNodes(storage)
-  if (!nodes) return 0
-  const ite = flat.iterator(nodes - 1)
-  let index = nodes - 1
-  while (await getStoredNode(storage, ite.parent(), null, false)) index = ite.index
-  return flat.rightSpan(index) / 2 + 1
-}
-
-function truncateMap (map, len) {
-  for (const node of map.values()) {
-    if (node.index >= 2 * len) map.delete(node.index)
-  }
-}
-
 function log2 (n) {
   let res = 1
 
@@ -1279,6 +1165,20 @@ function normalizeIndexed (block, hash) {
   return null
 }
 
+function getTreeNode (storage, index, error) {
+  const batch = storage.read()
+  const node = batch.getTreeNode(index, error)
+  batch.tryFlush()
+  return node
+}
+
+function hasTreeNode (storage, index) {
+  const batch = storage.read()
+  const has = batch.hasTreeNode(index)
+  batch.tryFlush()
+  return has
+}
+
 async function settleProof (p) {
   const result = [
     p.node && Promise.all(p.node),
@@ -1299,7 +1199,7 @@ async function settleProof (p) {
 }
 
 // tree can be either the merkle tree or a merkle tree batch
-async function generateProof (tree, block, hash, seek, upgrade) {
+async function generateProof (readBatch, tree, block, hash, seek, upgrade) {
   // Important that this does not throw inbetween making the promise arrays
   // and finalise being called, otherwise there will be lingering promises in the background
 
@@ -1308,17 +1208,13 @@ async function generateProof (tree, block, hash, seek, upgrade) {
     upgrade.length = upgrade.start < tree.prologue.length ? tree.prologue.length : upgrade.length
   }
 
-  const fork = tree.fork
-  const signature = tree.signature
   const head = 2 * tree.length
   const from = upgrade ? upgrade.start * 2 : 0
   const to = upgrade ? from + upgrade.length * 2 : head
   const node = normalizeIndexed(block, hash)
 
-  const result = { fork, block: null, hash: null, seek: null, upgrade: null, manifest: null }
-
   // can't do anything as we have no data...
-  if (head === 0) return result
+  if (head === 0) return new TreeProof(tree, null, null, null, null)
 
   if (from >= to || to > head) {
     throw INVALID_OPERATION('Invalid upgrade')
@@ -1329,61 +1225,22 @@ async function generateProof (tree, block, hash, seek, upgrade) {
 
   let subTree = head
 
-  const p = {
-    node: null,
-    seek: null,
-    upgrade: null,
-    additionalUpgrade: null
-  }
+  const p = new TreeProof(tree, block, hash, seek, upgrade)
 
   if (node !== null && (!upgrade || node.lastIndex < upgrade.start)) {
     subTree = nodesToRoot(node.index, node.nodes, to)
     const seekRoot = seek ? await seekUntrustedTree(tree, subTree, seek.bytes, seek.padding) : head
-    blockAndSeekProof(tree, node, seek, seekRoot, subTree, p)
+    blockAndSeekProof(tree, readBatch, node, seek, seekRoot, subTree, p.pending)
   } else if ((node || seek) && upgrade) {
     subTree = seek ? await seekFromHead(tree, to, seek.bytes, seek.padding) : node.index
   }
 
   if (upgrade) {
-    upgradeProof(tree, node, seek, from, to, subTree, p)
-    if (head > to) additionalUpgradeProof(tree, to, head, p)
-  }
-
-  const [pNode, pSeek, pUpgrade, pAdditional] = await settleProof(p)
-
-  if (block) {
-    if (pNode === null) throw INVALID_OPERATION('Invalid block request')
-    result.block = {
-      index: block.index,
-      value: null, // populated upstream, alloc it here for simplicity
-      nodes: pNode
-    }
-  } else if (hash) {
-    if (pNode === null) throw INVALID_OPERATION('Invalid hash request')
-    result.hash = {
-      index: hash.index,
-      nodes: pNode
-    }
-  }
-
-  if (seek && pSeek !== null) {
-    result.seek = {
-      bytes: seek.bytes,
-      nodes: pSeek
-    }
-  }
-
-  if (upgrade) {
-    result.upgrade = {
-      start: upgrade.start,
-      length: upgrade.length,
-      nodes: pUpgrade,
-      additionalNodes: pAdditional || [],
-      signature
-    }
+    upgradeProof(tree, readBatch, node, seek, from, to, subTree, p.pending)
+    if (head > to) additionalUpgradeProof(tree, readBatch, to, head, p.pending)
   }
 
-  return result
+  return p
 }
 
 function getUnpaddedSize (node, padding, ite) {