From 79b1b4bdbe0c4822d4a2e197a7bf59930ffe69c1 Mon Sep 17 00:00:00 2001
From: Hsiao-Wei Wang <hwwang156@gmail.com>
Date: Fri, 1 May 2020 10:52:45 +0800
Subject: [PATCH] Add `(shard, shard_root)` to `LatestMessage`

---
 specs/phase1/fork-choice.md                   | 32 ++++++++++++++++++-
 .../test/fork_choice/test_on_attestation.py   | 13 ++++++--
 2 files changed, 41 insertions(+), 4 deletions(-)

diff --git a/specs/phase1/fork-choice.md b/specs/phase1/fork-choice.md
index d8bf7fa09..f4e771ddb 100644
--- a/specs/phase1/fork-choice.md
+++ b/specs/phase1/fork-choice.md
@@ -10,6 +10,9 @@
 
 - [Introduction](#introduction)
 - [Fork choice](#fork-choice)
+  - [Helpers](#helpers)
+    - [Extended `LatestMessage`](#extended-latestmessage)
+    - [Updated `update_latest_messages`](#updated-update_latest_messages)
   - [Handlers](#handlers)
 
 <!-- END doctoc generated TOC please keep comment here to allow auto update -->
@@ -25,6 +28,33 @@ Due to the changes in the structure of `IndexedAttestation` in Phase 1, `on_atte
 
 The rest of the fork choice remains stable.
 
+### Helpers
+
+#### Extended `LatestMessage`
+
+```python
+@dataclass(eq=True, frozen=True)
+class LatestMessage(object):
+    epoch: Epoch
+    root: Root
+    shard: Shard
+    shard_root: Root
+```
+
+#### Updated `update_latest_messages`
+
+```python
+def update_latest_messages(store: Store, attesting_indices: Sequence[ValidatorIndex], attestation: Attestation) -> None:
+    target = attestation.data.target
+    beacon_block_root = attestation.data.beacon_block_root
+    shard = get_shard(store.block_states[beacon_block_root], attestation)
+    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.head_shard_root
+            )
+```
+
 ### Handlers
 
 ```python
@@ -49,4 +79,4 @@ def on_attestation(store: Store, attestation: Attestation) -> None:
         if attestation.aggregation_bits[i]
     ]
     update_latest_messages(store, attesting_indices, attestation)
-```
\ No newline at end of file
+```
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 360c18ccd..1bce42ca3 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
@@ -18,18 +18,25 @@ def run_on_attestation(spec, state, store, attestation, valid=True):
 
     if spec.fork == PHASE0:
         sample_index = indexed_attestation.attesting_indices[0]
+        latest_message = spec.LatestMessage(
+            epoch=attestation.data.target.epoch,
+            root=attestation.data.beacon_block_root,
+        )
     else:
         attesting_indices = [
             index for i, index in enumerate(indexed_attestation.committee)
             if attestation.aggregation_bits[i]
         ]
         sample_index = attesting_indices[0]
-    assert (
-        store.latest_messages[sample_index] ==
-        spec.LatestMessage(
+        latest_message = spec.LatestMessage(
             epoch=attestation.data.target.epoch,
             root=attestation.data.beacon_block_root,
+            shard=spec.get_shard(state, attestation),
+            shard_root=attestation.data.head_shard_root,
         )
+
+    assert (
+        store.latest_messages[sample_index] == latest_message
     )