Cache initial trackers

tests/core/pyspec/eth2spec/test/helpers/block.py

@@ -17,7 +17,11 @@ from py_ecc.bls.g2_primitives import (
     G1_to_pubkey as py_ecc_G1_to_bytes48,
     pubkey_to_G1 as py_ecc_bytes48_to_G1,
 )
-from eth2spec.test.helpers.whisk import get_whisk_tracker_and_commitment, is_first_proposal
+from eth2spec.test.helpers.whisk import (
+    compute_whisk_tracker_and_commitment,
+    is_first_proposal,
+    resolve_known_tracker
+)
 
 PointProjective = Optimized_Point3D[Optimized_Field]
 
@@ -179,7 +183,7 @@ def build_empty_block(spec, state, slot=None, proposer_index=None):
             k_final = whisk_ks_final[proposer_index]
             # TODO: Actual logic should pick a random r, but may need to do something fancy to locate trackers quickly
             r = 2
-            tracker, k_commitment = get_whisk_tracker_and_commitment(k_final, r)
+            tracker, k_commitment = compute_whisk_tracker_and_commitment(k_final, r)
             empty_block.body.whisk_registration_proof = GenerateWhiskTrackerProof(tracker, k_final)
             empty_block.body.whisk_tracker = tracker
             empty_block.body.whisk_k_commitment = k_commitment
@@ -208,8 +212,18 @@ def get_beacon_proposer_to_build(spec, state, proposer_index=None):
 
 
 def find_whisk_proposer(spec, state):
-    raise Exception("proposer not known without heavy math")
-    # proposer_tracker = state.whisk_proposer_trackers[state.slot % spec.WHISK_PROPOSER_TRACKERS_COUNT]
+    proposer_tracker = state.whisk_proposer_trackers[state.slot % spec.WHISK_PROPOSER_TRACKERS_COUNT]
+
+    # Check record of known trackers
+    # During the first shuffling phase (epoch < WHISK_EPOCHS_PER_SHUFFLING_PHASE)
+    # proposer trackers are those inserted on the genesis state, and have not gone
+    # through any shuffling. We cache those initial trackers and use `resolve_known_tracker`
+    # to check if the tracker is known, and skip the need to actually find the matching tracker
+    proposer_index = resolve_known_tracker(proposer_tracker)
+    if proposer_index is not None:
+        return proposer_index
+
+    print("proposer_tracker", proposer_tracker)
     # # First attempt direct equality with trackers
     # for i, validator in enumerate(state.validators):
     #     # # This is insanely slow
@@ -218,6 +232,7 @@ def find_whisk_proposer(spec, state):
     #         return i
     # # In Whisk where to get proposer from?
     # raise Exception("proposer_tracker not matched")
+    raise Exception("proposer not known without heavy math")
 
 
 def build_empty_block_for_next_slot(spec, state, proposer_index=None):

tests/core/pyspec/eth2spec/test/helpers/genesis.py

@@ -8,6 +8,7 @@ from eth2spec.test.helpers.forks import (
     is_post_altair, is_post_bellatrix, is_post_capella, is_post_eip6110, is_post_whisk
 )
 from eth2spec.test.helpers.keys import pubkeys
+from eth2spec.test.helpers.whisk import compute_whisk_initial_tracker_cached, compute_whisk_initial_k_commitment_cached
 
 
 def build_mock_validator(spec, i: int, balance: int):
@@ -142,10 +143,15 @@ def create_genesis_state(spec, validator_balances, activation_threshold):
         state.deposit_receipts_start_index = spec.UNSET_DEPOSIT_RECEIPTS_START_INDEX
 
     if is_post_whisk(spec):
-        zero_commitment = spec.BLSPubkey()
-        zero_tracker = spec.WhiskTracker()
-        for _ in range(len(state.validators)):
-            state.whisk_k_commitments.append(zero_commitment)
-            state.whisk_trackers.append(zero_tracker)
+        vc = len(state.validators)
+        for i in range(vc):
+            state.whisk_k_commitments.append(compute_whisk_initial_k_commitment_cached(i))
+            state.whisk_trackers.append(compute_whisk_initial_tracker_cached(i))
+
+        for i in range(spec.WHISK_CANDIDATE_TRACKERS_COUNT):
+            state.whisk_candidate_trackers[i] = compute_whisk_initial_tracker_cached(i % vc)
+
+        for i in range(spec.WHISK_PROPOSER_TRACKERS_COUNT):
+            state.whisk_proposer_trackers[i] = compute_whisk_initial_tracker_cached(i % vc)
 
     return state

tests/core/pyspec/eth2spec/test/helpers/whisk.py

@@ -1,21 +1,57 @@
-from typing import Tuple
+from typing import Tuple, Optional
 from eth_typing import BLSPubkey
 from py_ecc.optimized_bls12_381.optimized_curve import G1, multiply
 from py_ecc.bls.g2_primitives import G1_to_pubkey as py_ecc_G1_to_bytes48
 from curdleproofs import GenerateWhiskTrackerProof, WhiskTracker
