Repository: Chia-Network/chia-blockchain
File: chia/consensus/augmented_chain.py (Python)

Code before:
        _protocol_check: ClassVar[BlocksProtocol] = cast("AugmentedBlockchain", None)

    _underlying: BlocksProtocol
    _extra_blocks: dict[bytes32, tuple[FullBlock, BlockRecord]]
    _height_to_hash: dict[uint32, bytes32]

    def __init__(self, underlying: BlocksProtocol) -> None:
        self._underlying = underlying
        self._extra_blocks = {}
        self._height_to_hash = {}

    def _get_block_record(self, header_hash: bytes32) -> BlockRecord | None:
        eb = self._extra_blocks.get(header_hash)
        if eb is None:
            return None
        return eb[1]

    def add_extra_block(self, block: FullBlock, block_record: BlockRecord) -> None:
        if block.header_hash != block_record.header_hash:
            raise AugmentedBlockchainValidationError(
                f"Block header hash mismatch: block={block.header_hash.hex()[:16]}, "
                f"record={block_record.header_hash.hex()[:16]}"
            )
        if self._height_to_hash:
            max_height = max(self._height_to_hash.keys())
            last_header_hash = self._height_to_hash[max_height]

            # New block must chain to the last added block
            if block_record.prev_hash != last_header_hash:
                raise AugmentedBlockchainValidationError(
                    f"New block's prev_hash must match last added block. "
                    f"Expected {last_header_hash.hex()[:16]}, got {block_record.prev_hash.hex()[:16]}"
                )
        elif block_record.height > 0:
            if self._underlying.try_block_record(block_record.prev_hash) is None:
                raise AugmentedBlockchainValidationError(
                    f"First added block's prev_hash must exist in underlying blockchain. "
                    f"Block height {block_record.height}, prev_hash {block_record.prev_hash.hex()[:16]} not found"
                )

        self._extra_blocks[block_record.header_hash] = (block, block_record)
        self._height_to_hash[block_record.height] = block_record.header_hash

    def remove_extra_block(self, hh: bytes32) -> None:
        if hh not in self._extra_blocks:
            return

        block_record = self._extra_blocks.pop(hh)[1]
        if self._underlying.contains_block(block_record.header_hash, block_record.height):
            height_to_remove = block_record.height
            for h in range(height_to_remove, -1, -1):

Review comment (suggestion):
I don't think this is quite right. `_height_to_hash` will be empty if the underlying blockchain hasn't been augmented, right?

You may have to ask `self._blockchain` for the peak hash in that case.

I suspect if you fix this, you may not need to make the other changes to the test.

Code after:
        _protocol_check: ClassVar[BlocksProtocol] = cast("AugmentedBlockchain", None)

    _underlying: BlocksProtocol
    _extra_blocks: dict[bytes32, tuple[FullBlock, BlockRecord]]
    _height_to_hash: dict[uint32, bytes32]

    def __init__(self, underlying: BlocksProtocol) -> None:
        self._underlying = underlying
        self._extra_blocks = {}
        self._height_to_hash = {}

    def _get_block_record(self, header_hash: bytes32) -> BlockRecord | None:
        eb = self._extra_blocks.get(header_hash)
        if eb is None:
            return None
        return eb[1]

    def add_extra_block(self, block: FullBlock, block_record: BlockRecord) -> None:
        if block.header_hash != block_record.header_hash:
            raise AugmentedBlockchainValidationError(
                f"Block header hash mismatch: block={block.header_hash.hex()[:16]}, "
                f"record={block_record.header_hash.hex()[:16]}"
            )
        if self._height_to_hash:
            max_height = max(self._height_to_hash.keys())
            last_header_hash = self._height_to_hash[max_height]

            if block_record.prev_hash != last_header_hash:
                raise AugmentedBlockchainValidationError(
                    f"New block's prev_hash must match last added block. "
                    f"Expected {last_header_hash.hex()[:16]}, got {block_record.prev_hash.hex()[:16]}"
                )
        elif block_record.height > 0:
            if self._underlying.try_block_record(block_record.prev_hash) is None:
                raise AugmentedBlockchainValidationError(
                    f"First added block's prev_hash must exist in underlying blockchain. "
                    f"Block height {block_record.height}, prev_hash {block_record.prev_hash.hex()[:16]} not found"
                )

        self._extra_blocks[block_record.header_hash] = (block, block_record)
        self._height_to_hash[block_record.height] = block_record.header_hash

    def remove_extra_block(self, hh: bytes32) -> None:
        if hh not in self._extra_blocks:
            return

        block_record = self._extra_blocks.pop(hh)[1]
        if self._underlying.contains_block(block_record.header_hash, block_record.height):
            height_to_remove = block_record.height
            for h in range(height_to_remove, -1, -1):
                if h not in self._height_to_hash: