feat(dag-builder): improve DAGBuilder (#1225)

Author: glevco

hathor/dag_builder/artifacts.py

@@ -14,13 +14,16 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Iterator, NamedTuple
+from typing import TYPE_CHECKING, Iterator, NamedTuple, TypeVar
 
 from hathor.dag_builder.types import DAGNode
+from hathor.manager import HathorManager
 
 if TYPE_CHECKING:
     from hathor.transaction import BaseTransaction
 
+T = TypeVar('T', bound='BaseTransaction')
+
 
 class _Pair(NamedTuple):
     node: DAGNode
@@ -38,3 +41,37 @@ def __init__(self, items: Iterator[tuple[DAGNode, BaseTransaction]]) -> None:
             self.by_name[node.name] = p
 
         self.list: tuple[_Pair, ...] = tuple(v)
+        self._last_propagated: str | None = None
+
+    def get_typed_vertex(self, name: str, type_: type[T]) -> T:
+        """Get a vertex by name, asserting it is of the provided type."""
+        _, vertex = self.by_name[name]
+        assert isinstance(vertex, type_)
+        return vertex
+
+    def get_typed_vertices(self, names: list[str], type_: type[T]) -> list[T]:
+        """Get a list of vertices by name, asserting they are of the provided type."""
+        return [self.get_typed_vertex(name, type_) for name in names]
+
+    def propagate_with(self, manager: HathorManager, *, up_to: str | None = None) -> None:
+        """
+        Propagate vertices using the provided manager up to the provided node name, included.
+        Last propagation is preserved in memory so you can make a sequence of propagate_with().
+        """
+        found_begin = self._last_propagated is None
+        found_end = False
+
+        for node, vertex in self.list:
+            if found_begin:
+                assert manager.on_new_tx(vertex, fails_silently=False)
+                self._last_propagated = node.name
+
+            if node.name == self._last_propagated:
+                found_begin = True
+
+            if up_to and node.name == up_to:
+                found_end = True
+                break
+
+        assert found_begin, f'node "{self._last_propagated}" not found'
+        assert up_to is None or found_end, f'node "{up_to}" not found'

hathor/dag_builder/tokenizer.py

@@ -27,10 +27,18 @@
     a --> b --> c               # c is a parent of b which is a parent of a
     a.out[i] <<< b c d          # b, c, and d spend the i-th output of a
     a < b < c                   # a must be created before b and b must be created before c
-    a > b > c                   # a must be created after b and b must be creater after c
-    a.attr = value              # set value of attribute attr to a
+    a > b > c                   # a must be created after b and b must be created after c
+    a.attr1 = value             # set value of attribute attr to a
+    a.attr2 = "value"           # a string literal
+
+Special keywords:
+
+    b10 < dummy                 # `dummy` is a tx created automatically that spends genesis tokens and provides
+                                # outputs to txs defined by the user. It's usually useful to set it after some
+                                # block to pass the reward lock
 
 Special attributes:
+
     a.out[i] = 100 HTR          # set that the i-th output of a holds 100 HTR
     a.out[i] = 100 TOKEN        # set that the i-th output of a holds 100 TOKEN where TOKEN is a custom token
     a.weight = 50               # set vertex weight

hathor/dag_builder/vertex_exporter.py

@@ -291,6 +291,7 @@ def create_vertex(self, node: DAGNode) -> BaseTransaction:
         assert node.name not in self._vertices
         self._vertice_per_id[vertex.hash] = vertex
         self._vertices[node.name] = vertex
+        vertex.name = node.name
         return vertex
 
     def export(self) -> Iterator[tuple[DAGNode, BaseTransaction]]:

hathor/transaction/base_transaction.py

@@ -185,6 +185,9 @@ def __init__(
         self._hash: VertexId | None = hash  # Stored as bytes.
         self._static_metadata = None
 
+        # A name solely for debugging purposes.
+        self.name: str | None = None
+
     @classproperty
     def log(cls):
         """ This is a workaround because of a bug on structlog (or abc).

tests/consensus/test_first_block.py

@@ -1,3 +1,4 @@
+from hathor.transaction import Block, Transaction
 from tests import unittest
 
 
@@ -40,21 +41,11 @@ def test_first_block(self) -> None:
             b33 --> tx50
         """)
 
-        for node, vertex in artifacts.list:
-            self.manager.on_new_tx(vertex, fails_silently=False)
-
-        b31 = artifacts.by_name['b31'].vertex
-        b32 = artifacts.by_name['b32'].vertex
-        b33 = artifacts.by_name['b33'].vertex
-
-        tx10 = artifacts.by_name['tx10'].vertex
-        tx20 = artifacts.by_name['tx20'].vertex
-        tx30 = artifacts.by_name['tx30'].vertex
-        tx40 = artifacts.by_name['tx40'].vertex
-        tx41 = artifacts.by_name['tx41'].vertex
-        tx42 = artifacts.by_name['tx42'].vertex
-        tx43 = artifacts.by_name['tx43'].vertex
-        tx50 = artifacts.by_name['tx50'].vertex
+        artifacts.propagate_with(self.manager)
+
+        b31, b32, b33 = artifacts.get_typed_vertices(['b31', 'b32', 'b33'], Block)
+        txs = ['tx10', 'tx20', 'tx30', 'tx40', 'tx41', 'tx42', 'tx43', 'tx50']
+        tx10, tx20, tx30, tx40, tx41, tx42, tx43, tx50 = artifacts.get_typed_vertices(txs, Transaction)
 
         self.assertEqual(tx10.get_metadata().first_block, b31.hash)
 

tests/dag_builder/test_dag_builter.py renamed to tests/dag_builder/test_dag_builder.py

@@ -1,8 +1,9 @@
+from hathor.transaction import Block, Transaction
 from hathor.transaction.token_creation_tx import TokenCreationTransaction
 from tests import unittest
 
 
-class DAGCreatorTestCase(unittest.TestCase):
+class DAGBuilderTestCase(unittest.TestCase):
     def setUp(self):
         super().setUp()
 
@@ -26,14 +27,12 @@ def test_one_tx(self) -> None:
             b40 --> tx1
         """)
 
-        for node, vertex in artifacts.list:
-            self.manager.on_new_tx(vertex, fails_silently=False)
+        artifacts.propagate_with(self.manager)
 
         v_order = [node.name for node, _ in artifacts.list]
 
-        tx1 = artifacts.by_name['tx1'].vertex
-        b1 = artifacts.by_name['b1'].vertex
-        b40 = artifacts.by_name['b40'].vertex
+        b1, b40 = artifacts.get_typed_vertices(['b1', 'b40'], Block)
+        tx1 = artifacts.get_typed_vertex('tx1', Transaction)
 
         # blockchain genesis b[1..50]
         self.assertEqual(b1.parents[0], self._settings.GENESIS_BLOCK_HASH)
@@ -65,13 +64,11 @@ def test_weight(self) -> None:
             c1.weight = 80.6
         """)
 
