Allow a caller of persist_block to provide their own fork choice

eth2/beacon/chains/base.py

@@ -40,6 +40,9 @@
     BlockClassError,
     StateMachineNotFound,
 )
+from eth2.beacon.fork_choice import (
+    ForkChoiceScoring,
+)
 from eth2.beacon.types.attestations import (
     Attestation,
 )
@@ -226,18 +229,21 @@ def from_genesis(cls,
 
         chaindb = cls.get_chaindb_class()(db=base_db, genesis_config=genesis_config)
         chaindb.persist_state(genesis_state)
-        return cls._from_genesis_block(base_db, genesis_block, genesis_config)
+        state_machine = sm_class(chaindb, genesis_block, genesis_state)
+        fork_choice_scoring = state_machine.get_fork_choice_scoring()
+        return cls._from_genesis_block(base_db, genesis_block, fork_choice_scoring, genesis_config)
 
     @classmethod
     def _from_genesis_block(cls,
                             base_db: BaseAtomicDB,
                             genesis_block: BaseBeaconBlock,
+                            fork_choice_scoring: ForkChoiceScoring,
                             genesis_config: Eth2GenesisConfig) -> 'BaseBeaconChain':
         """
         Initialize the ``BeaconChain`` from the genesis block.
         """
         chaindb = cls.get_chaindb_class()(db=base_db, genesis_config=genesis_config)
-        chaindb.persist_block(genesis_block, genesis_block.__class__)
+        chaindb.persist_block(genesis_block, genesis_block.__class__, fork_choice_scoring)
         return cls(base_db, genesis_config)
 
     #
@@ -408,17 +414,15 @@ def import_block(
         # TODO: Now it just persists all state. Should design how to clean up the old state.
         self.chaindb.persist_state(state)
 
-        self.chaindb.persist_block_without_scoring(imported_block, imported_block.__class__)
-
         fork_choice_scoring = state_machine.get_fork_choice_scoring()
-        score = fork_choice_scoring(imported_block)
-
-        self.chaindb.set_score(imported_block, score)
-
         (
             new_canonical_blocks,
             old_canonical_blocks,
-        ) = self.chaindb.update_canonical_head_if_needed(imported_block, imported_block.__class__)
+        ) = self.chaindb.persist_block(
+            imported_block,
+            imported_block.__class__,
+            fork_choice_scoring,
+        )
 
         self.logger.debug(
             'IMPORTED_BLOCK: slot %s | signed root %s',

eth2/beacon/db/chain.py

@@ -39,6 +39,9 @@
 from eth.validation import (
     validate_word,
 )
+from eth2.beacon.fork_choice import (
+    ForkChoiceScoring,
+)
 from eth2.beacon.helpers import (
     slot_to_epoch,
 )
@@ -59,6 +62,7 @@
     AttestationRootNotFound,
     FinalizedHeadNotFound,
     JustifiedHeadNotFound,
+    MissingForkChoiceScorings,
 )
 from eth2.beacon.db.schema import SchemaV1
 
@@ -88,7 +92,8 @@ def __init__(self, db: BaseAtomicDB, genesis_config: Eth2GenesisConfig) -> None:
     def persist_block(
             self,
             block: BaseBeaconBlock,
-            block_class: Type[BaseBeaconBlock]
+            block_class: Type[BaseBeaconBlock],
+            fork_choice_scoring: ForkChoiceScoring,
     ) -> Tuple[Tuple[BaseBeaconBlock, ...], Tuple[BaseBeaconBlock, ...]]:
         pass
 
@@ -145,7 +150,8 @@ def block_exists(self, block_root: Hash32) -> bool:
     def persist_block_chain(
             self,
             blocks: Iterable[BaseBeaconBlock],
-            block_class: Type[BaseBeaconBlock]
+            block_class: Type[BaseBeaconBlock],
+            fork_choice_scoring: Iterable[ForkChoiceScoring],
     ) -> Tuple[Tuple[BaseBeaconBlock, ...], Tuple[BaseBeaconBlock, ...]]:
         pass
 
@@ -213,7 +219,8 @@ def _get_highest_justified_epoch(self, db: BaseDB) -> Epoch:
     def persist_block(
             self,
             block: BaseBeaconBlock,
-            block_class: Type[BaseBeaconBlock]
+            block_class: Type[BaseBeaconBlock],
+            fork_choice_scoring: ForkChoiceScoring,
     ) -> Tuple[Tuple[BaseBeaconBlock, ...], Tuple[BaseBeaconBlock, ...]]:
         """
         Persist the given block.
@@ -222,20 +229,23 @@ def persist_block(
             if block.is_genesis:
                 self._handle_exceptional_justification_and_finality(db, block)
 
-            return self._persist_block(db, block, block_class)
+            return self._persist_block(db, block, block_class, fork_choice_scoring)
 
     @classmethod
     def _persist_block(
             cls,
             db: 'BaseDB',
             block: BaseBeaconBlock,
-            block_class: Type[BaseBeaconBlock]
+            block_class: Type[BaseBeaconBlock],
+            fork_choice_scoring: ForkChoiceScoring,
     ) -> Tuple[Tuple[BaseBeaconBlock, ...], Tuple[BaseBeaconBlock, ...]]:
         block_chain = (block, )
+        scorings = (fork_choice_scoring, )
         new_canonical_blocks, old_canonical_blocks = cls._persist_block_chain(
             db,
             block_chain,
             block_class,
+            scorings,
         )
 
         return new_canonical_blocks, old_canonical_blocks
@@ -428,23 +438,22 @@ def _block_exists(db: BaseDB, block_root: Hash32) -> bool:
     def persist_block_chain(
             self,
             blocks: Iterable[BaseBeaconBlock],
-            block_class: Type[BaseBeaconBlock]
+            block_class: Type[BaseBeaconBlock],
+            fork_choice_scorings: Iterable[ForkChoiceScoring],
     ) -> Tuple[Tuple[BaseBeaconBlock, ...], Tuple[BaseBeaconBlock, ...]]:
         """
         Return two iterable of blocks, the first containing the new canonical blocks,
         the second containing the old canonical headers
         """
         with self.db.atomic_batch() as db:
-            return self._persist_block_chain(db, blocks, block_class)
+            return self._persist_block_chain(db, blocks, block_class, fork_choice_scorings)
 
     @staticmethod
     def _set_block_score_to_db(
             db: BaseDB,
-            block: BaseBeaconBlock
+            block: BaseBeaconBlock,
+            score: int,
     ) -> int:
-        # TODO: It's a stub before we implement fork choice rule
-        score = block.slot
-
         db.set(
             SchemaV1.make_block_root_to_score_lookup_key(block.signing_root),
             ssz.encode(score, sedes=ssz.sedes.uint64),
@@ -462,12 +471,15 @@ def _persist_block_chain(
             cls,
             db: BaseDB,
             blocks: Iterable[BaseBeaconBlock],
-            block_class: Type[BaseBeaconBlock]
+            block_class: Type[BaseBeaconBlock],
+            fork_choice_scorings: Iterable[ForkChoiceScoring],
     ) -> Tuple[Tuple[BaseBeaconBlock, ...], Tuple[BaseBeaconBlock, ...]]:
         blocks_iterator = iter(blocks)
+        scorings_iterator = iter(fork_choice_scorings)
 
         try:
             first_block = first(blocks_iterator)
+            first_scoring = first(scorings_iterator)
         except StopIteration:
             return tuple(), tuple()
 
@@ -488,15 +500,15 @@ def _persist_block_chain(
                 )
             )
 
-        score = first_block.slot
+        score = first_scoring(first_block)
 
         curr_block_head = first_block
         db.set(
             curr_block_head.signing_root,
             ssz.encode(curr_block_head),
         )
         cls._add_block_root_to_slot_lookup(db, curr_block_head)
-        cls._set_block_score_to_db(db, curr_block_head)
+        cls._set_block_score_to_db(db, curr_block_head, score)
         cls._add_attestations_root_to_block_lookup(db, curr_block_head)
 
         orig_blocks_seq = concat([(first_block,), blocks_iterator])
@@ -517,9 +529,17 @@ def _persist_block_chain(
                 ssz.encode(curr_block_head),
             )
             cls._add_block_root_to_slot_lookup(db, curr_block_head)
-            score = cls._set_block_score_to_db(db, curr_block_head)
             cls._add_attestations_root_to_block_lookup(db, curr_block_head)
 
+            # NOTE: len(scorings_iterator) should equal len(blocks_iterator)
+            try:
+                next_scoring = next(scorings_iterator)
+            except StopIteration:
+                raise MissingForkChoiceScorings
+
+            score = next_scoring(curr_block_head)
+            cls._set_block_score_to_db(db, curr_block_head, score)
+
         if no_canonical_head:
             return cls._set_as_canonical_chain_head(db, curr_block_head.signing_root, block_class)
 

eth2/beacon/db/exceptions.py

@@ -24,3 +24,11 @@ class AttestationRootNotFound(BeaconDBException):
     Exception raised if no attestation root is set in this database.
     """
     pass
+
+
+class MissingForkChoiceScorings(BeaconDBException):
+    """
+    Exception raised if a client tries to score a block without providing
+    the ability to generate a score via a ``scoring``.
+    """
+    pass

tests/eth2/core/beacon/chains/test_beacon_chain.py

@@ -42,7 +42,7 @@ def valid_chain(beacon_chain_with_block_validation):
         (100, 20, 10, 10),
     ]
 )
