modify fork choice to utilize epochs properly

specs/core/0_fork-choice.md

@@ -12,7 +12,7 @@
         - [Time parameters](#time-parameters)
     - [Fork choice](#fork-choice)
         - [Helpers](#helpers)
-            - [`Target`](#target)
+            - [`Checkpoint`](#checkpoint)
             - [`Store`](#store)
             - [`get_genesis_store`](#get_genesis_store)
             - [`get_ancestor`](#get_ancestor)
@@ -55,26 +55,37 @@ The head block root associated with a `store` is defined as `get_head(store)`. A
 
 ### Helpers
 
-#### `Target`
+#### `Checkpoint`
 
 ```python
-@dataclass
-class Target(object):
+@dataclass(eq=True, frozen=True)
+class Checkpoint(object):
     epoch: Epoch
     root: Hash
 ```
 
+#### `LatestMessage`
+
+```python
+@dataclass(eq=True, frozen=True)
+class LatestMessage(object):
+    epoch: Epoch
+    root: Hash
+```
+
+
 #### `Store`
 
 ```python
 @dataclass
 class Store(object):
+    time: int
+    justified_checkpoint: Checkpoint
+    finalized_checkpoint: Checkpoint
     blocks: Dict[Hash, BeaconBlock] = field(default_factory=dict)
-    states: Dict[Hash, BeaconState] = field(default_factory=dict)
-    time: int = 0
-    latest_targets: Dict[ValidatorIndex, Target] = field(default_factory=dict)
-    justified_root: Hash = ZERO_HASH
-    finalized_root: Hash = ZERO_HASH
+    block_states: Dict[Hash, BeaconState] = field(default_factory=dict)
+    checkpoint_states: Dict[Checkpoint, BeaconState] = field(default_factory=dict)
+    latest_messages: Dict[ValidatorIndex, LatestMessage] = field(default_factory=dict)
 ```
 
 #### `get_genesis_store`
@@ -83,12 +94,15 @@ class Store(object):
 def get_genesis_store(genesis_state: BeaconState) -> Store:
     genesis_block = BeaconBlock(state_root=hash_tree_root(genesis_state))
     root = signing_root(genesis_block)
+    justified_checkpoint = Checkpoint(GENESIS_EPOCH, root)
+    finalized_checkpoint = Checkpoint(GENESIS_EPOCH, root)
     return Store(
-        blocks={root: genesis_block},
-        states={root: genesis_state},
         time=genesis_state.genesis_time,
-        justified_root=root,
-        finalized_root=root,
+        justified_checkpoint=justified_checkpoint,
+        finalized_checkpoint=finalized_checkpoint,
+        blocks={root: genesis_block},
+        block_states={root: genesis_state.copy()},
+        checkpoint_states={justified_checkpoint: genesis_state.copy()},
     )
 ```
 
@@ -105,11 +119,12 @@ def get_ancestor(store: Store, root: Hash, slot: Slot) -> Hash:
 
 ```python
 def get_latest_attesting_balance(store: Store, root: Hash) -> Gwei:
-    state = store.states[store.justified_root]
-    active_indices = get_active_validator_indices(state.validator_registry, get_current_epoch(state))
+    state = store.checkpoint_states[store.justified_checkpoint]
+    active_indices = get_active_validator_indices(state, get_current_epoch(state))
     return Gwei(sum(
-        state.validator_registry[i].effective_balance for i in active_indices
-        if get_ancestor(store, store.latest_targets[i].root, store.blocks[root].slot) == root
+        state.validators[i].effective_balance for i in active_indices
+        if (i in store.latest_messages and
+            get_ancestor(store, store.latest_messages[i].root, store.blocks[root].slot) == root)
     ))
 ```
 
@@ -118,9 +133,13 @@ def get_latest_attesting_balance(store: Store, root: Hash) -> Gwei:
 ```python
 def get_head(store: Store) -> Hash:
     # Execute the LMD-GHOST fork choice
-    head = store.justified_root
+    head = store.justified_checkpoint.root
+    justified_slot = get_epoch_start_slot(store.justified_checkpoint.epoch)
     while True:
-        children = [root for root in store.blocks.keys() if store.blocks[root].parent_root == head]
+        children = [
+            root for root in store.blocks.keys()
+            if store.blocks[root].parent_root == head and store.blocks[root].slot > justified_slot
+        ]
         if len(children) == 0:
             return head
         # Sort by latest attesting balance with ties broken lexicographically
@@ -141,35 +160,65 @@ def on_tick(store: Store, time: int) -> None:
 ```python
 def on_block(store: Store, block: BeaconBlock) -> None:
     # Make a copy of the state to avoid mutability issues
-    pre_state = store.states[block.parent_root].copy()
+    assert block.parent_root in store.block_states
+    pre_state = store.block_states[block.parent_root].copy()
     # Blocks cannot be in the future. If they are, their consideration must be delayed until the are in the past.
     assert store.time >= pre_state.genesis_time + block.slot * SECONDS_PER_SLOT
     # Add new block to the store
     store.blocks[signing_root(block)] = block
     # Check block is a descendant of the finalized block
-    assert get_ancestor(store, signing_root(block), store.blocks[store.finalized_root].slot) == store.finalized_root
+    assert (
+        get_ancestor(store, signing_root(block), store.blocks[store.finalized_checkpoint.root].slot) ==
+        store.finalized_checkpoint.root
+    )
+    # Check that block is later than the finalized epoch slot
+    assert block.slot > get_epoch_start_slot(store.finalized_checkpoint.epoch)
     # Check the block is valid and compute the post-state
     state = state_transition(pre_state, block)
-    # Add new state to the store
-    store.states[signing_root(block)] = state
-    # Update justified block root
-    if state.current_justified_epoch > slot_to_epoch(store.blocks[store.justified_root].slot):
-        store.justified_root = state.current_justified_root
-    elif state.previous_justified_epoch > slot_to_epoch(store.blocks[store.justified_root].slot):
-        store.justified_root = state.previous_justified_root
-    # Update finalized block root
-    if state.finalized_epoch > slot_to_epoch(store.blocks[store.finalized_root].slot):
-        store.finalized_root = state.finalized_root
+    # Add new state for this block to the store
+    store.block_states[signing_root(block)] = state
+
+    # Update justified checkpoint
+    if state.current_justified_epoch > store.justified_checkpoint.epoch:
+        store.justified_checkpoint = Checkpoint(state.current_justified_epoch, state.current_justified_root)
+    elif state.previous_justified_epoch > store.justified_checkpoint.epoch:
+        store.justified_checkpoint = Checkpoint(state.previous_justified_epoch, state.previous_justified_root)
+
+    # Update finalized checkpoint
+    if state.finalized_epoch > state.finalized_epoch:
+        store.finalized_checkpoint = Checkpoint(state.finalized_epoch, state.finalized_root)
 ```
 
 #### `on_attestation`
 
 ```python
 def on_attestation(store: Store, attestation: Attestation) -> None:
-    state = store.states[get_head(store)]
-    indexed_attestation = convert_to_indexed(state, attestation)
-    validate_indexed_attestation(state, indexed_attestation)
+    target = Checkpoint(attestation.data.target_epoch, attestation.data.target_root)
+
+    # Cannot calculate the current shuffling if have not seen the target
+    assert target.root in store.blocks
+
+    # Attestations cannot be from future epochs. If they are, delay consideration until the epoch arrivesr
+    base_state = store.block_states[target.root].copy()
+    assert store.time >= base_state.genesis_time + get_epoch_start_slot(target.epoch) * SECONDS_PER_SLOT
+
+    # Store target checkpoint state if not yet seen
+    if target not in store.checkpoint_states:
+        process_slots(base_state, get_epoch_start_slot(target.epoch))
+        store.checkpoint_states[target] = base_state
+    target_state = store.checkpoint_states[target]
+
+    # Attestations can only affect the fork choice of subsequent slots.
+    # Delay consideration in the fork choice until their slot is in the past.
+    attestation_slot = get_attestation_data_slot(target_state, attestation.data)
+    assert store.time >= (attestation_slot + 1) * SECONDS_PER_SLOT
+
+    # Get state at the `target` to validate attestation and calculate the committees
+    indexed_attestation = convert_to_indexed(target_state, attestation)
+    validate_indexed_attestation(target_state, indexed_attestation)
+
+    # Update latest messages
     for i in indexed_attestation.custody_bit_0_indices + indexed_attestation.custody_bit_1_indices:
-        if i not in store.latest_targets or attestation.data.target_epoch > store.latest_targets[i].epoch:
-            store.latest_targets[i] = Target(attestation.data.target_epoch, attestation.data.target_root)
+        if i not in store.latest_messages or target.epoch > store.latest_messages[i].epoch:
+            store.latest_messages[i] = LatestMessage(epoch=target.epoch, root=attestation.data.beacon_block_root)
 ```

test_libs/pyspec/eth2spec/test/fork_choice/test_get_head.py

@@ -0,0 +1,118 @@
+from eth2spec.test.context import with_all_phases, with_state, bls_switch
+from eth2spec.test.helpers.attestations import get_valid_attestation
+from eth2spec.test.helpers.block import build_empty_block_for_next_slot
+from eth2spec.test.helpers.state import state_transition_and_sign_block
+
+
+def add_block_to_store(spec, store, block):
+    pre_state = store.block_states[block.parent_root]
+    block_time = pre_state.genesis_time + block.slot * spec.SECONDS_PER_SLOT
+
+    if store.time < block_time:
+        spec.on_tick(store, block_time)
+
+    spec.on_block(store, block)
+
+
+def add_attestation_to_store(spec, store, attestation):
+    parent_block = store.blocks[attestation.data.beacon_block_root]
+    pre_state = store.block_states[spec.signing_root(parent_block)]
+    block_time = pre_state.genesis_time + parent_block.slot * spec.SECONDS_PER_SLOT
+    next_epoch_time = block_time + spec.SLOTS_PER_EPOCH * spec.SECONDS_PER_SLOT
+
+    if store.time < next_epoch_time:
+        spec.on_tick(store, next_epoch_time)
+
+    spec.on_attestation(store, attestation)
+
+
+@with_all_phases
+@with_state
+@bls_switch
+def test_genesis(spec, state):
+    # Initialization
+    store = spec.get_genesis_store(state)
+    genesis_block = spec.BeaconBlock(state_root=state.hash_tree_root())
+    assert spec.get_head(store) == spec.signing_root(genesis_block)
+
+
+@with_all_phases
+@with_state
+@bls_switch
+def test_chain_no_attestations(spec, state):
+    # Initialization
+    store = spec.get_genesis_store(state)
+    genesis_block = spec.BeaconBlock(state_root=state.hash_tree_root())
+    assert spec.get_head(store) == spec.signing_root(genesis_block)
+
+    # On receiving a block of `GENESIS_SLOT + 1` slot
+    block_1 = build_empty_block_for_next_slot(spec, state)
+    state_transition_and_sign_block(spec, state, block_1)
+    add_block_to_store(spec, store, block_1)
+
+    # On receiving a block of next epoch
+    block_2 = build_empty_block_for_next_slot(spec, state)
+    state_transition_and_sign_block(spec, state, block_2)
+    add_block_to_store(spec, store, block_2)
+
+    assert spec.get_head(store) == spec.signing_root(block_2)
+
+
+@with_all_phases
+@with_state
+@bls_switch
+def test_split_tie_breaker_no_attestations(spec, state):
+    genesis_state = state.copy()
+
+    # Initialization
+    store = spec.get_genesis_store(state)
+    genesis_block = spec.BeaconBlock(state_root=state.hash_tree_root())
+    assert spec.get_head(store) == spec.signing_root(genesis_block)
+
+    # block at slot 1
+    block_1_state = genesis_state.copy()
+    block_1 = build_empty_block_for_next_slot(spec, block_1_state)
+    state_transition_and_sign_block(spec, block_1_state, block_1)
+    add_block_to_store(spec, store, block_1)
+
+    # additional block at slot 1
+    block_2_state = genesis_state.copy()
+    block_2 = build_empty_block_for_next_slot(spec, block_2_state)
+    block_2.body.graffiti = b'\x42' * 32
+    state_transition_and_sign_block(spec, block_2_state, block_2)
+    add_block_to_store(spec, store, block_2)
+
+    highest_root = max(spec.signing_root(block_1), spec.signing_root(block_2))
+
+    assert spec.get_head(store) == highest_root
+
+
+@with_all_phases
+@with_state
+@bls_switch
+def test_shorter_chain_but_heavier_weight(spec, state):
+    genesis_state = state.copy()
+
+    # Initialization
+    store = spec.get_genesis_store(state)
+    genesis_block = spec.BeaconBlock(state_root=state.hash_tree_root())
+    assert spec.get_head(store) == spec.signing_root(genesis_block)
+
+    # build longer tree
+    long_state = genesis_state.copy()
+    for i in range(3):
+        long_block = build_empty_block_for_next_slot(spec, long_state)
+        state_transition_and_sign_block(spec, long_state, long_block)
+        add_block_to_store(spec, store, long_block)
+
+    # build short tree
+    short_state = genesis_state.copy()
+    short_block = build_empty_block_for_next_slot(spec, short_state)
+    short_block.body.graffiti = b'\x42' * 32
+    state_transition_and_sign_block(spec, short_state, short_block)
+    add_block_to_store(spec, store, short_block)
+
+    short_attestation = get_valid_attestation(spec, short_state, short_block.slot, signed=True)
+    add_attestation_to_store(spec, store, short_attestation)
+
+    assert spec.get_head(store) == spec.signing_root(short_block)