leak state decorator, and test pre-state caching

tests/core/pyspec/eth2spec/test/context.py

@@ -9,6 +9,8 @@ from .utils import vector_test, with_meta_tags
 from random import Random
 from typing import Any, Callable, NewType, Sequence, TypedDict, Protocol
 
+from lru import LRU
+
 from importlib import reload
 
 
@@ -48,28 +50,45 @@ class SpecForks(TypedDict, total=False):
     PHASE1: SpecPhase1
 
 
+def _prepare_state(balances_fn: Callable[[Any], Sequence[int]], threshold_fn: Callable[[Any], int],
+                   spec: Spec, phases: SpecForks):
+
+    p0 = phases[PHASE0]
+    balances = balances_fn(p0)
+    activation_threshold = threshold_fn(p0)
+
+    state = create_genesis_state(spec=p0, validator_balances=balances,
+                                 activation_threshold=activation_threshold)
+    if spec.fork == PHASE1:
+        # TODO: instead of upgrading a test phase0 genesis state we can also write a phase1 state helper.
+        # Decide based on performance/consistency results later.
+        state = phases[PHASE1].upgrade_to_phase1(state)
+        # Shard state slot must lag behind BeaconState slot by at least 1
+        # Will handle this more elegantly with fork mechanics
+        spec.process_slots(state, state.slot + 1)
+
+    return state
+
+
+_custom_state_cache_dict = LRU(size=10)
+
+
 def with_custom_state(balances_fn: Callable[[Any], Sequence[int]],
                       threshold_fn: Callable[[Any], int]):
     def deco(fn):
+
         def entry(*args, spec: Spec, phases: SpecForks, **kw):
-            try:
+            # Use fork and file path to make a key for th
-                p0 = phases[PHASE0]
+            key = (spec.fork, spec.__file__, balances_fn, threshold_fn)
-                balances = balances_fn(p0)
+            global _custom_state_cache_dict
-                activation_threshold = threshold_fn(p0)
+            if key not in _custom_state_cache_dict:
+                state = _prepare_state(balances_fn, threshold_fn, spec, phases)
+                _custom_state_cache_dict[key] = state.get_backing()
 
-                state = create_genesis_state(spec=p0, validator_balances=balances,
+            # Take a copy out of the LRU cache result.
-                                             activation_threshold=activation_threshold)
+            # No copy is necessary, as we wrap the immutable backing with a new view.
-                if spec.fork == PHASE1:
+            state = spec.BeaconState(backing=_custom_state_cache_dict[key])
-                    # TODO: instead of upgrading a test phase0 genesis state we can also write a phase1 state helper.
+            kw['state'] = state
-                    # Decide based on performance/consistency results later.
-                    state = phases[PHASE1].upgrade_to_phase1(state)
-                    # Shard state slot must lag behind BeaconState slot by at least 1
-                    # Will handle this more elegantly with fork mechanics
-                    spec.process_slots(state, state.slot + 1)
-
-                kw['state'] = state
-            except KeyError:
-                raise TypeError('Spec decorator must come within state decorator to inject spec into state.')
             return fn(*args, spec=spec, phases=phases, **kw)
         return entry
     return deco

tests/core/pyspec/eth2spec/test/phase_0/rewards/test_leak.py

@@ -1,6 +1,7 @@
 from eth2spec.test.context import with_all_phases, spec_state_test
 from eth2spec.test.helpers.state import next_epoch
 import eth2spec.test.helpers.rewards as rewards_helpers
+from lru import LRU
 
 
 def transition_state_to_leak(spec, state, epochs=None):
@@ -12,80 +13,103 @@ def transition_state_to_leak(spec, state, epochs=None):
         next_epoch(spec, state)
 
 
+_cache_dict = LRU(size=10)
+
+
+def leaking(epochs=None):
+
+    def deco(fn):
+        def entry(*args, spec, state, **kw):
+            # If the pre-state is not already known in the LRU, then take it, make it leaking, and put it in the LRU.
+            # The input state is likely already cached, so the hash-tree-root is fine.
+            key = (state.hash_tree_root(), spec.MIN_EPOCHS_TO_INACTIVITY_PENALTY, spec.SLOTS_PER_EPOCH, epochs)
+            global _cache_dict
+            if key not in _cache_dict:
+                transition_state_to_leak(spec, state, epochs=epochs)
+                _cache_dict[key] = state.get_backing()
+
+            # Take a copy out of the LRU cache result.
+            # No copy is necessary, as we wrap the immutable backing with a new view.
+            state = spec.BeaconState(backing=_cache_dict[key])
+            return fn(*args, spec=spec, state=state, **kw)
+        return entry
+    return deco
+
+
 @with_all_phases
 @spec_state_test
+@leaking()
 def test_empty_leak(spec, state):
-    transition_state_to_leak(spec, state)
     yield from rewards_helpers.run_test_empty(spec, state)
 
 
 @with_all_phases
 @spec_state_test
+@leaking()
 def test_full_leak(spec, state):