-        for node, vertex in artifacts.list:
-            self.manager.on_new_tx(vertex, fails_silently=False)
+        artifacts.propagate_with(self.manager)
 
-        tx1 = artifacts.by_name['tx1'].vertex
-        tka = artifacts.by_name['TKA'].vertex
-        c1 = artifacts.by_name['c1'].vertex
-        b38 = artifacts.by_name['b38'].vertex
+        c1, b38 = artifacts.get_typed_vertices(['c1', 'b38'], Block)
+        tx1 = artifacts.get_typed_vertex('tx1', Transaction)
+        tka = artifacts.get_typed_vertex('TKA', TokenCreationTransaction)
 
         self.assertAlmostEqual(tka.weight, 31.8)
         self.assertAlmostEqual(tx1.weight, 25.2)
@@ -85,10 +82,9 @@ def test_spend_unspecified_utxo(self) -> None:
             tx1.out[0] <<< tx2
         """)
 
-        for node, vertex in artifacts.list:
-            self.manager.on_new_tx(vertex, fails_silently=False)
+        artifacts.propagate_with(self.manager)
 
-        tx1 = artifacts.by_name['tx1'].vertex
+        tx1 = artifacts.get_typed_vertex('tx1', Transaction)
         self.assertEqual(len(tx1.outputs), 1)
         # the default filler fills unspecified utxos with 1 HTR
         self.assertEqual(tx1.outputs[0].value, 1)
@@ -107,22 +103,11 @@ def test_block_parents(self) -> None:
             b36 --> tx4
         """)
 
