Add (shard, shard_root) to LatestMessage

ethereum · May 5, 2020 · 79b1b4b · 79b1b4b
1 parent fca1bbc
commit 79b1b4b
Show file tree

Hide file tree

Showing 2 changed files with 41 additions and 4 deletions.
diff --git 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)
-```
+```
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
@@ -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
     )