add storage

cryptarchia/cryptarchia.py

@@ -1,15 +1,15 @@
-from typing import TypeAlias, List, Dict
-from hashlib import sha256, blake2b
-from math import floor
-from copy import deepcopy
-import itertools
 import functools
-from dataclasses import dataclass, field, replace
+import itertools
 import logging
 from collections import defaultdict
+from copy import deepcopy
+from dataclasses import dataclass, field, replace
+from hashlib import blake2b, sha256
+from math import floor
+from typing import Dict, Generator, List, TypeAlias
 
 import numpy as np
-
+from sortedcontainers import SortedDict
 
 logger = logging.getLogger(__name__)
 
@@ -124,6 +124,9 @@ class Slot:
     def __lt__(self, other):
         return self.absolute_slot < other.absolute_slot
 
+    def __hash__(self):
+        return hash(self.absolute_slot)
+
 
 @dataclass
 class Coin:
@@ -367,6 +370,7 @@ class Follower:
         self.genesis_state = genesis_state
         self.ledger_state = {genesis_state.block.id(): genesis_state.copy()}
         self.epoch_state = {}
+        self.block_storage = BlockStorage()
 
     def validate_header(self, block: BlockHeader) -> bool:
         # TODO: verify blocks are not in the 'future'
@@ -471,6 +475,8 @@ class Follower:
             self.forks.remove(new_tip)
             self.local_chain = new_tip
 
+        self.block_storage.add_block(block)
+
     def unimported_orphans(self) -> list[BlockHeader]:
         """
         Returns all unimported orphans w.r.t. the given tip's state.
@@ -486,6 +492,8 @@ class Follower:
             for block_state in chain_suffix(fork, fork_depth, self.ledger_state):
                 b = block_state.block
                 if b.leader_proof.nullifier not in tip_state.nullifiers:
+                    # YJ: Why do this? This function is used only in tests.
+                    # Let's try to remove this line and run tests.
                     tip_state.nullifiers.add(b.leader_proof.nullifier)
                     orphans += [b]
 
@@ -593,6 +601,26 @@ class Follower:
         return int(prev_epoch.inferred_total_active_stake - h * blocks_per_slot_err)
 
 
+class BlockStorage:
+    def __init__(self):
+        self.blocks: dict[Id, BlockHeader] = dict()
+        self.ids_by_slot: SortedDict[Slot, set[Id]] = SortedDict()
+
+    def add_block(self, block: BlockHeader):
+        id = block.id()
+        self.blocks[id] = block
+        if block.slot not in self.ids_by_slot:
+            self.ids_by_slot[block.slot] = set()
+        self.ids_by_slot[block.slot].add(id)
+
+    def blocks_by_range(
+        self, from_slot: Slot, to_slot: Slot
+    ) -> Generator[BlockHeader, None, None]:
+        for slot in self.ids_by_slot.irange(from_slot, to_slot, inclusive=(True, True)):
+            for id in self.ids_by_slot[slot]:
+                yield self.blocks[id]
+
+
 def phi(f: float, alpha: float) -> float:
     """
     params:

cryptarchia/sync/full_sync.py

@@ -1,7 +1,6 @@
 from collections import defaultdict
-from typing import Generator
 
-from cryptarchia.cryptarchia import BlockHeader, Follower, Id, Slot
+from cryptarchia.cryptarchia import Follower, Id, Slot
 
 SLOT_TOLERANCE = 1
 
@@ -15,7 +14,7 @@ def full_sync(local: Follower, remotes: list[Follower], start_slot: Slot):
 
 
 def range_sync(local: Follower, remote: Follower, from_slot: Slot, to_slot: Slot):
-    for block in request_blocks_by_range(remote, from_slot, to_slot):
+    for block in remote.block_storage.blocks_by_range(from_slot, to_slot):
         local.on_block(block)
 
 
