Merge pull request #1495 from ethereum/filter-fc-justified

filter viable branches in fork choice

specs/core/0_fork-choice.md

@@ -136,17 +136,72 @@ def get_latest_attesting_balance(store: Store, root: Root) -> Gwei:
     ))
 ```
 
+#### `filter_block_tree`
+
+```python
+def filter_block_tree(store: Store, block_root: Root, blocks: Dict[Root, BeaconBlock]) -> bool:
+    block = store.blocks[block_root]
+    children = [
+        root for root in store.blocks.keys()
+        if store.blocks[root].parent_root == block_root
+    ]
+
+    # If any children branches contain expected finalized/justified checkpoints,
+    # add to filtered block-tree and signal viability to parent.
+    if any(children):
+        filter_block_tree_result = [filter_block_tree(store, child, blocks) for child in children]
+        if any(filter_block_tree_result):
+            blocks[block_root] = block
+            return True
+        return False
+
+    # If leaf block, check finalized/justified checkpoints as matching latest.
+    head_state = store.block_states[block_root]
+
+    correct_justified = (
+        store.justified_checkpoint.epoch == GENESIS_EPOCH
+        or head_state.current_justified_checkpoint == store.justified_checkpoint
+    )
+    correct_finalized = (
+        store.finalized_checkpoint.epoch == GENESIS_EPOCH
+        or head_state.finalized_checkpoint == store.finalized_checkpoint
+    )
+    # If expected finalized/justified, add to viable block-tree and signal viability to parent.
+    if correct_justified and correct_finalized:
+        blocks[block_root] = block
+        return True
+
+    # Otherwise, branch not viable
+    return False
+```
+
+#### `get_filtered_block_tree`
+
+```python
+def get_filtered_block_tree(store: Store) -> Dict[Root, BeaconBlock]:
+    """
+    Retrieve a filtered block true from ``store``, only returning branches
+    whose leaf state's justified/finalized info agrees with that in ``store``.
+    """
+    base = store.justified_checkpoint.root
+    blocks: Dict[Root, BeaconBlock] = {}
+    filter_block_tree(store, base, blocks)
+    return blocks
+```
+
 #### `get_head`
 
 ```python
 def get_head(store: Store) -> Root:
+    # Get filtered block tree that only includes viable branches
+    blocks = get_filtered_block_tree(store)
     # Execute the LMD-GHOST fork choice
     head = store.justified_checkpoint.root
     justified_slot = compute_start_slot_at_epoch(store.justified_checkpoint.epoch)
     while True:
         children = [
-            root for root in store.blocks.keys()
-            if store.blocks[root].parent_root == head and store.blocks[root].slot > justified_slot
+            root for root in blocks.keys()
+            if blocks[root].parent_root == head and blocks[root].slot > justified_slot
         ]
         if len(children) == 0:
             return head
@@ -172,8 +227,8 @@ def should_update_justified_checkpoint(store: Store, new_justified_checkpoint: C
     if new_justified_block.slot <= compute_start_slot_at_epoch(store.justified_checkpoint.epoch):
         return False
     if not (
-        get_ancestor(store, new_justified_checkpoint.root, store.blocks[store.justified_checkpoint.root].slot) ==
-        store.justified_checkpoint.root
+        get_ancestor(store, new_justified_checkpoint.root, store.blocks[store.justified_checkpoint.root].slot)
+        == store.justified_checkpoint.root
     ):
         return False
 

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

@@ -1,7 +1,11 @@
 from eth2spec.test.context import with_all_phases, spec_state_test
 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
+from eth2spec.test.helpers.state import (
+    next_epoch,
+    next_epoch_with_attestations,
+    state_transition_and_sign_block,
+)
 
 
 def add_block_to_store(spec, store, signed_block):
@@ -112,3 +116,79 @@ def test_shorter_chain_but_heavier_weight(spec, state):
     add_attestation_to_store(spec, store, short_attestation)
 
     assert spec.get_head(store) == spec.hash_tree_root(short_block)
+
+
+@with_all_phases
+@spec_state_test
+def test_filtered_block_tree(spec, state):
+    # Initialization
+    genesis_state_root = state.hash_tree_root()
+    store = spec.get_genesis_store(state)
+    genesis_block = spec.BeaconBlock(state_root=genesis_state_root)
+
+    # transition state past initial couple of epochs
+    next_epoch(spec, state)
+    next_epoch(spec, state)
+
+    assert spec.get_head(store) == spec.hash_tree_root(genesis_block)
+
+    # fill in attestations for entire epoch, justifying the recent epoch
+    prev_state, signed_blocks, state = next_epoch_with_attestations(spec, state, True, False)
+    attestations = [
+        attestation for signed_block in signed_blocks
+        for attestation in signed_block.message.body.attestations
+    ]
+    assert state.current_justified_checkpoint.epoch > prev_state.current_justified_checkpoint.epoch
+
+    # tick time forward and add blocks and attestations to store
+    current_time = state.slot * spec.SECONDS_PER_SLOT + store.genesis_time
+    spec.on_tick(store, current_time)
+    for signed_block in signed_blocks:
+        spec.on_block(store, signed_block)
+    for attestation in attestations:
+        spec.on_attestation(store, attestation)
+
+    assert store.justified_checkpoint == state.current_justified_checkpoint
+
+    # the last block in the branch should be the head
+    expected_head_root = spec.hash_tree_root(signed_blocks[-1].message)
+    assert spec.get_head(store) == expected_head_root
+
+    #
+    # create branch containing the justified block but not containing enough on
+    # chain votes to justify that block
+    #
+
+    # build a chain without attestations off of previous justified block
+    non_viable_state = store.block_states[store.justified_checkpoint.root].copy()
+
+    # ensure that next wave of votes are for future epoch
+    next_epoch(spec, non_viable_state)
+    next_epoch(spec, non_viable_state)
+    next_epoch(spec, non_viable_state)
+    assert spec.get_current_epoch(non_viable_state) > store.justified_checkpoint.epoch
+
+    # create rogue block that will be attested to in this non-viable branch
+    rogue_block = build_empty_block_for_next_slot(spec, non_viable_state)
+    signed_rogue_block = state_transition_and_sign_block(spec, non_viable_state, rogue_block)
+
+    # create an epoch's worth of attestations for the rogue block
+    next_epoch(spec, non_viable_state)
+    attestations = []
+    for i in range(spec.SLOTS_PER_EPOCH):
+        slot = rogue_block.slot + i
+        for index in range(spec.get_committee_count_at_slot(non_viable_state, slot)):
+            attestation = get_valid_attestation(spec, non_viable_state, rogue_block.slot + i, index)
+            attestations.append(attestation)
+
+    # tick time forward to be able to include up to the latest attestation
+    current_time = (attestations[-1].data.slot + 1) * spec.SECONDS_PER_SLOT + store.genesis_time
+    spec.on_tick(store, current_time)
+
+    # include rogue block and associated attestations in the store
+    spec.on_block(store, signed_rogue_block)
+    for attestation in attestations:
+        spec.on_attestation(store, attestation)
+
+    # ensure that get_head still returns the head from the previous branch
+    assert spec.get_head(store) == expected_head_root