From 859a0dbbbe459c6fca1e1cc5c0e454c9ae144529 Mon Sep 17 00:00:00 2001
From: Hsiao-Wei Wang <hwwang156@gmail.com>
Date: Mon, 20 Aug 2018 16:32:46 +0800
Subject: [PATCH] Added `get_indices_for_slot` and `get_block_hash`

---
 beacon_chain/state/helpers.py |  23 ++++++++
 tests/state/test_helpers.py   | 101 +++++++++++++++++++++++++++++++++-
 2 files changed, 122 insertions(+), 2 deletions(-)

diff --git a/beacon_chain/state/helpers.py b/beacon_chain/state/helpers.py
index e511545..5323221 100644
--- a/beacon_chain/state/helpers.py
+++ b/beacon_chain/state/helpers.py
@@ -131,3 +131,26 @@ def get_new_shuffling(seed,
             committee=indices
         ) for j, indices in enumerate(shard_indices)])
     return o
+
+
+def get_indices_for_slot(
+        crystallized_state,
+        slot,
+        config=DEFAULT_CONFIG):
+    cycle_length = config['cycle_length']
+
+    ifh_start = crystallized_state.last_state_recalc - cycle_length
+    assert ifh_start <= slot < ifh_start + cycle_length * 2
+    return crystallized_state.indices_for_slots[slot - ifh_start]
+
+
+def get_block_hash(
+        active_state,
+        current_block,
+        slot,
+        config=DEFAULT_CONFIG):
+    cycle_length = config['cycle_length']
+
+    sback = current_block.slot_number - cycle_length * 2
+    assert sback <= slot < sback + cycle_length * 2
+    return active_state.recent_block_hashes[slot - sback]
diff --git a/tests/state/test_helpers.py b/tests/state/test_helpers.py
index 45646c5..6196aef 100644
--- a/tests/state/test_helpers.py
+++ b/tests/state/test_helpers.py
@@ -1,7 +1,13 @@
 import pytest
 
+from beacon_chain.utils.blake import blake
+
+from beacon_chain.state.active_state import ActiveState
+from beacon_chain.state.block import Block
 from beacon_chain.state.helpers import (
-    get_new_shuffling
+    get_new_shuffling,
+    get_indices_for_slot,
+    get_block_hash,
 )
 
 
@@ -28,7 +34,12 @@ def test_get_new_shuffling_is_complete(genesis_validators, config):
     )
 
     assert len(shuffling) == config['cycle_length']
-
+    print([
+        [j.committee 
+        for j in item]
+        for item in shuffling]
+    )
+    print('shuffling: {}'.format(shuffling))
     validators = set()
     shards = set()
     for height_indices in shuffling:
@@ -39,3 +50,89 @@ def test_get_new_shuffling_is_complete(genesis_validators, config):
 
     # assert len(shards) == config['shard_count']
     assert len(validators) == len(genesis_validators)
+
+
+@pytest.mark.parametrize(
+    (
+        'num_validators,slot,success'
+    ),
+    [
+        (100, 0, True),
+        (100, 63, True),
+        (100, 64, False),
+    ],
+)
+def test_get_indices_for_slot(
+        genesis_crystallized_state,
+        num_validators,
+        slot,
+        success,
+        config):
+    crystallized_state = genesis_crystallized_state
+
+    if success:
+        indices_for_slot = get_indices_for_slot(
+            crystallized_state,
+            slot,
+            config=config,
+        )
+        assert len(indices_for_slot) > 0
+    else:
+        with pytest.raises(AssertionError):
+            get_indices_for_slot(
+                crystallized_state,
+                slot,
+                config=config,
+            )
+
+
+@pytest.mark.parametrize(
+    (
+        'slot,success'
+    ),
+    [
+        (0, True),
+        (127, True),
+        (128, False),
+    ],
+)
+def test_get_block_hash(
+        genesis_block,
+        slot,
+        success,
+        config):
+    cycle_length = config['cycle_length']
+
+    blocks = get_empty_chain(cycle_length * 3)
+    active_state = ActiveState(
+        recent_block_hashes=[block.hash for block in blocks[:cycle_length*2]]
+    )
+    if success:
+        block_hash = get_block_hash(
+            active_state,
+            blocks[cycle_length*2],
+            slot,
+            config=config,
+        )
+        assert block_hash == blocks[slot].hash
+    else:
+        with pytest.raises(AssertionError):
+            get_block_hash(
+                active_state,
+                blocks[cycle_length*2],
+                slot,
+                config=config,
+            )
+
+
+def get_empty_chain(length):
+    blocks = []
+    for slot in range(length * 3):
+        blocks.append(
+            Block(
+                slot_number=slot,
+                parent_hash=blocks[slot-1].hash if slot > 0 else b'00'*32
+            )
+        )
+
+    return blocks