+from eth2spec.test.helpers.keys import whisk_ks_initial
 
 
-def get_whisk_k_commitment(k: int) -> BLSPubkey:
+# Map of validator index to initial WhiskTracker (r = 1, k = index)
+whisk_initial_tracker_cache_by_index = {}
+# Map of validator index to k commitment (k = index)
+whisk_initial_k_commitment_cache_by_index = {}
+# Map of k_r_G to validator index
+whisk_initial_tracker_cache_by_k_r_G = {}
+INITIAL_R = 1
+
+
+def compute_whisk_initial_tracker_cached(i: int) -> WhiskTracker:
+    if i in whisk_initial_tracker_cache_by_index:
+        return whisk_initial_tracker_cache_by_index[i]
+
+    tracker = compute_whisk_tracker(whisk_ks_initial[i], INITIAL_R)
+    whisk_initial_tracker_cache_by_index[i] = tracker
+    whisk_initial_tracker_cache_by_k_r_G[tracker.k_r_G] = i
+    return tracker
+
+
+def compute_whisk_initial_k_commitment_cached(i: int) -> BLSPubkey:
+    if i in whisk_initial_k_commitment_cache_by_index:
+        return whisk_initial_k_commitment_cache_by_index[i]
+
+    commitment = compute_whisk_k_commitment(whisk_ks_initial[i])
+    whisk_initial_k_commitment_cache_by_index[i] = commitment
+    return commitment
+
+
+def resolve_known_tracker(tracker: WhiskTracker) -> Optional[int]:
+    if tracker.k_r_G in whisk_initial_tracker_cache_by_k_r_G:
+        return whisk_initial_tracker_cache_by_k_r_G[tracker.k_r_G]
+    else:
+        return None
+
+
+def compute_whisk_k_commitment(k: int) -> BLSPubkey:
     return py_ecc_G1_to_bytes48(multiply(G1, int(k)))
 
 
-def get_whisk_tracker(k: int, r: int) -> WhiskTracker:
+def compute_whisk_tracker(k: int, r: int) -> WhiskTracker:
     r_G = multiply(G1, int(r))
     k_r_G = multiply(r_G, int(k))
     return WhiskTracker(py_ecc_G1_to_bytes48(r_G), py_ecc_G1_to_bytes48(k_r_G))
 
 
-def get_whisk_tracker_and_commitment(k: int, r: int) -> Tuple[WhiskTracker, BLSPubkey]:
+def compute_whisk_tracker_and_commitment(k: int, r: int) -> Tuple[WhiskTracker, BLSPubkey]:
     k_G = multiply(G1, int(k))
     r_G = multiply(G1, int(r))
     k_r_G = multiply(r_G, int(k))
@@ -36,14 +72,14 @@ def is_first_proposal(spec, state, proposer_index: int) -> bool:
 
 
 def set_registration(body, k: int, r: int):
-    tracker, k_commitment = get_whisk_tracker_and_commitment(k, r)
+    tracker, k_commitment = compute_whisk_tracker_and_commitment(k, r)
     body.whisk_registration_proof = GenerateWhiskTrackerProof(tracker, k)
     body.whisk_tracker = tracker
     body.whisk_k_commitment = k_commitment
 
 
 def set_opening_proof(spec, state, block, proposer_index: int, k: int, r: int):
-    tracker, k_commitment = get_whisk_tracker_and_commitment(k, r)
+    tracker, k_commitment = compute_whisk_tracker_and_commitment(k, r)
     state.whisk_proposer_trackers[state.slot % spec.WHISK_PROPOSER_TRACKERS_COUNT] = tracker
     state.whisk_k_commitments[proposer_index] = k_commitment
     block.proposer_index = proposer_index

tests/core/pyspec/eth2spec/test/whisk/block_processing/test_process_shuffled_trackers.py

@@ -1,6 +1,6 @@
 from eth2spec.test.context import spec_state_test, with_whisk_and_later, expect_assertion_error
 from eth2spec.test.helpers.keys import whisk_ks_initial
-from eth2spec.test.helpers.whisk import get_whisk_tracker
+from eth2spec.test.helpers.whisk import compute_whisk_tracker
 from curdleproofs import GenerateWhiskShuffleProof
 
 
@@ -18,7 +18,7 @@ def get_and_populate_pre_shuffle_trackers(spec, state, body):
     pre_shuffle_trackers = []
     for i in shuffle_indices:
         # Set r to some value > 1 ( = 2+i)
-        tracker = get_whisk_tracker(whisk_ks_initial[i], 2 + i)
+        tracker = compute_whisk_tracker(whisk_ks_initial[i], 2 + i)
         state.whisk_candidate_trackers[i] = tracker
         pre_shuffle_trackers.append(tracker)
     return pre_shuffle_trackers

tests/core/pyspec/eth2spec/test/whisk/block_processing/test_process_whisk_opening_proof.py

