Simplify previous hash lookups w/ cached iterable

Author: carver

eth/_utils/generator.py

@@ -1,5 +1,5 @@
 import itertools
-from typing import (
+from typing import (  # noqa: F401
     Generic,
     Iterable,
     Iterator,
@@ -12,14 +12,14 @@
 
 
 class CachedIterable(Generic[TItem], Iterable[TItem]):
-    def __init__(self, iterator: Iterable[TItem]) -> None:
-        self._cached_results: List[TItem] = []
-        self._iterator = iterator
+    def __init__(self, iterable: Iterable[TItem]) -> None:
+        self._cached_results = []  # type: List[TItem]
+        self._iterator = iter(iterable)
 
     def __iter__(self) -> Iterator[TItem]:
         return itertools.chain(self._cached_results, self._cache_and_yield())
 
-    def _cache_and_yield(self) -> Iterable[TItem]:
+    def _cache_and_yield(self) -> Iterator[TItem]:
         for item in self._iterator:
             self._cached_results.append(item)
             yield item

eth/vm/execution_context.py

@@ -7,6 +7,8 @@
     Hash32,
 )
 
+from eth._utils.generator import CachedIterable
+
 
 class ExecutionContext:
     _coinbase = None
@@ -30,7 +32,7 @@ def __init__(
         self._block_number = block_number
         self._difficulty = difficulty
         self._gas_limit = gas_limit
-        self._prev_hashes = prev_hashes
+        self._prev_hashes = CachedIterable(prev_hashes)
 
     @property
     def coinbase(self) -> Address:

eth/vm/state.py

@@ -3,7 +3,6 @@
     abstractmethod
 )
 import contextlib
-import itertools
 import logging
 from typing import (  # noqa: F401
     cast,
@@ -20,6 +19,7 @@
     Address,
     Hash32,
 )
+from eth_utils.toolz import nth
 
 from eth.constants import (
     BLANK_ROOT_HASH,
@@ -85,14 +85,11 @@ class for vm execution.
     transaction_context_class = None  # type: Type[BaseTransactionContext]
     account_db_class = None  # type: Type[BaseAccountDB]
     transaction_executor = None  # type: Type[BaseTransactionExecutor]
-    prev_hashes_generator_backup = None
 
     def __init__(self, db: BaseDB, execution_context: ExecutionContext, state_root: bytes) -> None:
         self._db = db
         self.execution_context = execution_context
         self.account_db = self.get_account_db_class()(self._db, state_root)
-        if execution_context:
-            self.prev_hashes_generator_backup = execution_context.prev_hashes
 
     #
     # Logging
@@ -201,12 +198,6 @@ def get_ancestor_hash(self, block_number: int) -> Hash32:
         Return the empty bytestring ``b''`` if the block number is outside of the
         range of available block numbers (typically the last 255 blocks).
         """
-        # Create a copy of the generator so that traversing it doesn't change the
-        # original generator
-        self.prev_hashes_generator_backup, temp_prev_hashes_generator = itertools.tee(
-            self.prev_hashes_generator_backup
-        )
-
         ancestor_depth = self.block_number - block_number - 1
         is_ancestor_depth_out_of_range = (
             ancestor_depth >= MAX_PREV_HEADER_DEPTH or
@@ -216,19 +207,11 @@ def get_ancestor_hash(self, block_number: int) -> Hash32:
         if is_ancestor_depth_out_of_range:
             return Hash32(b'')
 
-        current_ancestor_depth = 0
-        current_ancestor_hash = None
-        while True:
-            try:
-                current_ancestor_hash = next(temp_prev_hashes_generator)
-            except StopIteration:
-                # Ancestor with specified depth not present
-                return Hash32(b'')
-            if current_ancestor_depth == ancestor_depth:
-                return current_ancestor_hash
-            current_ancestor_depth += 1
-
-        raise Exception("Invariant: unreachable code path")
+        try:
+            return nth(ancestor_depth, self.execution_context.prev_hashes)
+        except StopIteration:
+            # Ancestor with specified depth not present
+            return Hash32(b'')
 
     #
     # Computation

tests/core/generator-utils/test_cached_iteratable.py renamed to tests/core/generator-utils/test_cached_iterable.py

@@ -23,3 +23,10 @@ def crash_after_first_val():
     repeated_use = CachedIterable(crash_after_first_val())
     assert first(repeated_use) == 1
     assert first(repeated_use) == 1
+
+
+def test_cached_generator_iterable():
+    input_vals = [2]
+    repeated_use = CachedIterable(input_vals)
+    assert list(repeated_use) == input_vals
+    assert list(repeated_use) == input_vals