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

@@ -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
 ```

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
         )
     ))

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

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