From bb3c360734d21e477f5d155ebce8b40b7e100828 Mon Sep 17 00:00:00 2001
From: Hsiao-Wei Wang <hwwang156@gmail.com>
Date: Fri, 3 Jul 2020 10:48:57 +0800
Subject: [PATCH 1/7] Handle the dependencies order of dataclass objects

---
 setup.py | 54 ++++++++++++++++++++++++++++++++++++------------------
 1 file changed, 36 insertions(+), 18 deletions(-)

diff --git a/setup.py b/setup.py
index 214559d22..b080dbe23 100644
--- a/setup.py
+++ b/setup.py
@@ -1,3 +1,4 @@
+from enum import Enum, auto
 from setuptools import setup, find_packages, Command
 from setuptools.command.build_py import build_py
 from distutils import dir_util
@@ -14,6 +15,13 @@ class SpecObject(NamedTuple):
     custom_types: Dict[str, str]
     constants: Dict[str, str]
     ssz_objects: Dict[str, str]
+    dataclasses: Dict[str, str]
+
+
+class CodeBlockType(Enum):
+    SSZ = auto()
+    DATACLASS = auto()
+    FUNCTION = auto()
 
 
 def get_spec(file_name: str) -> SpecObject:
@@ -28,8 +36,9 @@ def get_spec(file_name: str) -> SpecObject:
     functions: Dict[str, str] = {}
     constants: Dict[str, str] = {}
     ssz_objects: Dict[str, str] = {}
+    dataclasses: Dict[str, str] = {}
     function_matcher = re.compile(FUNCTION_REGEX)
-    is_ssz = False
+    block_type = CodeBlockType.FUNCTION
     custom_types: Dict[str, str] = {}
     for linenum, line in enumerate(open(file_name).readlines()):
         line = line.rstrip()
@@ -43,20 +52,26 @@ def get_spec(file_name: str) -> SpecObject:
         else:
             # Handle function definitions & ssz_objects
             if pulling_from is not None:
-                # SSZ Object
                 if len(line) > 18 and line[:6] == 'class ' and line[-12:] == '(Container):':
                     name = line[6:-12]
                     # Check consistency with markdown header
                     assert name == current_name
-                    is_ssz = True
-                # function definition
+                    block_type = CodeBlockType.SSZ
+                elif line[:10] == '@dataclass':
+                    block_type = CodeBlockType.DATACLASS
                 elif function_matcher.match(line) is not None:
                     current_name = function_matcher.match(line).group(0)
-                    is_ssz = False
-                if is_ssz:
+                    block_type = CodeBlockType.FUNCTION
+
+                if block_type == CodeBlockType.SSZ:
                     ssz_objects[current_name] = ssz_objects.get(current_name, '') + line + '\n'
-                else:
+                elif block_type == CodeBlockType.DATACLASS:
+                    dataclasses[current_name] = dataclasses.get(current_name, '') + line + '\n'
+                elif block_type == CodeBlockType.FUNCTION:
                     functions[current_name] = functions.get(current_name, '') + line + '\n'
+                else:
+                    pass
+
             # Handle constant and custom types table entries
             elif pulling_from is None and len(line) > 0 and line[0] == '|':
                 row = line[1:].split('|')
@@ -75,7 +90,7 @@ def get_spec(file_name: str) -> SpecObject:
                         constants[row[0]] = row[1].replace('**TBD**', '2**32')
                     elif row[1].startswith('uint') or row[1].startswith('Bytes'):
                         custom_types[row[0]] = row[1]
-    return SpecObject(functions, custom_types, constants, ssz_objects)
+    return SpecObject(functions, custom_types, constants, ssz_objects, dataclasses)
 
 
 CONFIG_LOADER = '''
@@ -237,7 +252,7 @@ get_start_shard = cache_this(
     _get_start_shard, lru_size=SLOTS_PER_EPOCH * 3)'''
 
 
-def objects_to_spec(spec_object: SpecObject, imports: str, fork: str) -> str:
+def objects_to_spec(spec_object: SpecObject, imports: str, fork: str, ordered_class_objects: Dict[str, str]) -> str:
     """
     Given all the objects that constitute a spec, combine them into a single pyfile.
     """
@@ -257,7 +272,7 @@ def objects_to_spec(spec_object: SpecObject, imports: str, fork: str) -> str:
         if k == "BLS12_381_Q":
             spec_object.constants[k] += "  # noqa: E501"
     constants_spec = '\n'.join(map(lambda x: '%s = %s' % (x, spec_object.constants[x]), spec_object.constants))
-    ssz_objects_instantiation_spec = '\n\n'.join(spec_object.ssz_objects.values())
+    ordered_class_objects_spec = '\n\n'.join(ordered_class_objects.values())
     spec = (
             imports
             + '\n\n' + f"fork = \'{fork}\'\n"
@@ -265,7 +280,7 @@ def objects_to_spec(spec_object: SpecObject, imports: str, fork: str) -> str:
             + '\n' + SUNDRY_CONSTANTS_FUNCTIONS
             + '\n\n' + constants_spec
             + '\n\n' + CONFIG_LOADER
-            + '\n\n' + ssz_objects_instantiation_spec
+            + '\n\n' + ordered_class_objects_spec
             + '\n\n' + functions_spec
             + '\n' + PHASE0_SUNDRY_FUNCTIONS
     )
@@ -291,11 +306,12 @@ ignored_dependencies = [
     'bit', 'boolean', 'Vector', 'List', 'Container', 'BLSPubkey', 'BLSSignature',
     'Bytes1', 'Bytes4', 'Bytes32', 'Bytes48', 'Bytes96', 'Bitlist', 'Bitvector',
     'uint8', 'uint16', 'uint32', 'uint64', 'uint128', 'uint256',
-    'bytes', 'byte', 'ByteList', 'ByteVector'
+    'bytes', 'byte', 'ByteList', 'ByteVector',
+    'Dict', 'dict', 'field',
 ]
 
 
-def dependency_order_ssz_objects(objects: Dict[str, str], custom_types: Dict[str, str]) -> None:
+def dependency_order_class_objects(objects: Dict[str, str], custom_types: Dict[str, str]) -> None:
     """
     Determines which SSZ Object is dependent on which other and orders them appropriately
     """
