refactor(daa): externalize block dependencies

Author: glevco

hathor/daa.py

@@ -27,10 +27,12 @@
 
 from hathor.conf.settings import HathorSettings
 from hathor.profiler import get_cpu_profiler
-from hathor.util import iwindows
+from hathor.types import VertexId
+from hathor.util import iwindows, not_none
 
 if TYPE_CHECKING:
     from hathor.transaction import Block, Transaction
+    from hathor.transaction.storage.simple_memory_storage import SimpleMemoryStorage
 
 logger = get_logger()
 cpu = get_cpu_profiler()
@@ -57,17 +59,38 @@ def __init__(self, *, settings: HathorSettings, test_mode: TestMode = TestMode.D
         DifficultyAdjustmentAlgorithm.singleton = self
 
     @cpu.profiler(key=lambda _, block: 'calculate_block_difficulty!{}'.format(block.hash.hex()))
-    def calculate_block_difficulty(self, block: 'Block') -> float:
-        """ Calculate block weight according to the ascendents of `block`, using calculate_next_weight."""
+    def calculate_block_difficulty(self, block: 'Block', memory_storage: 'SimpleMemoryStorage') -> float:
+        """ Calculate block weight according to the ascendants of `block`, using calculate_next_weight."""
         if self.TEST_MODE & TestMode.TEST_BLOCK_WEIGHT:
             return 1.0
 
         if block.is_genesis:
             return self.MIN_BLOCK_WEIGHT
 
-        return self.calculate_next_weight(block.get_block_parent(), block.timestamp)
+        parent_block = memory_storage.get_parent_block(block)
 
-    def calculate_next_weight(self, parent_block: 'Block', timestamp: int) -> float:
+        return self.calculate_next_weight(parent_block, block.timestamp, memory_storage)
+
+    def calculate_N(self, parent_block: 'Block') -> int:
+        return min(2 * self._settings.BLOCK_DIFFICULTY_N_BLOCKS, parent_block.get_height() - 1)
+
+    def get_block_dependencies(self, block: 'Block') -> list[VertexId]:
+        parent_block = block.get_block_parent()
+        N = self.calculate_N(parent_block)
+        ids: list[VertexId] = [not_none(parent_block.hash)]
+
+        while len(ids) <= N + 1:
+            parent_block = parent_block.get_block_parent()
+            ids.append(not_none(parent_block.hash))
+
+        return ids
+
+    def calculate_next_weight(
+        self,
+        parent_block: 'Block',
+        timestamp: int,
+        memory_storage: Optional['SimpleMemoryStorage'] = None,
+    ) -> float:
         """ Calculate the next block weight, aka DAA/difficulty adjustment algorithm.
 
         The algorithm used is described in [RFC 22](https://gitlab.com/HathorNetwork/rfcs/merge_requests/22).
@@ -80,7 +103,7 @@ def calculate_next_weight(self, parent_block: 'Block', timestamp: int) -> float:
         from hathor.transaction import sum_weights
 
         root = parent_block
-        N = min(2 * self._settings.BLOCK_DIFFICULTY_N_BLOCKS, parent_block.get_height() - 1)
+        N = self.calculate_N(parent_block)
         K = N // 2
         T = self.AVG_TIME_BETWEEN_BLOCKS
         S = 5
@@ -90,7 +113,7 @@ def calculate_next_weight(self, parent_block: 'Block', timestamp: int) -> float:
         blocks: list['Block'] = []
         while len(blocks) < N + 1:
             blocks.append(root)
-            root = root.get_block_parent()
+            root = memory_storage.get_parent_block(root) if memory_storage else root.get_block_parent()
             assert root is not None
 
         # TODO: revise if this assertion can be safely removed

hathor/transaction/storage/simple_memory_storage.py

@@ -0,0 +1,65 @@
+#  Copyright 2023 Hathor Labs
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#  http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+from hathor.transaction import BaseTransaction, Block, Transaction
+from hathor.transaction.storage import TransactionStorage
+from hathor.transaction.storage.exceptions import TransactionDoesNotExist
+from hathor.types import VertexId
+
+
+class SimpleMemoryStorage:
+    __slots__ = ('_blocks', '_transactions',)
+
+    def __init__(self) -> None:
+        self._blocks: dict[VertexId, Block] = {}
+        self._transactions: dict[VertexId, Transaction] = {}
+
+    @property
+    def _vertices(self) -> dict[VertexId, BaseTransaction]:
+        return {**self._blocks, **self._transactions}
+
+    def get_block(self, block_id: VertexId) -> Block:
+        if block := self._blocks.get(block_id):
+            return block
+
+        raise TransactionDoesNotExist(f'Block "{block_id.hex()}" does not exist in this SimpleMemoryStorage.')
+
+    def get_transaction(self, tx_id: VertexId) -> Transaction:
+        if tx := self._transactions.get(tx_id):
+            return tx
+
+        raise TransactionDoesNotExist(f'Transaction "{tx_id.hex()}" does not exist in this SimpleMemoryStorage.')
+
+    def get_parent_block(self, block: Block) -> Block:
+        parent_hash = block.get_block_parent_hash()
+
+        return self.get_block(parent_hash)
+
+    def add_vertices_from_storage(self, storage: TransactionStorage, ids: list[VertexId]) -> None:
+        for vertex_id in ids:
+            self.add_vertex_from_storage(storage, vertex_id)
+
+    def add_vertex_from_storage(self, storage: TransactionStorage, vertex_id: VertexId) -> None:
+        if vertex_id in self._vertices:
+            return
+
+        vertex = storage.get_transaction(vertex_id)
+
+        match vertex:
+            case Block():
+                self._blocks[vertex_id] = vertex
+            case Transaction():
+                self._transactions[vertex_id] = vertex
+            case _:
+                raise NotImplementedError

hathor/verification/block_verifier.py

@@ -25,6 +25,8 @@
     TransactionDataError,
     WeightError,
 )
+from hathor.transaction.storage.simple_memory_storage import SimpleMemoryStorage
+from hathor.util import not_none
 
 
 class BlockVerifier:
@@ -51,7 +53,11 @@ def verify_height(self, block: Block) -> None:
 
     def verify_weight(self, block: Block) -> None:
         """Validate minimum block difficulty."""
-        min_block_weight = self._daa.calculate_block_difficulty(block)
+        memory_storage = SimpleMemoryStorage()
+        dependencies = self._daa.get_block_dependencies(block)
+        memory_storage.add_vertices_from_storage(not_none(block.storage), dependencies)
+
+        min_block_weight = self._daa.calculate_block_difficulty(block, memory_storage)
         if block.weight < min_block_weight - self._settings.WEIGHT_TOL:
             raise WeightError(f'Invalid new block {block.hash_hex}: weight ({block.weight}) is '
                               f'smaller than the minimum weight ({min_block_weight})')

tests/tx/test_genesis.py

@@ -1,3 +1,5 @@
+from unittest.mock import Mock
+
 from hathor.conf import HathorSettings
 from hathor.daa import DifficultyAdjustmentAlgorithm, TestMode
 from hathor.transaction.storage import TransactionMemoryStorage
@@ -72,9 +74,9 @@ def test_genesis_weight(self):
         # Validate the block and tx weight
         # in test mode weight is always 1
         self._daa.TEST_MODE = TestMode.TEST_ALL_WEIGHT
-        self.assertEqual(self._daa.calculate_block_difficulty(genesis_block), 1)
+        self.assertEqual(self._daa.calculate_block_difficulty(genesis_block, Mock()), 1)
         self.assertEqual(self._daa.minimum_tx_weight(genesis_tx), 1)
 
         self._daa.TEST_MODE = TestMode.DISABLED
-        self.assertEqual(self._daa.calculate_block_difficulty(genesis_block), genesis_block.weight)
+        self.assertEqual(self._daa.calculate_block_difficulty(genesis_block, Mock()), genesis_block.weight)
         self.assertEqual(self._daa.minimum_tx_weight(genesis_tx), genesis_tx.weight)