@@ -27,18 +26,3 @@ def group_targets(
         if target.tip().slot.absolute_slot - start_slot.absolute_slot > SLOT_TOLERANCE:
             groups[target.tip_id()].append(target)
     return groups
-
-
-def request_blocks_by_range(
-    remote: Follower, from_slot: Slot, to_slot: Slot
-) -> Generator[BlockHeader, None, None]:
-    # TODO: Optimize this by keeping blocks by slot in the Follower
-    blocks_by_slot: dict[int, list[BlockHeader]] = defaultdict(list)
-    for ledger_state in remote.ledger_state.values():
-        if from_slot <= ledger_state.block.slot <= to_slot:
-            blocks_by_slot[ledger_state.block.slot.absolute_slot].append(
-                ledger_state.block
-            )
-    for slot in range(from_slot.absolute_slot, to_slot.absolute_slot + 1):
-        for block in blocks_by_slot[slot]:
-            yield block

cryptarchia/sync/test_full_sync.py

@@ -75,8 +75,12 @@ class TestFullSync(TestCase):
 
         new_follower = Follower(genesis, config)
         full_sync(new_follower, [follower], genesis.block.slot)
-        assert new_follower.tip() == follower.tip()
-        assert new_follower.forks == follower.forks
+        # Since the length of two forks is the same, the tip is chosen
+        # depending on the order of block delivery.
+        assert new_follower.tip() in [b2, b4]
+        expected_forks = [b2.id(), b4.id()]
+        expected_forks.remove(new_follower.tip_id())
+        assert new_follower.forks == expected_forks
 
     def test_continue_syncing_forks(self):
         # b0 - b1 - b2 == tip
@@ -137,7 +141,7 @@ class TestFullSync(TestCase):
         b0, c_b = mk_block(genesis.block, 1, c_b), c_b.evolve()
         b3, c_b = mk_block(b0, 2, c_b), c_b.evolve()
         b4, c_b = mk_block(b3, 3, c_b), c_b.evolve()
-        b5, c_b = mk_block(b4, 3, c_b), c_b.evolve()
+        b5, c_b = mk_block(b4, 4, c_b), c_b.evolve()
         for b in [b0, b3, b4, b5]:
             peer_1.on_block(b)
         assert peer_1.tip() == b5

cryptarchia/test_ledger_state_update.py

@@ -2,9 +2,8 @@ from unittest import TestCase
 
 import numpy as np
 
-from .cryptarchia import Follower, Coin, iter_chain
-
-from .test_common import mk_config, mk_block, mk_genesis_state
+from .cryptarchia import Coin, Follower, iter_chain
+from .test_common import mk_block, mk_config, mk_genesis_state
 
 
 class TestLedgerStateUpdate(TestCase):
@@ -139,7 +138,7 @@ class TestLedgerStateUpdate(TestCase):
 
         follower = Follower(genesis, config)
 
-        # We assume an epoch length of 10 slots in this test.
+        # We assume an epoch length of 20 slots in this test.
         assert config.epoch_length == 20, f"epoch len: {config.epoch_length}"
 
         # ---- EPOCH 0 ----
@@ -164,14 +163,14 @@ class TestLedgerStateUpdate(TestCase):
         # ---- EPOCH 2 ----
 
         # when trying to propose a block for epoch 2, the stake distribution snapshot should be taken
-        # at the end of epoch 0, i.e. slot 9
+        # at the end of epoch 0, i.e. slot 19
         # To ensure this is the case, we add a new coin just to the state associated with that slot,
         # so that the new block can be accepted only if that is the snapshot used
         # first, verify that if we don't change the state, the block is not accepted
         block_4 = mk_block(slot=40, parent=block_3, coin=Coin(sk=4, value=100))
         follower.on_block(block_4)
         assert follower.tip() == block_3
-        # then we add the coin to "spendable commitments" associated with slot 9
+        # then we add the coin to "spendable commitments" associated with slot 19
         follower.ledger_state[block_2.id()].commitments_spend.add(
             Coin(sk=4, value=100).commitment()
         )

requirements.txt

@@ -13,3 +13,4 @@ portalocker==2.8.2 # portable file locking
 keum==0.2.0 # for CL's use of more obscure curves
 poseidon-hash==0.1.4 # used as the algebraic hash in CL
 hypothesis==6.103.0
+sortedcontainers==2.4.0