Fix import_block and add state_root update (#241)

* Fix `state_root` setting

* Remove `parent_block_class` from `BeaconStateMachine.__init__` parameters

* Revert

* Test skipped slot

* PR feedback

* Refactor `create_block_on_state`

eth2/beacon/chains/base.py

@@ -252,11 +252,10 @@ def get_state_machine(self, at_block: BaseBeaconBlock=None) -> 'BaseBeaconStateM
         """
         block = self.ensure_block(at_block)
         sm_class = self.get_state_machine_class_for_block_slot(block.slot)
-        parent_block_class = self.get_block_class(block.parent_root)
+
         return sm_class(
             chaindb=self.chaindb,
             block=block,
-            parent_block_class=parent_block_class,
         )
 
     #
@@ -373,9 +372,7 @@ def import_block(
             parent_block,
             FromBlockParams(),
         )
-        state, imported_block = self.get_state_machine(
-            base_block_for_import,
-        ).import_block(block)
+        state, imported_block = self.get_state_machine(base_block_for_import).import_block(block)
 
         # TODO: Now it just persit all state. Should design how to clean up the old state.
         self.chaindb.persist_state(state)

eth2/beacon/constants.py

@@ -1,3 +1,7 @@
+from eth.constants import (
+    ZERO_HASH32,
+)
+
 from eth2.beacon.typing import (
     BLSSignature,
     SlotNumber,
@@ -20,3 +24,5 @@
 EMPTY_SIGNATURE = BLSSignature(b'\x00' * 96)
 GWEI_PER_ETH = 10**9
 FAR_FUTURE_SLOT = SlotNumber(2**64 - 1)
+
+GENESIS_PARENT_ROOT = ZERO_HASH32

eth2/beacon/db/chain.py

@@ -577,8 +577,7 @@ def _get_state_by_root(db: BaseDB, state_root: Hash32) -> BeaconState:
         try:
             state_rlp = db[state_root]
         except KeyError:
-            raise StateRootNotFound("No state with root {0} found".format(
-                encode_hex(state_rlp)))
+            raise StateRootNotFound(f"No state with root {encode_hex(state_root)} found")
         return _decode_state(state_rlp)
 
     def persist_state(self,

eth2/beacon/state_machines/base.py

@@ -72,7 +72,9 @@ def state_transition(self) -> BaseStateTransition:
     # Import block API
     #
     @abstractmethod
-    def import_block(self, block: BaseBeaconBlock) -> Tuple[BeaconState, BaseBeaconBlock]:
+    def import_block(self,
+                     block: BaseBeaconBlock,
+                     check_proposer_signature: bool=False) -> Tuple[BeaconState, BaseBeaconBlock]:
         pass
 
     @staticmethod
@@ -85,13 +87,9 @@ def create_block_from_parent(parent_block: BaseBeaconBlock,
 class BeaconStateMachine(BaseBeaconStateMachine):
     def __init__(self,
                  chaindb: BaseBeaconChainDB,
-                 block: BaseBeaconBlock,
-                 parent_block_class: Type[BaseBeaconBlock]) -> None:
+                 block: BaseBeaconBlock) -> None:
         self.chaindb = chaindb
-        self.block = self.get_block_class().from_parent(
-            parent_block=self.chaindb.get_block_by_root(block.parent_root, parent_block_class),
-            block_params=FromBlockParams(slot=block.slot),
-        )
+        self.block = block
 
     @property
     def state(self) -> BeaconState:
@@ -139,12 +137,17 @@ def state_transition(self) -> BaseStateTransition:
     #
     # Import block API
     #
-    def import_block(self, block: BaseBeaconBlock) -> Tuple[BeaconState, BaseBeaconBlock]:
+    def import_block(self,
+                     block: BaseBeaconBlock,
+                     check_proposer_signature: bool=False) -> Tuple[BeaconState, BaseBeaconBlock]:
         state = self.state_transition.apply_state_transition(
             self.state,
             block,
+            check_proposer_signature,
+        )
+
+        block = block.copy(
+            state_root=state.root,
         )
-        # TODO: Validate state roots
-        # TODO: Update self.state
-        # TODO: persist states in BeaconChain
+
         return state, block

eth2/beacon/state_machines/forks/serenity/state_transitions.py

@@ -28,11 +28,14 @@ class SerenityStateTransition(BaseStateTransition):
     def __init__(self, config: BeaconConfig):
         self.config = config
 
-    def apply_state_transition(self, state: BeaconState, block: BaseBeaconBlock) -> BeaconState:
+    def apply_state_transition(self,
+                               state: BeaconState,
+                               block: BaseBeaconBlock,
+                               check_proposer_signature: bool=False) -> BeaconState:
         while state.slot != block.slot:
             state = self.per_slot_transition(state, block.parent_root)
             if state.slot == block.slot:
-                state = self.per_block_transition(state, block)
+                state = self.per_block_transition(state, block, check_proposer_signature)
             if state.slot % self.config.EPOCH_LENGTH == 0:
                 state = self.per_epoch_transition(state, block)
 
@@ -87,17 +90,21 @@ def per_slot_transition(self,
         )
         return state
 
-    def per_block_transition(self, state: BeaconState, block: BaseBeaconBlock) -> BeaconState:
+    def per_block_transition(self,
+                             state: BeaconState,
+                             block: BaseBeaconBlock,
+                             check_proposer_signature: bool=False) -> BeaconState:
         # TODO: finish per-block processing logic as the spec
         validate_block_slot(state, block)
-        validate_proposer_signature(
-            state,
-            block,
-            beacon_chain_shard_number=self.config.BEACON_CHAIN_SHARD_NUMBER,
-            epoch_length=self.config.EPOCH_LENGTH,
-            target_committee_size=self.config.TARGET_COMMITTEE_SIZE,
-            shard_count=self.config.SHARD_COUNT
-        )
+        if not check_proposer_signature:
+            validate_proposer_signature(
+                state,
+                block,
+                beacon_chain_shard_number=self.config.BEACON_CHAIN_SHARD_NUMBER,
+                epoch_length=self.config.EPOCH_LENGTH,
+                target_committee_size=self.config.TARGET_COMMITTEE_SIZE,
+                shard_count=self.config.SHARD_COUNT
+            )
         # TODO: state = process_randao(state, block, self.config)
         # TODO: state = process_eth1_data(state, block, self.config)
 

eth2/beacon/state_machines/state_transitions.py

@@ -23,7 +23,10 @@ def __init__(self, config: BeaconConfig):
         self.config = config
 
     @abstractmethod
-    def apply_state_transition(self, state: BeaconState, block: BaseBeaconBlock) -> BeaconState:
+    def apply_state_transition(self,
+                               state: BeaconState,
+                               block: BaseBeaconBlock,
+                               check_proposer_signature: bool=False) -> BeaconState:
         pass
 
     @abstractmethod
@@ -33,7 +36,10 @@ def per_slot_transition(self,
         pass
 
     @abstractmethod
-    def per_block_transition(self, state: BeaconState, block: BaseBeaconBlock) -> BeaconState:
+    def per_block_transition(self,
+                             state: BeaconState,
+                             block: BaseBeaconBlock,
+                             check_proposer_signature: bool=False) -> BeaconState:
         pass
 
     @abstractmethod

eth2/beacon/tools/builder/proposer.py

@@ -20,6 +20,9 @@
     get_domain,
 )
 
+from eth2.beacon.state_machines.base import (
+    BaseBeaconStateMachine,
+)
 from eth2.beacon.state_machines.configs import BeaconConfig
 
 from eth2.beacon.types.attestations import Attestation
@@ -34,22 +37,14 @@
     BLSPubkey,
     FromBlockParams,
     SlotNumber,
+    ValidatorIndex,
 )
 
 
-def create_block_on_state(
-        state: BeaconState,
-        config: BeaconConfig,
-        block_class: BaseBeaconBlock,
-        parent_block: BaseBeaconBlock,
-        slot: SlotNumber,
-        validator_index: int,
-        privkey: int,
-        attestations: Sequence[Attestation]):
-    """
-    Create a beacon block with the given parameters.
-    """
-    # Check proposer
+def validate_proposer_index(state: BeaconState,
+                            config: BeaconConfig,
+                            slot: SlotNumber,
+                            validator_index: ValidatorIndex):
     beacon_proposer_index = get_beacon_proposer_index(
         state.copy(
             slot=slot,
@@ -63,6 +58,26 @@ def create_block_on_state(
     if validator_index != beacon_proposer_index:
         raise ProposerIndexError
 
+
+def create_block_on_state(
+        *,
+        state: BeaconState,
+        config: BeaconConfig,
+        state_machine: BaseBeaconStateMachine,
+        block_class: BaseBeaconBlock,
+        parent_block: BaseBeaconBlock,
+        slot: SlotNumber,
+        validator_index: ValidatorIndex,
+        privkey: int,
+        attestations: Sequence[Attestation],
+        check_proposer_index: bool=True) -> BaseBeaconBlock:
+    """
+    Create a beacon block with the given parameters.
+    """
+    # Check proposer
+    if check_proposer_index:
+        validate_proposer_index(state, config, slot, validator_index)
+
     # Prepare block: slot and parent_root
     block = block_class.from_parent(
         parent_block=parent_block,
@@ -82,6 +97,9 @@ def create_block_on_state(
         body=body,
     )
 
+    # Apply state transition to get state root
+    state, block = state_machine.import_block(block, check_proposer_signature=True)
+
     # Sign
     empty_signature_block_root = block.block_without_signature_root
     proposal_root = ProposalSignedData(
@@ -105,15 +123,19 @@ def create_block_on_state(
     return block
 
 
-def create_mock_block(state: BeaconState,
+def create_mock_block(*,
+                      state: BeaconState,
                       config: BeaconConfig,
+                      state_machine: BaseBeaconStateMachine,
                       block_class: Type[BaseBeaconBlock],
                       parent_block: BaseBeaconBlock,
                       keymap: Dict[BLSPubkey, int],
                       slot: SlotNumber=None,
                       attestations: Sequence[Attestation]=()) -> BaseBeaconBlock:
     """
     Create a mocking block with the given block parameters and ``keymap``.
+
+    Note that it doesn't return the correct ``state_root``.
     """
     proposer_index = get_beacon_proposer_index(
         state.copy(
@@ -127,15 +149,17 @@ def create_mock_block(state: BeaconState,
     proposer_pubkey = state.validator_registry[proposer_index].pubkey
     proposer_privkey = keymap[proposer_pubkey]
 
-    block = create_block_on_state(
-        state,
-        config,
-        block_class,
-        parent_block,
-        slot,
+    result_block = create_block_on_state(
+        state=state,
+        config=config,
+        state_machine=state_machine,
+        block_class=block_class,
+        parent_block=parent_block,
+        slot=slot,
         validator_index=proposer_index,
         privkey=proposer_privkey,
         attestations=attestations,
+        check_proposer_index=False,
     )
 
-    return block
+    return result_block

eth2/beacon/types/states.py

@@ -92,7 +92,7 @@ class BeaconState(rlp.Serializable):
         ('latest_index_roots', CountableList(hash32)),
         ('latest_penalized_balances', CountableList(uint64)),  # Balances penalized at every withdrawal period  # noqa: E501
         ('latest_attestations', CountableList(PendingAttestationRecord)),
-        ('batched_block_roots', CountableList(Hash32)),  # allow for a log-sized Merkle proof from any block to any historical block root"  # noqa: E501
+        ('batched_block_roots', CountableList(hash32)),  # allow for a log-sized Merkle proof from any block to any historical block root"  # noqa: E501
 
         # Ethereum 1.0 chain
         ('latest_eth1_data', Eth1Data),

tests/eth2/beacon/chains/test_chain.py

@@ -67,7 +67,7 @@ def test_canonical_chain(valid_chain):
         'num_validators,epoch_length,target_committee_size,shard_count'
     ),
     [
-        (100, 20, 10, 10),
+        (100, 16, 10, 10),
     ]
 )
 def test_import_blocks(valid_chain,
@@ -76,20 +76,22 @@ def test_import_blocks(valid_chain,
                        config,
                        keymap):
     state = genesis_state
-    blocks = tuple()
-
+    blocks = (genesis_block,)
     valid_chain_2 = copy.deepcopy(valid_chain)
     for i in range(3):
         block = create_mock_block(
             state=state,
             config=config,
+            state_machine=valid_chain.get_state_machine(blocks[-1]),
             block_class=genesis_block.__class__,
-            parent_block=genesis_block,
+            parent_block=blocks[-1],
             keymap=keymap,
             slot=state.slot + 2,
         )
 
         valid_chain.import_block(block)
+        assert valid_chain.get_canonical_head() == block
+
         state = valid_chain.get_state_machine(block).state
 
         assert block == valid_chain.get_canonical_block_by_slot(
@@ -102,10 +104,11 @@ def test_import_blocks(valid_chain,
 
     assert valid_chain.get_canonical_head() != valid_chain_2.get_canonical_head()
 
-    for block in blocks:
+    for block in blocks[1:]:
         valid_chain_2.import_block(block)
 
     assert valid_chain.get_canonical_head() == valid_chain_2.get_canonical_head()
+    assert valid_chain.get_state_machine(blocks[-1]).state.slot != 0
     assert (
         valid_chain.get_state_machine(blocks[-1]).state ==
         valid_chain_2.get_state_machine(blocks[-1]).state

tests/eth2/beacon/state_machines/test_demo.py

@@ -47,17 +47,23 @@ def test_demo(base_db,
     chaindb.persist_state(genesis_state)
 
     state = genesis_state
+    block = genesis_block
 
     current_slot = 1
     chain_length = 3 * config.EPOCH_LENGTH
     attestations = ()
+    blocks = (block,)
     for current_slot in range(chain_length):
         # two epochs
         block = create_mock_block(
             state=state,
             config=config,
+            state_machine=fixture_sm_class(
+                chaindb,
+                blocks[-1],
+            ),
             block_class=SerenityBeaconBlock,
-            parent_block=genesis_block,
+            parent_block=block,
             keymap=keymap,
             slot=current_slot,
             attestations=attestations,
@@ -71,19 +77,14 @@ def test_demo(base_db,
         # Get state machine instance
         sm = fixture_sm_class(
             chaindb,
-            block,
-            parent_block_class=SerenityBeaconBlock,
+            blocks[-1],
         )
         state, _ = sm.import_block(block)
 
-        # TODO: move to chain level?
-        block = block.copy(
-            state_root=state.root,
-        )
-
         chaindb.persist_state(state)
         chaindb.persist_block(block, SerenityBeaconBlock)
 
+        blocks += (block,)
         if current_slot > config.MIN_ATTESTATION_INCLUSION_DELAY:
             attestation_slot = current_slot - config.MIN_ATTESTATION_INCLUSION_DELAY
             attestations = create_mock_signed_attestations_at_slot(