batched merkle proofs

scribe/db/db.py

@@ -1020,29 +1020,30 @@ class HubDB:
 
         if needed_confirmed:
             needed_heights = set()
-            tx_heights_and_positions = {}
+            tx_heights_and_positions = defaultdict(list)
             for (tx_hash_bytes, tx_num), tx in zip(needed_confirmed, await run_in_executor(
                     self._executor, self.prefix_db.tx.multi_get, [(tx_hash,) for tx_hash, _ in needed_confirmed],
                     True, False)):
                 tx_height = bisect_right(self.tx_counts, tx_num)
                 needed_heights.add(tx_height)
                 tx_pos = tx_num - self.tx_counts[tx_height - 1]
-                tx_heights_and_positions[tx_hash_bytes] = (tx, tx_num, tx_height, tx_pos)
+                tx_heights_and_positions[tx_height].append((tx_hash_bytes, tx, tx_num, tx_pos))
 
             sorted_heights = list(sorted(needed_heights))
             block_txs = await run_in_executor(
                 self._executor, self.prefix_db.block_txs.multi_get, [(height,) for height in sorted_heights]
             )
             block_txs = {height: v.tx_hashes for height, v in zip(sorted_heights, block_txs)}
-            for tx_hash_bytes, (tx, tx_num, tx_height, tx_pos) in tx_heights_and_positions.items():
-                branch, root = self.merkle.branch_and_root(
-                    block_txs[tx_height], tx_pos
+            for tx_height, v in tx_heights_and_positions.items():
+                branches, root = self.merkle.branches_and_root(
+                    block_txs[tx_height], [tx_pos for (tx_hash_bytes, tx, tx_num, tx_pos) in v]
                 )
+                for (tx_hash_bytes, tx, tx_num, tx_pos) in v:
                     merkle = {
                         'block_height': tx_height,
                         'merkle': [
                             hash_to_hex_str(_hash)
-                        for _hash in branch
+                            for _hash in branches[tx_pos]
                         ],
                         'pos': tx_pos
                     }

scribe/db/merkle.py

@@ -25,7 +25,7 @@
 # and warranty status of this software.
 
 """Merkle trees, branches, proofs and roots."""
-
+import typing
 from asyncio import Event
 from math import ceil, log
 
@@ -87,6 +87,26 @@ class Merkle:
 
         return branch, hashes[0]
 
+    @staticmethod
+    def branches_and_root(block_tx_hashes: typing.List[bytes], tx_positions: typing.List[int]):
+        block_tx_hashes = list(block_tx_hashes)
+        positions = list(tx_positions)
+        length = ceil(log(len(block_tx_hashes), 2))
+        branches = [[] for _ in range(len(tx_positions))]
+        for _ in range(length):
+            if len(block_tx_hashes) & 1:
+                h = block_tx_hashes[-1]
+                block_tx_hashes.append(h)
+            for idx, tx_position in enumerate(tx_positions):
+                h = block_tx_hashes[tx_position ^ 1]
+                branches[idx].append(h)
+                tx_positions[idx] >>= 1
+            block_tx_hashes = [
+                double_sha256(block_tx_hashes[n] + block_tx_hashes[n + 1]) for n in
+                range(0, len(block_tx_hashes), 2)
+            ]
+        return {tx_position: branch for tx_position, branch in zip(positions, branches)}, block_tx_hashes[0]
+
     @staticmethod
     def root(hashes, length=None):
         """Return the merkle root of a non-empty iterable of binary hashes."""