-        for node, vertex in artifacts.list:
-            self.manager.on_new_tx(vertex, fails_silently=False)
+        artifacts.propagate_with(self.manager)
 
-        b0 = artifacts.by_name['b30'].vertex
-        b1 = artifacts.by_name['b31'].vertex
-        b2 = artifacts.by_name['b32'].vertex
-        b3 = artifacts.by_name['b33'].vertex
-        b4 = artifacts.by_name['b34'].vertex
-        b5 = artifacts.by_name['b35'].vertex
-        b6 = artifacts.by_name['b36'].vertex
-        b7 = artifacts.by_name['b37'].vertex
-
-        tx1 = artifacts.by_name['tx1'].vertex
-        tx2 = artifacts.by_name['tx2'].vertex
-        tx3 = artifacts.by_name['tx3'].vertex
-        tx4 = artifacts.by_name['tx4'].vertex
+        blocks = ['b30', 'b31', 'b32', 'b33', 'b34', 'b35', 'b36', 'b37']
+        b0, b1, b2, b3, b4, b5, b6, b7 = artifacts.get_typed_vertices(blocks, Block)
+        tx1, tx2, tx3, tx4 = artifacts.get_typed_vertices(['tx1', 'tx2', 'tx3', 'tx4'], Transaction)
 
         self.assertEqual(b2.parents[0], b1.hash)
         self.assertEqual(b3.parents[0], b2.hash)
@@ -149,14 +134,12 @@ def test_custom_token(self) -> None:
             b40 --> tx1
         """)
 
-        for node, vertex in artifacts.list:
-            self.manager.on_new_tx(vertex, fails_silently=False)
+        artifacts.propagate_with(self.manager)
 
-        tka = artifacts.by_name['TKA'].vertex
-        tx1 = artifacts.by_name['tx1'].vertex
+        tx1 = artifacts.get_typed_vertex('tx1', Transaction)
+        tka = artifacts.get_typed_vertex('TKA', TokenCreationTransaction)
 
         # TKA token creation transaction
-        self.assertIsInstance(tka, TokenCreationTransaction)
         self.assertEqual(tka.token_name, 'TKA')
         self.assertEqual(tka.token_symbol, 'TKA')
 
@@ -201,8 +184,7 @@ def test_big_dag(self) -> None:
             b16 < tx4
         """)
 
-        for node, vertex in artifacts.list:
-            self.manager.on_new_tx(vertex, fails_silently=False)
+        artifacts.propagate_with(self.manager)
 
     def test_no_hash_conflict(self) -> None:
         artifacts = self.dag_builder.build_from_str("""
@@ -212,9 +194,24 @@ def test_no_hash_conflict(self) -> None:
 
             tx10.out[0] <<< tx20 tx30 tx40
         """)
+        artifacts.propagate_with(self.manager)
+
+    def test_propagate_with(self) -> None:
+        tx_storage = self.manager.tx_storage
+        artifacts = self.dag_builder.build_from_str('''
+            blockchain genesis b[1..10]
+            b10 < dummy
+            tx1 <-- tx2
+        ''')
+
+        artifacts.propagate_with(self.manager, up_to='b5')
+        assert len(list(tx_storage.get_all_transactions())) == 8  # 3 genesis + 5 blocks
+
+        artifacts.propagate_with(self.manager, up_to='b10')
+        assert len(list(tx_storage.get_all_transactions())) == 13  # 3 genesis + 10 blocks
+
+        artifacts.propagate_with(self.manager, up_to='tx1')
+        assert len(list(tx_storage.get_all_transactions())) == 15  # 3 genesis + 10 blocks + dummy + tx1
 
-        for node, vertex in artifacts.list:
-            print()
-            print(node.name)
-            print()
-            self.manager.on_new_tx(vertex, fails_silently=False)
+        artifacts.propagate_with(self.manager)
+        assert len(list(tx_storage.get_all_transactions())) == 16  # 3 genesis + 10 blocks + dummy + tx1 + tx2