diff --git a/scribe/db/db.py b/scribe/db/db.py index af5a24d..6316eba 100644 --- a/scribe/db/db.py +++ b/scribe/db/db.py @@ -18,7 +18,7 @@ from scribe.schema.claim import guess_stream_type from scribe.schema.result import Censor from scribe.blockchain.transaction import TxInput from scribe.common import hash_to_hex_str, hash160, LRUCacheWithMetrics -from scribe.db.merkle import Merkle, MerkleCache +from scribe.db.merkle import Merkle, MerkleCache, FastMerkleCacheItem from scribe.db.common import ResolveResult, STREAM_TYPES, CLAIM_TYPES, ExpandedResolveResult, DBError, UTXO from scribe.db.prefixes import PendingActivationValue, ClaimTakeoverValue, ClaimToTXOValue, PrefixDB from scribe.db.prefixes import ACTIVATED_CLAIM_TXO_TYPE, ACTIVATED_SUPPORT_TXO_TYPE, EffectiveAmountKey diff --git a/scribe/db/merkle.py b/scribe/db/merkle.py index 072f810..c3a3a08 100644 --- a/scribe/db/merkle.py +++ b/scribe/db/merkle.py @@ -276,3 +276,30 @@ class MerkleCache: level = await self._level_for(length) return self.merkle.branch_and_root_from_level( level, leaf_hashes, index, self.depth_higher) + + +class FastMerkleCacheItem: + __slots__ = ['tx_hashes', 'tree', 'root_hash'] + + def __init__(self, tx_hashes: typing.List[bytes]): + self.tx_hashes = tuple(tx_hashes) + self.tree: typing.List[typing.List[bytes]] = [] + self.root_hash = self._walk_merkle(tx_hashes, self.tree.append) + + @staticmethod + def _walk_merkle(items: typing.List[bytes], append_layer) -> bytes: + if len(items) == 1: + return items[0] + append_layer(items) + layer = [ + double_sha256(items[index] + items[index]) + if index + 1 == len(items) else double_sha256(items[index] + items[index + 1]) + for index in range(0, len(items), 2) + ] + return FastMerkleCacheItem._walk_merkle(layer, append_layer) + + def branch(self, tx_position: int) -> typing.List[str]: + return [ + (layer[-1] if (tx_position >> shift) ^ 1 == len(layer) else layer[(tx_position >> shift) ^ 1])[::-1].hex() + for shift, layer in enumerate(self.tree) + ]