remove contains height from BlocksProtocol

chia/_tests/blockchain/test_augmented_chain.py

@@ -22,6 +22,7 @@ class NullBlockchain:
 
     added_blocks: set[bytes32] = field(default_factory=set)
     heights: dict[uint32, bytes32] = field(default_factory=dict)
+    _peak_height: Optional[uint32] = None
 
     # BlocksProtocol
     async def lookup_block_generators(self, header_hash: bytes32, generator_refs: set[uint32]) -> dict[uint32, bytes]:
@@ -52,6 +53,9 @@ def contains_block(self, header_hash: bytes32) -> bool:
     def contains_height(self, height: uint32) -> bool:
         return height in self.heights.keys()
 
+    def get_peak_height(self) -> Optional[uint32]:
+        return self._peak_height
+
     async def prev_block_hash(self, header_hashes: list[bytes32]) -> list[bytes32]:
         raise KeyError("no block records in NullBlockchain")  # pragma: no cover
 

chia/_tests/util/blockchain_mock.py

@@ -38,6 +38,10 @@ def __init__(
         self._sub_epoch_summaries = sub_epoch_summaries
         self._sub_epoch_segments: dict[bytes32, SubEpochSegments] = {}
         self.log = logging.getLogger(__name__)
+        self._peak_height = None
+        for block in blocks.values():
+            if self._peak_height is None or block.height > self._peak_height:
+                self._peak_height = block.height
 
     def get_peak(self) -> Optional[BlockRecord]:
         return None
@@ -73,7 +77,10 @@ async def contains_block_from_db(self, header_hash: bytes32) -> bool:
         return header_hash in self._block_records
 
     def contains_height(self, height: uint32) -> bool:
-        return height in self._height_to_hash
+        peak_height = self.get_peak_height()
+        if peak_height is None:
+            return False
+        return height <= peak_height
 
     async def warmup(self, fork_point: uint32) -> None:
         return

chia/consensus/blockchain.py

@@ -754,7 +754,10 @@ def height_to_hash(self, height: uint32) -> Optional[bytes32]:
         return self.__height_map.get_hash(height)
 
     def contains_height(self, height: uint32) -> bool:
-        return self.__height_map.contains_height(height)
+        peak_height = self.get_peak_height()
+        if peak_height is None:
+            return False
+        return height <= peak_height
 
     def get_peak_height(self) -> Optional[uint32]:
         return self._peak_height

chia/consensus/blockchain_interface.py

@@ -14,6 +14,7 @@ class BlockRecordsProtocol(Protocol):
     def try_block_record(self, header_hash: bytes32) -> Optional[BlockRecord]: ...
     def block_record(self, header_hash: bytes32) -> BlockRecord: ...
     def contains_height(self, height: uint32) -> bool: ...
+    def get_peak_height(self) -> Optional[uint32]: ...
     def contains_block(self, header_hash: bytes32) -> bool: ...
     def height_to_hash(self, height: uint32) -> Optional[bytes32]: ...
     def height_to_block_record(self, height: uint32) -> BlockRecord: ...

chia/util/augmented_chain.py

@@ -28,11 +28,13 @@ class AugmentedBlockchain:
     _underlying: BlocksProtocol
     _extra_blocks: dict[bytes32, tuple[FullBlock, BlockRecord]]
     _height_to_hash: dict[uint32, bytes32]
+    _peak_height: Optional[uint32]
 
     def __init__(self, underlying: BlocksProtocol) -> None:
         self._underlying = underlying
         self._extra_blocks = {}
         self._height_to_hash = {}
+        self._peak_height = None
 
     def _get_block_record(self, header_hash: bytes32) -> Optional[BlockRecord]:
         eb = self._extra_blocks.get(header_hash)
@@ -44,9 +46,14 @@ def add_extra_block(self, block: FullBlock, block_record: BlockRecord) -> None:
         assert block.header_hash == block_record.header_hash
         self._extra_blocks[block_record.header_hash] = (block, block_record)
         self._height_to_hash[block_record.height] = block_record.header_hash
+        if self._peak_height is None:
+            self._peak_height = block.height
+        else:
+            self._peak_height = max(block.height, self._peak_height)
 
     def remove_extra_block(self, hh: bytes32) -> None:
         if hh in self._extra_blocks:
+            assert self._underlying.contains_block(hh)
             block_record = self._extra_blocks.pop(hh)[1]
             del self._height_to_hash[block_record.height]
 
@@ -121,7 +128,10 @@ def contains_block(self, header_hash: bytes32) -> bool:
         return (header_hash in self._extra_blocks) or self._underlying.contains_block(header_hash)
 
     def contains_height(self, height: uint32) -> bool:
-        return (height in self._height_to_hash) or self._underlying.contains_height(height)
+        peak_height = self.get_peak_height()
+        if peak_height is None:
+            return False
+        return height <= peak_height
 
     async def prev_block_hash(self, header_hashes: list[bytes32]) -> list[bytes32]:
         ret: list[bytes32] = []
@@ -132,3 +142,6 @@ async def prev_block_hash(self, header_hashes: list[bytes32]) -> list[bytes32]:
             else:
                 ret.extend(await self._underlying.prev_block_hash([hh]))
         return ret
+
+    def get_peak_height(self) -> Optional[uint32]:
+        return self._peak_height

chia/util/block_cache.py

@@ -16,18 +16,26 @@ class BlockCache:
 
     _block_records: dict[bytes32, BlockRecord]
     _height_to_hash: dict[uint32, bytes32]
+    _peak_height: Optional[uint32]
 
     def __init__(
         self,
         blocks: dict[bytes32, BlockRecord],
     ):
         self._block_records = blocks
-        self._height_to_hash = {block.height: hh for hh, block in blocks.items()}
+        self._height_to_hash = {}
+        self._peak_height = uint32(0)
+        for hh, block in blocks.items():
+            self._height_to_hash[block.height] = hh
+            if self._peak_height is None or block.height > self._peak_height:
+                self._peak_height = block.height
 
     def add_block(self, block: BlockRecord) -> None:
         hh = block.header_hash
         self._block_records[hh] = block
         self._height_to_hash[block.height] = hh
+        if self._peak_height is None or block.height > self._peak_height:
+            self._peak_height = block.height
 
     def block_record(self, header_hash: bytes32) -> BlockRecord:
         return self._block_records[header_hash]
@@ -47,7 +55,13 @@ def contains_block(self, header_hash: bytes32) -> bool:
         return header_hash in self._block_records
 
     def contains_height(self, height: uint32) -> bool:
-        return height in self._height_to_hash
+        peak_height = self.get_peak_height()
+        if peak_height is None:
+            return False
+        return height <= peak_height
+
+    def get_peak_height(self) -> Optional[uint32]:
+        return self._peak_height
 
     def try_block_record(self, header_hash: bytes32) -> Optional[BlockRecord]:
         return self._block_records.get(header_hash)

chia/wallet/wallet_blockchain.py

@@ -157,7 +157,6 @@ async def _rollback_to_height(self, height: int) -> None:
             return
         for h in range(max(0, height + 1), self._peak.height + 1):
             del self._height_to_hash[uint32(h)]
-
         await self._basic_store.remove_object("PEAK_BLOCK")
 
     async def set_peak_block(self, block: HeaderBlock, timestamp: Optional[uint64] = None) -> None:
@@ -198,6 +197,11 @@ def contains_block(self, header_hash: bytes32) -> bool:
     def contains_height(self, height: uint32) -> bool:
         return height in self._height_to_hash
 
+    def get_peak_height(self) -> Optional[uint32]:
+        if self._peak is None:
+            return None
+        return self._peak.height
+
     def height_to_hash(self, height: uint32) -> bytes32:
         return self._height_to_hash[height]