-    transition_state_to_leak(spec, state)
     yield from rewards_helpers.run_test_full_all_correct(spec, state)
 
 
 @with_all_phases
 @spec_state_test
+@leaking()
 def test_half_full_leak(spec, state):
-    transition_state_to_leak(spec, state)
     yield from rewards_helpers.run_test_half_full(spec, state)
 
 
 @with_all_phases
 @spec_state_test
+@leaking()
 def test_quarter_full_leak(spec, state):
-    transition_state_to_leak(spec, state)
     yield from rewards_helpers.run_test_partial(spec, state, 0.25)
 
 
 @with_all_phases
 @spec_state_test
+@leaking()
 def test_full_but_partial_participation_leak(spec, state):
-    transition_state_to_leak(spec, state)
     yield from rewards_helpers.run_test_full_but_partial_participation(spec, state)
 
 
 @with_all_phases
 @spec_state_test
+@leaking()
 def test_one_attestation_one_correct_leak(spec, state):
-    transition_state_to_leak(spec, state)
     yield from rewards_helpers.run_test_one_attestation_one_correct(spec, state)
 
 
 @with_all_phases
 @spec_state_test
+@leaking()
 def test_with_not_yet_activated_validators_leak(spec, state):
-    transition_state_to_leak(spec, state)
     yield from rewards_helpers.run_test_with_not_yet_activated_validators(spec, state)
 
 
 @with_all_phases
 @spec_state_test
+@leaking()
 def test_with_exited_validators_leak(spec, state):
-    transition_state_to_leak(spec, state)
     yield from rewards_helpers.run_test_with_exited_validators(spec, state)
 
 
 @with_all_phases
 @spec_state_test
+@leaking()
 def test_with_slashed_validators_leak(spec, state):
-    transition_state_to_leak(spec, state)
     yield from rewards_helpers.run_test_with_slashed_validators(spec, state)
 
 
 @with_all_phases
 @spec_state_test
+@leaking()
 def test_some_very_low_effective_balances_that_attested_leak(spec, state):
-    transition_state_to_leak(spec, state)
     yield from rewards_helpers.run_test_some_very_low_effective_balances_that_attested(spec, state)
 
 
 @with_all_phases
 @spec_state_test
+@leaking()
 def test_some_very_low_effective_balances_that_did_not_attest_leak(spec, state):
-    transition_state_to_leak(spec, state)
     yield from rewards_helpers.run_test_some_very_low_effective_balances_that_did_not_attest(spec, state)
 
 
@@ -98,8 +122,8 @@ def test_some_very_low_effective_balances_that_did_not_attest_leak(spec, state):
 
 @with_all_phases
 @spec_state_test
+@leaking()
 def test_full_half_correct_target_incorrect_head_leak(spec, state):
-    transition_state_to_leak(spec, state)
     yield from rewards_helpers.run_test_full_fraction_incorrect(
         spec, state,
         correct_target=True,
@@ -110,8 +134,8 @@ def test_full_half_correct_target_incorrect_head_leak(spec, state):
 
 @with_all_phases
 @spec_state_test
+@leaking()
 def test_full_correct_target_incorrect_head_leak(spec, state):
-    transition_state_to_leak(spec, state)
     yield from rewards_helpers.run_test_full_fraction_incorrect(
         spec, state,
         correct_target=True,
@@ -122,8 +146,8 @@ def test_full_correct_target_incorrect_head_leak(spec, state):
 
 @with_all_phases
 @spec_state_test
+@leaking()
 def test_full_half_incorrect_target_incorrect_head_leak(spec, state):
-    transition_state_to_leak(spec, state)
     yield from rewards_helpers.run_test_full_fraction_incorrect(
         spec, state,
         correct_target=False,
@@ -134,8 +158,8 @@ def test_full_half_incorrect_target_incorrect_head_leak(spec, state):
 
 @with_all_phases
 @spec_state_test
+@leaking()
 def test_full_half_incorrect_target_correct_head_leak(spec, state):
-    transition_state_to_leak(spec, state)
     yield from rewards_helpers.run_test_full_fraction_incorrect(
         spec, state,
         correct_target=False,
@@ -146,20 +170,20 @@ def test_full_half_incorrect_target_correct_head_leak(spec, state):
 
 @with_all_phases
 @spec_state_test
+@leaking()
 def test_full_random_leak(spec, state):
-    transition_state_to_leak(spec, state)
     yield from rewards_helpers.run_test_full_random(spec, state)
 
 
 @with_all_phases
 @spec_state_test
+@leaking(epochs=5)
 def test_full_random_five_epoch_leak(spec, state):
-    transition_state_to_leak(spec, state, epochs=5)
     yield from rewards_helpers.run_test_full_random(spec, state)
 
 
 @with_all_phases
 @spec_state_test
+@leaking(epochs=10)
 def test_full_random_ten_epoch_leak(spec, state):
-    transition_state_to_leak(spec, state, epochs=10)
     yield from rewards_helpers.run_test_full_random(spec, state)