@@ -1,7 +1,7 @@
 from eth2spec.test.context import spec_state_test, with_whisk_and_later, expect_assertion_error
 from eth2spec.test.helpers.whisk import (
-    get_whisk_k_commitment,
-    get_whisk_tracker,
+    compute_whisk_k_commitment,
+    compute_whisk_tracker,
     set_opening_proof
 )
 
@@ -44,7 +44,7 @@ def test_valid_proof(spec, state):
 def test_wrong_commitment(spec, state):
     block = empty_block(spec)
     set_opening_proof(spec, state, block, PROPOSER_INDEX, K_OK, R_OK)
-    state.whisk_k_commitments[PROPOSER_INDEX] = get_whisk_k_commitment(K_WRONG)
+    state.whisk_k_commitments[PROPOSER_INDEX] = compute_whisk_k_commitment(K_WRONG)
     run_process_whisk_opening_proof(spec, state, block, valid=False)
 
 
@@ -53,7 +53,8 @@ def test_wrong_commitment(spec, state):
 def test_wrong_tracker_r(spec, state):
     block = empty_block(spec)
     set_opening_proof(spec, state, block, PROPOSER_INDEX, K_OK, R_OK)
-    state.whisk_proposer_trackers[state.slot % spec.WHISK_PROPOSER_TRACKERS_COUNT] = get_whisk_tracker(K_OK, R_WRONG)
+    wrong_tracker = compute_whisk_tracker(K_OK, R_WRONG)
+    state.whisk_proposer_trackers[state.slot % spec.WHISK_PROPOSER_TRACKERS_COUNT] = wrong_tracker
     run_process_whisk_opening_proof(spec, state, block, valid=False)
 
 

tests/core/pyspec/eth2spec/test/whisk/block_processing/test_process_whisk_registration.py

@@ -1,5 +1,5 @@
 from eth2spec.test.context import spec_state_test, with_whisk_and_later, expect_assertion_error
-from eth2spec.test.helpers.whisk import set_as_first_proposal, get_whisk_k_commitment, set_registration
+from eth2spec.test.helpers.whisk import set_as_first_proposal, compute_whisk_k_commitment, set_registration
 
 
 def empty_block_body(spec):
@@ -57,7 +57,7 @@ def test_first_proposal_indentity_tracker(spec, state):
 def test_first_proposal_non_unique_k_other(spec, state):
     body = empty_block_body(spec)
     set_as_first_proposal_and_proposer(spec, state, PROPOSER_INDEX)
-    state.whisk_k_commitments[OTHER_INDEX] = get_whisk_k_commitment(OTHER_K)
+    state.whisk_k_commitments[OTHER_INDEX] = compute_whisk_k_commitment(OTHER_K)
     set_registration(body, OTHER_K, OTHER_R)
     yield from run_process_whisk_registration(spec, state, body, valid=False)
 
@@ -67,7 +67,7 @@ def test_first_proposal_non_unique_k_other(spec, state):
 def test_first_proposal_non_unique_k_self(spec, state):
     body = empty_block_body(spec)
     set_as_first_proposal_and_proposer(spec, state, PROPOSER_INDEX)
-    state.whisk_k_commitments[PROPOSER_INDEX] = get_whisk_k_commitment(OTHER_K)
+    state.whisk_k_commitments[PROPOSER_INDEX] = compute_whisk_k_commitment(OTHER_K)
     set_registration(body, OTHER_K, OTHER_R)
     yield from run_process_whisk_registration(spec, state, body, valid=False)
 

tests/core/pyspec/eth2spec/test/whisk/sanity/blocks/test_whisk.py

@@ -2,7 +2,7 @@ from eth2spec.test.helpers.block import build_empty_block
 from eth2spec.test.context import spec_state_test, with_whisk_and_later
 from eth2spec.test.helpers.keys import whisk_ks_initial
 from eth2spec.test.helpers.state import state_transition_and_sign_block
-from eth2spec.test.helpers.whisk import get_whisk_tracker_and_commitment
+from eth2spec.test.helpers.whisk import compute_whisk_tracker_and_commitment
 from curdleproofs import WhiskTracker
 
 known_whisk_trackers = {}
@@ -36,7 +36,7 @@ def fill_candidate_trackers(spec, state, tracker: WhiskTracker):
 def test_whisk__process_block_single_initial(spec, state):
     assert state.slot == 0
     proposer_slot_1 = 0
-    tracker_slot_1, k_commitment = get_whisk_tracker_and_commitment(whisk_ks_initial[proposer_slot_1], 1)
+    tracker_slot_1, k_commitment = compute_whisk_tracker_and_commitment(whisk_ks_initial[proposer_slot_1], 1)
     state.whisk_k_commitments[proposer_slot_1] = k_commitment
     state.whisk_proposer_trackers[1] = tracker_slot_1
     fill_candidate_trackers(spec, state, tracker_slot_1)