-def test_canonical_chain(valid_chain, genesis_slot):
+def test_canonical_chain(valid_chain, genesis_slot, fork_choice_scoring):
     genesis_block = valid_chain.get_canonical_block_by_slot(genesis_slot)
 
     # Our chain fixture is created with only the genesis header, so initially that's the head of
@@ -53,7 +53,7 @@ def test_canonical_chain(valid_chain, genesis_slot):
         slot=genesis_block.slot + 1,
         previous_block_root=genesis_block.signing_root,
     )
-    valid_chain.chaindb.persist_block(block, block.__class__)
+    valid_chain.chaindb.persist_block(block, block.__class__, fork_choice_scoring)
 
     assert valid_chain.get_canonical_head() == block
 

tests/eth2/core/beacon/conftest.py

@@ -16,6 +16,9 @@
 from eth2.beacon.constants import (
     FAR_FUTURE_EPOCH,
 )
+from eth2.beacon.fork_choice import (
+    higher_slot_scoring,
+)
 from eth2.beacon.helpers import (
     slot_to_epoch,
 )
@@ -803,6 +806,11 @@ def fixture_sm_class(config):
     )
 
 
+@pytest.fixture
+def fork_choice_scoring():
+    return higher_slot_scoring
+
+
 @pytest.fixture
 def genesis_config(config):
     return Eth2GenesisConfig(config)