@@ -332,13 +348,14 @@ def combine_spec_objects(spec0: SpecObject, spec1: SpecObject) -> SpecObject:
     """
     Takes in two spec variants (as tuples of their objects) and combines them using the appropriate combiner function.
     """
-    functions0, custom_types0, constants0, ssz_objects0 = spec0
-    functions1, custom_types1, constants1, ssz_objects1 = spec1
+    functions0, custom_types0, constants0, ssz_objects0, dataclasses0 = spec0
+    functions1, custom_types1, constants1, ssz_objects1, dataclasses1 = spec1
     functions = combine_functions(functions0, functions1)
     custom_types = combine_constants(custom_types0, custom_types1)
     constants = combine_constants(constants0, constants1)
     ssz_objects = combine_ssz_objects(ssz_objects0, ssz_objects1, custom_types)
-    return SpecObject(functions, custom_types, constants, ssz_objects)
+    dataclasses = combine_functions(dataclasses0, dataclasses1)
+    return SpecObject(functions, custom_types, constants, ssz_objects, dataclasses)
 
 
 fork_imports = {
@@ -354,9 +371,10 @@ def build_spec(fork: str, source_files: List[str]) -> str:
     for value in all_specs[1:]:
         spec_object = combine_spec_objects(spec_object, value)
 
-    dependency_order_ssz_objects(spec_object.ssz_objects, spec_object.custom_types)
+    class_objects = {**spec_object.ssz_objects, **spec_object.dataclasses}
+    dependency_order_class_objects(class_objects, spec_object.custom_types)
 
-    return objects_to_spec(spec_object, fork_imports[fork], fork)
+    return objects_to_spec(spec_object, fork_imports[fork], fork, class_objects)
 
 
 class PySpecCommand(Command):

From 43ef9aa294214fd89e53430db5ccb234aa71d8c7 Mon Sep 17 00:00:00 2001
From: Hsiao-Wei Wang <hwwang156@gmail.com>
Date: Fri, 3 Jul 2020 10:49:19 +0800
Subject: [PATCH 2/7] Make ShardLatestMessage per shard per validator

1. Add `ShardLatestMessage` dataclass
2. To make it compatible with phase 0 tests and APIs, add `Store.shard_stores: Dict[Shard, ShardStore]`
3. Update `get_forkchoice_store` and `update_latest_messages`
---
 specs/phase1/fork-choice.md                   | 80 ++++++++++++++++---
 specs/phase1/shard-fork-choice.md             | 19 ++---
 .../test/fork_choice/test_on_attestation.py   |  7 +-
 .../test/fork_choice/test_on_shard_head.py    |  2 +-
 4 files changed, 82 insertions(+), 26 deletions(-)

diff --git a/specs/phase1/fork-choice.md b/specs/phase1/fork-choice.md
index 41787cfd0..d0b06ffa9 100644
--- a/specs/phase1/fork-choice.md
+++ b/specs/phase1/fork-choice.md
@@ -9,8 +9,13 @@
 
 
 - [Introduction](#introduction)
-  - [Helpers](#helpers)
-    - [Extended `LatestMessage`](#extended-latestmessage)
+  - [Updated data structures](#updated-data-structures)
+    - [Extended `Store`](#extended-store)
+  - [New data structures](#new-data-structures)
+    - [`ShardLatestMessage`](#shardlatestmessage)
+    - [`ShardStore`](#shardstore)
+  - [Updated helpers](#updated-helpers)
+    - [Updated `get_forkchoice_store`](#updated-get_forkchoice_store)
     - [Updated `update_latest_messages`](#updated-update_latest_messages)
 
 <!-- END doctoc generated TOC please keep comment here to allow auto update -->
@@ -20,17 +25,74 @@
 
 This document is the beacon chain fork choice spec for part of Ethereum 2.0 Phase 1.
 
-### Helpers
+### Updated data structures
 
-#### Extended `LatestMessage`
+#### Extended `Store`
+
+```python
+@dataclass
+class Store(object):
+    time: uint64
+    genesis_time: uint64
+    justified_checkpoint: Checkpoint
+    finalized_checkpoint: Checkpoint
+    best_justified_checkpoint: Checkpoint
+    blocks: Dict[Root, BeaconBlock] = field(default_factory=dict)
+    block_states: Dict[Root, BeaconState] = field(default_factory=dict)
+    checkpoint_states: Dict[Checkpoint, BeaconState] = field(default_factory=dict)
+    latest_messages: Dict[ValidatorIndex, LatestMessage] = field(default_factory=dict)
+    shard_stores: Dict[Shard, ShardStore] = field(default_factory=dict)
+```
+
+### New data structures
+
+#### `ShardLatestMessage`
 
 ```python
 @dataclass(eq=True, frozen=True)
-class LatestMessage(object):
+class ShardLatestMessage(object):
     epoch: Epoch
     root: Root
+```
+
+#### `ShardStore`
+
+```python
+@dataclass
+class ShardStore:
     shard: Shard
-    shard_root: Root
+    signed_blocks: Dict[Root, SignedShardBlock] = field(default_factory=dict)
+    block_states: Dict[Root, ShardState] = field(default_factory=dict)
+    latest_messages: Dict[ValidatorIndex, ShardLatestMessage] = field(default_factory=dict)
+```
+
+### Updated helpers
+
+#### Updated `get_forkchoice_store`
+
+```python
+def get_forkchoice_store(anchor_state: BeaconState) -> Store:
+    anchor_block_header = anchor_state.latest_block_header.copy()
+    if anchor_block_header.state_root == Bytes32():
+        anchor_block_header.state_root = hash_tree_root(anchor_state)
+    anchor_root = hash_tree_root(anchor_block_header)
+    anchor_epoch = get_current_epoch(anchor_state)
+    justified_checkpoint = Checkpoint(epoch=anchor_epoch, root=anchor_root)
+    finalized_checkpoint = Checkpoint(epoch=anchor_epoch, root=anchor_root)
+    return Store(
+        time=anchor_state.genesis_time + SECONDS_PER_SLOT * anchor_state.slot,
+        genesis_time=anchor_state.genesis_time,
+        justified_checkpoint=justified_checkpoint,
+        finalized_checkpoint=finalized_checkpoint,
+        best_justified_checkpoint=justified_checkpoint,
+        blocks={anchor_root: anchor_block_header},
+        block_states={anchor_root: anchor_state.copy()},
+        checkpoint_states={justified_checkpoint: anchor_state.copy()},
+        shard_stores={
+            Shard(shard): get_forkchoice_shard_store(anchor_state, Shard(shard))
+            for shard in range(get_active_shard_count(anchor_state))
+        }
+    )
 ```
 
 #### Updated `update_latest_messages`
@@ -43,7 +105,7 @@ def update_latest_messages(store: Store, attesting_indices: Sequence[ValidatorIn
     shard = attestation.data.shard
     for i in attesting_indices:
         if i not in store.latest_messages or target.epoch > store.latest_messages[i].epoch:
-            store.latest_messages[i] = LatestMessage(
-                epoch=target.epoch, root=beacon_block_root, shard=shard, shard_root=attestation.data.shard_head_root
-            )
+            store.latest_messages[i] = LatestMessage(epoch=target.epoch, root=beacon_block_root)
+            shard_latest_message = ShardLatestMessage(epoch=target.epoch, root=attestation.data.shard_head_root)
+            store.shard_stores[shard].latest_messages[i] = shard_latest_message
 ```
diff --git a/specs/phase1/shard-fork-choice.md b/specs/phase1/shard-fork-choice.md
index 8416009d7..1835df432 100644
--- a/specs/phase1/shard-fork-choice.md
+++ b/specs/phase1/shard-fork-choice.md
@@ -11,7 +11,6 @@
 - [Introduction](#introduction)
 - [Fork choice](#fork-choice)
   - [Helpers](#helpers)
-    - [`ShardStore`](#shardstore)
     - [`get_forkchoice_shard_store`](#get_forkchoice_shard_store)
     - [`get_shard_latest_attesting_balance`](#get_shard_latest_attesting_balance)
     - [`get_shard_head`](#get_shard_head)
@@ -30,16 +29,6 @@ This document is the shard chain fork choice spec for part of Ethereum 2.0 Phase
 
 ### Helpers
 
-#### `ShardStore`
-
-```python
-@dataclass
-class ShardStore:
-    shard: Shard
-    signed_blocks: Dict[Root, SignedShardBlock] = field(default_factory=dict)
-    block_states: Dict[Root, ShardState] = field(default_factory=dict)
-```
-
 #### `get_forkchoice_shard_store`
 
 ```python
@@ -64,12 +53,14 @@ def get_shard_latest_attesting_balance(store: Store, shard_store: ShardStore, ro
     return Gwei(sum(
         state.validators[i].effective_balance for i in active_indices
         if (
-            i in store.latest_messages
+            i in shard_store.latest_messages
             # TODO: check the latest message logic: currently, validator's previous vote of another shard
             # would be ignored once their newer vote is accepted. Check if it makes sense.
-            and store.latest_messages[i].shard == shard_store.shard
             and get_shard_ancestor(
-                store, shard_store, store.latest_messages[i].shard_root, shard_store.signed_blocks[root].message.slot
+                store,
+                shard_store,
+                shard_store.latest_messages[i].root,
+                shard_store.signed_blocks[root].message.slot,
             ) == root
         )
     ))
diff --git a/tests/core/pyspec/eth2spec/test/fork_choice/test_on_attestation.py b/tests/core/pyspec/eth2spec/test/fork_choice/test_on_attestation.py
index ffd8b417d..6fa842255 100644
--- a/tests/core/pyspec/eth2spec/test/fork_choice/test_on_attestation.py
+++ b/tests/core/pyspec/eth2spec/test/fork_choice/test_on_attestation.py
@@ -26,9 +26,12 @@ def run_on_attestation(spec, state, store, attestation, valid=True):
         latest_message = spec.LatestMessage(
             epoch=attestation.data.target.epoch,
             root=attestation.data.beacon_block_root,
-            shard=attestation.data.shard,
-            shard_root=attestation.data.shard_head_root,
         )
+        shard_latest_message = spec.ShardLatestMessage(
+            epoch=attestation.data.target.epoch,
+            root=attestation.data.shard_head_root,
+        )
+        assert store.shard_stores[attestation.data.shard].latest_messages[sample_index] == shard_latest_message
 
     assert (
         store.latest_messages[sample_index] == latest_message
diff --git a/tests/core/pyspec/eth2spec/test/fork_choice/test_on_shard_head.py b/tests/core/pyspec/eth2spec/test/fork_choice/test_on_shard_head.py
index a87a85917..3b03906d9 100644
--- a/tests/core/pyspec/eth2spec/test/fork_choice/test_on_shard_head.py
+++ b/tests/core/pyspec/eth2spec/test/fork_choice/test_on_shard_head.py
@@ -111,7 +111,7 @@ def test_basic(spec, state):
     anchor_root = get_anchor_root(spec, state)
     assert spec.get_head(store) == anchor_root
 
-    shard_store = spec.get_forkchoice_shard_store(state, shard)
+    shard_store = store.shard_stores[shard]
     shard_head_root = spec.get_shard_head(store, shard_store)
     assert shard_head_root == state.shard_states[shard].latest_block_root
     assert shard_store.block_states[shard_head_root].slot == 1

From 2da331a34528566a042287f69b9f0722df5ee913 Mon Sep 17 00:00:00 2001
From: Hsiao-Wei Wang <hwwang156@gmail.com>
Date: Wed, 15 Jul 2020 19:01:20 +0800
Subject: [PATCH 3/7] Rename `test_on_shard_head.py` to
 `test_on_shard_block.py`

---
 .../fork_choice/{test_on_shard_head.py => test_on_shard_block.py} | 0
 1 file changed, 0 insertions(+), 0 deletions(-)
 rename tests/core/pyspec/eth2spec/test/fork_choice/{test_on_shard_head.py => test_on_shard_block.py} (100%)

diff --git a/tests/core/pyspec/eth2spec/test/fork_choice/test_on_shard_head.py b/tests/core/pyspec/eth2spec/test/fork_choice/test_on_shard_block.py
similarity index 100%
rename from tests/core/pyspec/eth2spec/test/fork_choice/test_on_shard_head.py
rename to tests/core/pyspec/eth2spec/test/fork_choice/test_on_shard_block.py

From f6b1fe61727e7e8c4a6bcc43eaec58ae5ce2ead3 Mon Sep 17 00:00:00 2001
From: Hsiao-Wei Wang <hwwang156@gmail.com>
Date: Tue, 14 Jul 2020 17:36:27 +0800
Subject: [PATCH 4/7] Refactor tests and avoiding passing `shart_store` to
 helper functions

---
 specs/phase1/shard-fork-choice.md             |  30 +--
 .../test/fork_choice/test_on_shard_block.py   | 178 ++++++++++++++----
 2 files changed, 159 insertions(+), 49 deletions(-)

diff --git a/specs/phase1/shard-fork-choice.md b/specs/phase1/shard-fork-choice.md
index 1835df432..39e957819 100644
--- a/specs/phase1/shard-fork-choice.md
+++ b/specs/phase1/shard-fork-choice.md
@@ -47,7 +47,8 @@ def get_forkchoice_shard_store(anchor_state: BeaconState, shard: Shard) -> Shard
 #### `get_shard_latest_attesting_balance`
 
 ```python
-def get_shard_latest_attesting_balance(store: Store, shard_store: ShardStore, root: Root) -> Gwei:
+def get_shard_latest_attesting_balance(store: Store, shard: Shard, root: Root) -> Gwei:
+    shard_store = store.shard_stores[shard]
     state = store.checkpoint_states[store.justified_checkpoint]
     active_indices = get_active_validator_indices(state, get_current_epoch(state))
     return Gwei(sum(
@@ -58,7 +59,7 @@ def get_shard_latest_attesting_balance(store: Store, shard_store: ShardStore, ro
             # would be ignored once their newer vote is accepted. Check if it makes sense.
             and get_shard_ancestor(
                 store,
-                shard_store,
+                shard,
                 shard_store.latest_messages[i].root,
                 shard_store.signed_blocks[root].message.slot,
             ) == root
@@ -69,10 +70,14 @@ def get_shard_latest_attesting_balance(store: Store, shard_store: ShardStore, ro
 #### `get_shard_head`
 
 ```python
-def get_shard_head(store: Store, shard_store: ShardStore) -> Root:
+def get_shard_head(store: Store, shard: Shard) -> Root:
     # Execute the LMD-GHOST fork choice
+    """
+    Execute the LMD-GHOST fork choice.
+    """
+    shard_store = store.shard_stores[shard]
     beacon_head_root = get_head(store)
-    shard_head_state = store.block_states[beacon_head_root].shard_states[shard_store.shard]
+    shard_head_state = store.block_states[beacon_head_root].shard_states[shard]
     shard_head_root = shard_head_state.latest_block_root
     shard_blocks = {
         root: signed_shard_block.message for root, signed_shard_block in shard_store.signed_blocks.items()
@@ -88,17 +93,18 @@ def get_shard_head(store: Store, shard_store: ShardStore) -> Root:
             return shard_head_root
         # Sort by latest attesting balance with ties broken lexicographically
         shard_head_root = max(
-            children, key=lambda root: (get_shard_latest_attesting_balance(store, shard_store, root), root)
+            children, key=lambda root: (get_shard_latest_attesting_balance(store, shard, root), root)
         )
 ```
 
 #### `get_shard_ancestor`
 
 ```python
-def get_shard_ancestor(store: Store, shard_store: ShardStore, root: Root, slot: Slot) -> Root:
+def get_shard_ancestor(store: Store, shard: Shard, root: Root, slot: Slot) -> Root:
+    shard_store = store.shard_stores[shard]
     block = shard_store.signed_blocks[root].message
     if block.slot > slot:
-        return get_shard_ancestor(store, shard_store, block.shard_parent_root, slot)
+        return get_shard_ancestor(store, shard, block.shard_parent_root, slot)
     elif block.slot == slot:
         return root
     else:
@@ -109,17 +115,17 @@ def get_shard_ancestor(store: Store, shard_store: ShardStore, root: Root, slot:
 #### `get_pending_shard_blocks`
 
 ```python
-def get_pending_shard_blocks(store: Store, shard_store: ShardStore) -> Sequence[SignedShardBlock]:
+def get_pending_shard_blocks(store: Store, shard: Shard) -> Sequence[SignedShardBlock]:
     """
     Return the canonical shard block branch that has not yet been crosslinked.
     """
-    shard = shard_store.shard
+    shard_store = store.shard_stores[shard]
 
     beacon_head_root = get_head(store)
     beacon_head_state = store.block_states[beacon_head_root]
     latest_shard_block_root = beacon_head_state.shard_states[shard].latest_block_root
 
-    shard_head_root = get_shard_head(store, shard_store)
+    shard_head_root = get_shard_head(store, shard)
     root = shard_head_root
     signed_shard_blocks = []
     while root != latest_shard_block_root:
@@ -136,9 +142,9 @@ def get_pending_shard_blocks(store: Store, shard_store: ShardStore) -> Sequence[
 #### `on_shard_block`
 
 ```python
-def on_shard_block(store: Store, shard_store: ShardStore, signed_shard_block: SignedShardBlock) -> None:
+def on_shard_block(store: Store, shard: Shard, signed_shard_block: SignedShardBlock) -> None:
+    shard_store = store.shard_stores[shard]
     shard_block = signed_shard_block.message
-    shard = shard_store.shard
 
     # Check shard
     # TODO: check it in networking spec
diff --git a/tests/core/pyspec/eth2spec/test/fork_choice/test_on_shard_block.py b/tests/core/pyspec/eth2spec/test/fork_choice/test_on_shard_block.py
index 3b03906d9..dae9246fb 100644
--- a/tests/core/pyspec/eth2spec/test/fork_choice/test_on_shard_block.py
+++ b/tests/core/pyspec/eth2spec/test/fork_choice/test_on_shard_block.py
@@ -8,27 +8,43 @@ from eth2spec.test.helpers.shard_block import (
     get_committee_index_of_shard,
 )
 from eth2spec.test.helpers.fork_choice import add_block_to_store, get_anchor_root
+from eth2spec.test.helpers.shard_transitions import is_full_crosslink
 from eth2spec.test.helpers.state import state_transition_and_sign_block
 from eth2spec.test.helpers.block import build_empty_block
 
 
-def run_on_shard_block(spec, store, shard_store, signed_block, valid=True):
+def run_on_shard_block(spec, store, shard, signed_block, valid=True):
     if not valid:
         try:
-            spec.on_shard_block(store, shard_store, signed_block)
+            spec.on_shard_block(store, shard, signed_block)
         except AssertionError:
             return
         else:
             assert False
 
-    spec.on_shard_block(store, shard_store, signed_block)
+    spec.on_shard_block(store, shard, signed_block)
+    shard_store = store.shard_stores[shard]
     assert shard_store.signed_blocks[hash_tree_root(signed_block.message)] == signed_block
 
 
-def apply_shard_block(spec, store, shard_store, beacon_parent_state, shard_blocks_buffer):
-    shard = shard_store.shard
+def initialize_store(spec, state, shard):
+    store = spec.get_forkchoice_store(state)
+    anchor_root = get_anchor_root(spec, state)
+    assert spec.get_head(store) == anchor_root
+
+    shard_head_root = spec.get_shard_head(store, shard)
+    assert shard_head_root == state.shard_states[shard].latest_block_root
+    shard_store = store.shard_stores[shard]
+    assert shard_store.block_states[shard_head_root].slot == 1
+    assert shard_store.block_states[shard_head_root] == state.shard_states[shard]
+
+    return store
+
+
+def create_and_apply_shard_block(spec, store, shard, beacon_parent_state, shard_blocks_buffer):
     body = b'\x56' * 4
-    shard_head_root = spec.get_shard_head(store, shard_store)
+    shard_head_root = spec.get_shard_head(store, shard)
+    shard_store = store.shard_stores[shard]
     shard_parent_state = shard_store.block_states[shard_head_root]
     assert shard_parent_state.slot != beacon_parent_state.slot
     shard_block = build_shard_block(
@@ -36,12 +52,12 @@ def apply_shard_block(spec, store, shard_store, beacon_parent_state, shard_block
         shard_parent_state=shard_parent_state, slot=beacon_parent_state.slot, body=body, signed=True
     )
     shard_blocks_buffer.append(shard_block)
-    run_on_shard_block(spec, store, shard_store, shard_block)
-    assert spec.get_shard_head(store, shard_store) == shard_block.message.hash_tree_root()
+    run_on_shard_block(spec, store, shard, shard_block)
+    assert spec.get_shard_head(store, shard) == shard_block.message.hash_tree_root()
 
 
-def check_pending_shard_blocks(spec, store, shard_store, shard_blocks_buffer):
-    pending_shard_blocks = spec.get_pending_shard_blocks(store, shard_store)
+def check_pending_shard_blocks(spec, store, shard, shard_blocks_buffer):
+    pending_shard_blocks = spec.get_pending_shard_blocks(store, shard)
     assert pending_shard_blocks == shard_blocks_buffer
 
 
@@ -52,10 +68,22 @@ def is_in_offset_sets(spec, beacon_head_state, shard):
     return beacon_head_state.slot in offset_slots
 
 
-def apply_shard_and_beacon(spec, state, store, shard_store, shard_blocks_buffer):
-    store.time = store.time + spec.SECONDS_PER_SLOT * spec.SLOTS_PER_EPOCH
+def create_attestation_for_shard_blocks(spec, beacon_parent_state, shard, committee_index, blocks,
+                                        filter_participant_set=None):
+    shard_transition = spec.get_shard_transition(beacon_parent_state, shard, blocks)
+    attestation = get_valid_on_time_attestation(
+        spec,
+        beacon_parent_state,
+        index=committee_index,
+        shard_transition=shard_transition,
+        signed=False,
+    )
+    return attestation
 
-    shard = shard_store.shard
+
+def create_beacon_block_with_shard_transition(
+        spec, state, store, shard, shard_blocks_buffer, is_checking_pending_shard_blocks=True):
+    beacon_block = build_empty_block(spec, state, slot=state.slot + 1)
     committee_index = get_committee_index_of_shard(spec, state, state.slot, shard)
     has_shard_committee = committee_index is not None  # has committee of `shard` at this slot
 
@@ -63,14 +91,12 @@ def apply_shard_and_beacon(spec, state, store, shard_store, shard_blocks_buffer)
 
     # If next slot has committee of `shard`, add `shard_transtion` to the proposing beacon block
     if has_shard_committee and len(shard_blocks_buffer) > 0:
-        # Sanity check `get_pending_shard_blocks` function
-        check_pending_shard_blocks(spec, store, shard_store, shard_blocks_buffer)
+        # Sanity check `get_pending_shard_blocks`
+        # Assert that the pending shard blocks set in the store equal to shard_blocks_buffer
+        if is_checking_pending_shard_blocks:
+            check_pending_shard_blocks(spec, store, shard, shard_blocks_buffer)
         # Use temporary next state to get ShardTransition of shard block
-        shard_transitions = get_shard_transitions(
-            spec,
-            state,
-            shard_block_dict={shard: shard_blocks_buffer},
-        )
+        shard_transitions = get_shard_transitions(spec, state, shard_block_dict={shard: shard_blocks_buffer})
         shard_transition = shard_transitions[shard]
         attestation = get_valid_on_time_attestation(
             spec,
@@ -86,15 +112,31 @@ def apply_shard_and_beacon(spec, state, store, shard_store, shard_blocks_buffer)
         # Clear buffer
         shard_blocks_buffer.clear()
 
-    signed_beacon_block = state_transition_and_sign_block(spec, state, beacon_block)  # transition!
-    add_block_to_store(spec, store, signed_beacon_block)
-    assert spec.get_head(store) == beacon_block.hash_tree_root()
+    return beacon_block
 
-    # On shard block at transitioned `state.slot`
+
+def apply_all_attestation_to_store(spec, store, attestations):
+    for attestation in attestations:
+        spec.on_attestation(store, attestation)
+
+
+def apply_beacon_block_to_store(spec, state, store, beacon_block):
+    signed_beacon_block = state_transition_and_sign_block(spec, state, beacon_block)  # transition!
+    store.time = store.time + spec.SECONDS_PER_SLOT
+    add_block_to_store(spec, store, signed_beacon_block)
+    apply_all_attestation_to_store(spec, store, signed_beacon_block.message.body.attestations)
+
+
+def create_and_apply_beacon_and_shard_blocks(spec, state, store, shard, shard_blocks_buffer):
+    beacon_block = create_beacon_block_with_shard_transition(spec, state, store, shard, shard_blocks_buffer)
+    apply_beacon_block_to_store(spec, state, store, beacon_block)
+
+    # On shard block at the transitioned `state.slot`
     if is_in_offset_sets(spec, state, shard):
         # The created shard block would be appended to `shard_blocks_buffer`
-        apply_shard_block(spec, store, shard_store, state, shard_blocks_buffer)
+        create_and_apply_shard_block(spec, store, shard, state, shard_blocks_buffer)
 
+    has_shard_committee = get_committee_index_of_shard(spec, state, state.slot, shard) is not None
     return has_shard_committee
 
 
@@ -107,23 +149,85 @@ def test_basic(spec, state):
     shard = spec.Shard(1)
 
     # Initialization
-    store = spec.get_forkchoice_store(state)
-    anchor_root = get_anchor_root(spec, state)
-    assert spec.get_head(store) == anchor_root
-
-    shard_store = store.shard_stores[shard]
-    shard_head_root = spec.get_shard_head(store, shard_store)
-    assert shard_head_root == state.shard_states[shard].latest_block_root
-    assert shard_store.block_states[shard_head_root].slot == 1
-    assert shard_store.block_states[shard_head_root] == state.shard_states[shard]
+    store = initialize_store(spec, state, shard)
 
     # For mainnet config, it's possible that only one committee of `shard` per epoch.
     # we set this counter to test more rounds.
     shard_committee_counter = 2
-    shard_blocks_buffer = []
+    shard_blocks_buffer = []  # the accumulated shard blocks that haven't been crosslinked yet
     while shard_committee_counter > 0:
-        has_shard_committee = apply_shard_and_beacon(
-            spec, state, store, shard_store, shard_blocks_buffer
+        has_shard_committee = create_and_apply_beacon_and_shard_blocks(
+            spec, state, store, shard, shard_blocks_buffer
         )
         if has_shard_committee:
             shard_committee_counter -= 1
+
+
+def create_simple_fork(spec, state, store, shard):
+    # Beacon block
+    beacon_block = create_beacon_block_with_shard_transition(spec, state, store, shard, [])
+    apply_beacon_block_to_store(spec, state, store, beacon_block)
+
+    beacon_head_root = spec.get_head(store)
+    assert beacon_head_root == beacon_block.hash_tree_root()
+    beacon_parent_state = store.block_states[beacon_head_root]
+    shard_store = store.shard_stores[shard]
+    shard_parent_state = shard_store.block_states[spec.get_shard_head(store, shard)]
+
+    # Shard block A
+    body = b'\x56' * 4
+    forking_block_child = build_shard_block(
+        spec, beacon_parent_state, shard,
+        shard_parent_state=shard_parent_state, slot=beacon_parent_state.slot, body=body, signed=True
+    )
+    run_on_shard_block(spec, store, shard, forking_block_child)
+
+    # Shard block B
+    body = b'\x78' * 4  # different body
+    shard_block_b = build_shard_block(
+        spec, beacon_parent_state, shard,
+        shard_parent_state=shard_parent_state, slot=beacon_parent_state.slot, body=body, signed=True
+    )
+    run_on_shard_block(spec, store, shard, shard_block_b)
+
+    # Set forking_block
+    current_head = spec.get_shard_head(store, shard)
+    if current_head == forking_block_child.message.hash_tree_root():
+        head_block = forking_block_child
+        forking_block = shard_block_b
+    else:
+        assert current_head == shard_block_b.message.hash_tree_root()
+        head_block = shard_block_b
+        forking_block = forking_block_child
+
+    return head_block, forking_block
+
+
+@with_all_phases_except([PHASE0])
+@spec_state_test
+@never_bls  # Set to never_bls for testing `check_pending_shard_blocks`
+def test_shard_simple_fork(spec, state):
+    if not is_full_crosslink(spec, state):
+        # skip
+        return
+
+    spec.PHASE_1_GENESIS_SLOT = 0  # NOTE: mock genesis slot here
+    state = spec.upgrade_to_phase1(state)
+    shard = spec.Shard(1)
+
+    # Initialization
+    store = initialize_store(spec, state, shard)
+
+    # Create fork
+    _, forking_block = create_simple_fork(spec, state, store, shard)
+
+    # Vote for forking_block
+    state = store.block_states[spec.get_head(store)].copy()
+    beacon_block = create_beacon_block_with_shard_transition(spec, state, store, shard, [forking_block],
+                                                             is_checking_pending_shard_blocks=False)
+    # apply_beacon_block_to_store(spec, state, store, beacon_block)
+    store.time = store.time + spec.SECONDS_PER_SLOT
+    apply_all_attestation_to_store(spec, store, beacon_block.body.attestations)
+
+    # Head block has been changed
+    assert spec.get_shard_head(store, shard) == forking_block.message.hash_tree_root()

From 4dcf5e249814e761095ca868e5a16f85ee4a9791 Mon Sep 17 00:00:00 2001
From: Hsiao-Wei Wang <hwwang156@gmail.com>
Date: Fri, 17 Jul 2020 14:23:58 +0800
Subject: [PATCH 5/7] Add test case of different shards

---
 .../test/fork_choice/test_on_shard_block.py   | 80 +++++++++++++++----
 1 file changed, 63 insertions(+), 17 deletions(-)

diff --git a/tests/core/pyspec/eth2spec/test/fork_choice/test_on_shard_block.py b/tests/core/pyspec/eth2spec/test/fork_choice/test_on_shard_block.py
index dae9246fb..d08f6393b 100644
--- a/tests/core/pyspec/eth2spec/test/fork_choice/test_on_shard_block.py
+++ b/tests/core/pyspec/eth2spec/test/fork_choice/test_on_shard_block.py
@@ -27,16 +27,17 @@ def run_on_shard_block(spec, store, shard, signed_block, valid=True):
     assert shard_store.signed_blocks[hash_tree_root(signed_block.message)] == signed_block
 
 
-def initialize_store(spec, state, shard):
+def initialize_store(spec, state, shards):
     store = spec.get_forkchoice_store(state)
     anchor_root = get_anchor_root(spec, state)
     assert spec.get_head(store) == anchor_root
 
-    shard_head_root = spec.get_shard_head(store, shard)
-    assert shard_head_root == state.shard_states[shard].latest_block_root
-    shard_store = store.shard_stores[shard]
-    assert shard_store.block_states[shard_head_root].slot == 1
-    assert shard_store.block_states[shard_head_root] == state.shard_states[shard]
+    for shard in shards:
+        shard_head_root = spec.get_shard_head(store, shard)
+        assert shard_head_root == state.shard_states[shard].latest_block_root
+        shard_store = store.shard_stores[shard]
+        assert shard_store.block_states[shard_head_root].slot == 1
+        assert shard_store.block_states[shard_head_root] == state.shard_states[shard]
 
     return store
 
@@ -76,7 +77,7 @@ def create_attestation_for_shard_blocks(spec, beacon_parent_state, shard, commit
         beacon_parent_state,
         index=committee_index,
         shard_transition=shard_transition,
-        signed=False,
+        signed=True,
     )
     return attestation
 
@@ -103,7 +104,7 @@ def create_beacon_block_with_shard_transition(
             state,
             index=committee_index,
             shard_transition=shard_transition,
-            signed=False,
+            signed=True,
         )
         assert attestation.data.shard == shard
         beacon_block.body.attestations = [attestation]
@@ -127,8 +128,12 @@ def apply_beacon_block_to_store(spec, state, store, beacon_block):
     apply_all_attestation_to_store(spec, store, signed_beacon_block.message.body.attestations)
 
 
-def create_and_apply_beacon_and_shard_blocks(spec, state, store, shard, shard_blocks_buffer):
-    beacon_block = create_beacon_block_with_shard_transition(spec, state, store, shard, shard_blocks_buffer)
+def create_and_apply_beacon_and_shard_blocks(spec, state, store, shard, shard_blocks_buffer,
+                                             is_checking_pending_shard_blocks=True):
+    beacon_block = create_beacon_block_with_shard_transition(
+        spec, state, store, shard, shard_blocks_buffer,
+        is_checking_pending_shard_blocks=is_checking_pending_shard_blocks
+    )
     apply_beacon_block_to_store(spec, state, store, beacon_block)
 
     # On shard block at the transitioned `state.slot`
@@ -149,16 +154,14 @@ def test_basic(spec, state):
     shard = spec.Shard(1)
 
     # Initialization
-    store = initialize_store(spec, state, shard)
+    store = initialize_store(spec, state, [shard])
 
     # For mainnet config, it's possible that only one committee of `shard` per epoch.
     # we set this counter to test more rounds.
     shard_committee_counter = 2
     shard_blocks_buffer = []  # the accumulated shard blocks that haven't been crosslinked yet
     while shard_committee_counter > 0:
-        has_shard_committee = create_and_apply_beacon_and_shard_blocks(
-            spec, state, store, shard, shard_blocks_buffer
-        )
+        has_shard_committee = create_and_apply_beacon_and_shard_blocks(spec, state, store, shard, shard_blocks_buffer)
         if has_shard_committee:
             shard_committee_counter -= 1
 
@@ -205,7 +208,6 @@ def create_simple_fork(spec, state, store, shard):
 
 @with_all_phases_except([PHASE0])
 @spec_state_test
-@never_bls  # Set to never_bls for testing `check_pending_shard_blocks`
 def test_shard_simple_fork(spec, state):
     if not is_full_crosslink(spec, state):
         # skip
@@ -216,7 +218,7 @@ def test_shard_simple_fork(spec, state):
     shard = spec.Shard(1)
 
     # Initialization
-    store = initialize_store(spec, state, shard)
+    store = initialize_store(spec, state, [shard])
 
     # Create fork
     _, forking_block = create_simple_fork(spec, state, store, shard)
@@ -225,9 +227,53 @@ def test_shard_simple_fork(spec, state):
     state = store.block_states[spec.get_head(store)].copy()
     beacon_block = create_beacon_block_with_shard_transition(spec, state, store, shard, [forking_block],
                                                              is_checking_pending_shard_blocks=False)
-    # apply_beacon_block_to_store(spec, state, store, beacon_block)
     store.time = store.time + spec.SECONDS_PER_SLOT
     apply_all_attestation_to_store(spec, store, beacon_block.body.attestations)
 
     # Head block has been changed
     assert spec.get_shard_head(store, shard) == forking_block.message.hash_tree_root()
+
+
+@with_all_phases_except([PHASE0])
+@spec_state_test
+def test_shard_latest_messages_for_different_shards(spec, state):
+    if not is_full_crosslink(spec, state):
+        # skip
+        return
+
+    spec.PHASE_1_GENESIS_SLOT = 0  # NOTE: mock genesis slot here
+    state = spec.upgrade_to_phase1(state)
+    shard_0 = spec.Shard(0)
+    shard_1 = spec.Shard(1)
+
+    # Initialization
+    store = initialize_store(spec, state, [shard_0, shard_1])
+
+    # Shard 0 ----------------------------------
+    # Create fork on shard 0
+    _, forking_block = create_simple_fork(spec, state, store, shard_0)
+
+    # Vote for forking_block on shard 0
+    state = store.block_states[spec.get_head(store)].copy()
+    beacon_block = create_beacon_block_with_shard_transition(spec, state, store, shard_0, [forking_block],
+                                                             is_checking_pending_shard_blocks=False)
+    store.time = store.time + spec.SECONDS_PER_SLOT
+    apply_all_attestation_to_store(spec, store, beacon_block.body.attestations)
+
+    # Head block of shard 0 has been changed due to the shard latest messages
+    assert spec.get_shard_head(store, shard_0) == forking_block.message.hash_tree_root()
+
+    # Shard 1 ----------------------------------
+    # Run shard 1 after 1~2 epochs
+    shard_committee_counter = 2
+    shard_blocks_buffer = []  # the accumulated shard blocks that haven't been crosslinked yet
+    while shard_committee_counter > 0:
+        has_shard_committee = create_and_apply_beacon_and_shard_blocks(
+            spec, state, store, shard_1, shard_blocks_buffer
+        )
+        if has_shard_committee:
+            shard_committee_counter -= 1
+
+    # Go back to see shard 0 ----------------------------------
+    # The head block of shard 0 should be unchanged.
+    assert spec.get_shard_head(store, shard_0) == forking_block.message.hash_tree_root()

From 33e56b57425a71a13ab20d7b4940135d71d689cc Mon Sep 17 00:00:00 2001
From: Hsiao-Wei Wang <hwwang156@gmail.com>
Date: Wed, 22 Jul 2020 22:48:52 +0800
Subject: [PATCH 6/7] Apply Terence's suggestion: refactor on_shard_block
 interface

---
 specs/phase1/shard-fork-choice.md                   |  9 +++------
 .../test/fork_choice/test_on_shard_block.py         | 13 +++++++------
 2 files changed, 10 insertions(+), 12 deletions(-)

diff --git a/specs/phase1/shard-fork-choice.md b/specs/phase1/shard-fork-choice.md
index 39e957819..ce71ea7de 100644
--- a/specs/phase1/shard-fork-choice.md
+++ b/specs/phase1/shard-fork-choice.md
@@ -142,13 +142,10 @@ def get_pending_shard_blocks(store: Store, shard: Shard) -> Sequence[SignedShard
 #### `on_shard_block`
 
 ```python
-def on_shard_block(store: Store, shard: Shard, signed_shard_block: SignedShardBlock) -> None:
-    shard_store = store.shard_stores[shard]
+def on_shard_block(store: Store, signed_shard_block: SignedShardBlock) -> None:
     shard_block = signed_shard_block.message
-
-    # Check shard
-    # TODO: check it in networking spec
-    assert shard_block.shard == shard
+    shard = shard_block.shard
+    shard_store = store.shard_stores[shard]
 
     # Check shard parent exists
     assert shard_block.shard_parent_root in shard_store.block_states
diff --git a/tests/core/pyspec/eth2spec/test/fork_choice/test_on_shard_block.py b/tests/core/pyspec/eth2spec/test/fork_choice/test_on_shard_block.py
index d08f6393b..0b33517c1 100644
--- a/tests/core/pyspec/eth2spec/test/fork_choice/test_on_shard_block.py
+++ b/tests/core/pyspec/eth2spec/test/fork_choice/test_on_shard_block.py
@@ -13,16 +13,17 @@ from eth2spec.test.helpers.state import state_transition_and_sign_block
 from eth2spec.test.helpers.block import build_empty_block
 
 
-def run_on_shard_block(spec, store, shard, signed_block, valid=True):
+def run_on_shard_block(spec, store, signed_block, valid=True):
+    shard = signed_block.message.shard
     if not valid:
         try:
-            spec.on_shard_block(store, shard, signed_block)
+            spec.on_shard_block(store, signed_block)
         except AssertionError:
             return
         else:
             assert False
 
-    spec.on_shard_block(store, shard, signed_block)
+    spec.on_shard_block(store, signed_block)
     shard_store = store.shard_stores[shard]
     assert shard_store.signed_blocks[hash_tree_root(signed_block.message)] == signed_block
 
@@ -53,7 +54,7 @@ def create_and_apply_shard_block(spec, store, shard, beacon_parent_state, shard_
         shard_parent_state=shard_parent_state, slot=beacon_parent_state.slot, body=body, signed=True
     )
     shard_blocks_buffer.append(shard_block)
-    run_on_shard_block(spec, store, shard, shard_block)
+    run_on_shard_block(spec, store, shard_block)
     assert spec.get_shard_head(store, shard) == shard_block.message.hash_tree_root()
 
 
@@ -183,7 +184,7 @@ def create_simple_fork(spec, state, store, shard):
         spec, beacon_parent_state, shard,
         shard_parent_state=shard_parent_state, slot=beacon_parent_state.slot, body=body, signed=True
     )
-    run_on_shard_block(spec, store, shard, forking_block_child)
+    run_on_shard_block(spec, store, forking_block_child)
 
     # Shard block B
     body = b'\x78' * 4  # different body
@@ -191,7 +192,7 @@ def create_simple_fork(spec, state, store, shard):
         spec, beacon_parent_state, shard,
         shard_parent_state=shard_parent_state, slot=beacon_parent_state.slot, body=body, signed=True
     )
-    run_on_shard_block(spec, store, shard, shard_block_b)
+    run_on_shard_block(spec, store, shard_block_b)
 
     # Set forking_block
     current_head = spec.get_shard_head(store, shard)

From 5d388f7b9bc1dec41eaa79b1f8b11083cb3dabb6 Mon Sep 17 00:00:00 2001
From: Hsiao-Wei Wang <hwwang156@gmail.com>
Date: Wed, 29 Jul 2020 02:03:52 +0800
Subject: [PATCH 7/7] Fix the conflict of #1971

---
 .../{ => phase1/unittests}/fork_choice/test_on_shard_block.py   | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)
 rename tests/core/pyspec/eth2spec/test/{ => phase1/unittests}/fork_choice/test_on_shard_block.py (99%)

diff --git a/tests/core/pyspec/eth2spec/test/fork_choice/test_on_shard_block.py b/tests/core/pyspec/eth2spec/test/phase1/unittests/fork_choice/test_on_shard_block.py
similarity index 99%
rename from tests/core/pyspec/eth2spec/test/fork_choice/test_on_shard_block.py
rename to tests/core/pyspec/eth2spec/test/phase1/unittests/fork_choice/test_on_shard_block.py
index 0b33517c1..de47aaa39 100644
--- a/tests/core/pyspec/eth2spec/test/fork_choice/test_on_shard_block.py
+++ b/tests/core/pyspec/eth2spec/test/phase1/unittests/fork_choice/test_on_shard_block.py
@@ -37,7 +37,7 @@ def initialize_store(spec, state, shards):
         shard_head_root = spec.get_shard_head(store, shard)
         assert shard_head_root == state.shard_states[shard].latest_block_root
         shard_store = store.shard_stores[shard]
-        assert shard_store.block_states[shard_head_root].slot == 1
+        assert shard_store.block_states[shard_head_root].slot == 0
         assert shard_store.block_states[shard_head_root] == state.shard_states[shard]
 
     return store