modify fork choice to utilize epochs as first class citizens

specs/core/0_fork-choice.md

@@ -63,17 +63,30 @@ class Target(object):
     root: Hash
 ```
 
+#### `RootSlot`
+
+```python
+@dataclass
+class RootSlot(object):
+    slot: Slot
+    root: Hash
+```
+
+
+
 #### `Store`
 
 ```python
 @dataclass
 class Store(object):
     blocks: Dict[Hash, BeaconBlock] = field(default_factory=dict)
-    states: Dict[Hash, BeaconState] = field(default_factory=dict)
+    states: Dict[RootSlot, BeaconState] = field(default_factory=dict)
     time: int = 0
     latest_targets: Dict[ValidatorIndex, Target] = field(default_factory=dict)
     justified_root: Hash = ZERO_HASH
+    justified_epoch: Epoch = 0
     finalized_root: Hash = ZERO_HASH
+    finalized_epoch: Epoch = 0
 ```
 
 #### `get_genesis_store`
@@ -98,8 +111,8 @@ 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, slot_to_epoch(state.slot))
+    state = store.states[RootSlot(store.justified_root, get_epoch_start_slot(store.justified_epoch)]
+    active_indices = get_active_validator_indices(state.validator_registry, 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
@@ -112,8 +125,12 @@ def get_latest_attesting_balance(store: Store, root: Hash) -> Gwei:
 def get_head(store: Store) -> Hash:
     # Execute the LMD-GHOST fork choice
     head = store.justified_root
+    justified_slot = get_epoch_start_slot(store.justified_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
@@ -137,30 +154,58 @@ def on_block(store: Store, block: BeaconBlock) -> None:
     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
+    # Check that block is later than the finalized epoch slot
+    assert blocks.slot > get_epoch_start_slot(store.finalized_epoch)
     # Check block slot against Unix time
     pre_state = deepcopy(store.states[block.parent_root])
     assert store.time >= pre_state.genesis_time + block.slot * SECONDS_PER_SLOT
     # 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):
+    # Add new state for this block to the store
+    store.states[RootSlot(signing_root(block), block.slot)] = state
+    # Update justified block root and epoch
+    previous_justified_epoch = store.justified_epoch
+    justified_root_slot = None
+    if state.current_justified_epoch > store.justified_epoch:
         store.justified_root = state.current_justified_root
-    elif state.previous_justified_epoch > slot_to_epoch(store.blocks[store.justified_root].slot):
+        state.justified_epoch = state.current_justified_epoch
+    elif state.previous_justified_epoch > store.justified_epoch:
         store.justified_root = state.previous_justified_root
-    # Update finalized block root
-    if state.finalized_epoch > slot_to_epoch(store.blocks[store.finalized_root].slot):
+        state.justified_epoch = state.previous_justified_epoch
+
+    # Store justified state
+    if previous_justified_epoch != store.justified_epoch:
+        justified_slot = get_epoch_start_slot(store.justified_epoch)
+        justified_root_slot = RootSlot(store.justified_root, justified_slot)
+        if justified_root_slot not in store.states:
+            base_state = store.states[RootSlot(store.justified_root, store.blocks[store.justified_root].slot)]
+            store.states[justified_root_slot] = process_slots(deepcopy(base_state), justified_slot)
+
+    # Update finalized block root and epoch, and store finalized state
+    if state.finalized_epoch > state.finalized_epoch:
         store.finalized_root = state.finalized_root
+        store.finalized_epoch = state.finalized_epoch
+        finalized_slot = get_epoch_start_slot(store.finalized_epoch)
+        finalized_root_slot = RootSlot(store.finalized_root, finalized_slot)
+        if finalized_root_slot not in store.states:
+            base_state = store.states[RootSlot(store.finalized_root, store.blocks[store.finalized_root].slot)]
+            store.states[finalized_root_slot] = process_slots(deepcopy(base_state), finalized_slot)
 ```
 
 #### `on_attestation`
 
 ```python
 def on_attestation(store: Store, attestation: Attestation) -> None:
-    state = store.states[get_head(store)]
+    # cannot calculate the current shuffling if have not seen the target
+    assert attestation.data.target_root in store.blocks
+
+    # get state at the `target_root`/`target_epoch` to validate attestation and calculate the committees
+    state = store.states[(attestation.data.target_root, get_epoch_start_slot(attestation.data.target_epoch))]
+
     indexed_attestation = convert_to_indexed(state, attestation)
     validate_indexed_attestation(state, indexed_attestation)
+
+    # update latest targets
     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)