More refactoring around withdrawals

- Check touched for empty accounts to delete after applying all withdrawals, much in the same way we check at the end of transaction computations.
- Add a test for testing withdrawal retrieval from the chaindb
- Minor refactoring and fixes

eth/abc.py

@@ -3146,11 +3146,7 @@ def get_transaction_context(self,
     def apply_withdrawal(self, withdrawal: WithdrawalAPI) -> None:
         ...
 
-    def apply_all_withdrawals(
-        self,
-        withdrawals: Sequence[WithdrawalAPI],
-        base_header: BlockHeaderAPI,
-    ) -> None:
+    def apply_all_withdrawals(self, withdrawals: Sequence[WithdrawalAPI]) -> None:
         ...
 
 
@@ -3365,6 +3361,15 @@ def apply_all_transactions(
         """
         ...
 
+    def apply_all_withdrawals(self, withdrawals: Sequence[WithdrawalAPI]) -> None:
+        """
+        Updates the state by applying all withdrawals.
+        This does *not* update the current block or header of the VM.
+
+        :param withdrawals: an iterable of all withdrawals to apply
+        """
+        ...
+
     @abstractmethod
     def make_receipt(self,
                      base_header: BlockHeaderAPI,
@@ -3401,13 +3406,16 @@ def mine_block(self, block: BlockAPI, *args: Any, **kwargs: Any) -> BlockAndMeta
         ...
 
     @abstractmethod
-    def set_block_transactions(self,
-                               base_block: BlockAPI,
-                               new_header: BlockHeaderAPI,
-                               transactions: Sequence[SignedTransactionAPI],
-                               receipts: Sequence[ReceiptAPI]) -> BlockAPI:
+    def set_block_transactions_and_withdrawals(
+        self,
+        base_block: BlockAPI,
+        new_header: BlockHeaderAPI,
+        transactions: Sequence[SignedTransactionAPI],
+        receipts: Sequence[ReceiptAPI],
+        withdrawals: Sequence[WithdrawalAPI] = None,
+    ) -> BlockAPI:
         """
-        Create a new block with the given ``transactions``.
+        Create a new block with the given ``transactions`` and/or ``withdrawals``.
         """
         ...
 
@@ -3995,18 +4003,20 @@ def get_canonical_block_hash(self, block_number: BlockNumber) -> Hash32:
         ...
 
     @abstractmethod
-    def build_block_with_transactions(
+    def build_block_with_transactions_and_withdrawals(
             self,
             transactions: Tuple[SignedTransactionAPI, ...],
-            parent_header: BlockHeaderAPI = None
+            parent_header: BlockHeaderAPI = None,
+            withdrawals: Tuple[WithdrawalAPI, ...] = None,
     ) -> Tuple[BlockAPI, Tuple[ReceiptAPI, ...], Tuple[ComputationAPI, ...]]:
         """
         Generate a block with the provided transactions. This does *not* import
         that block into your chain. If you want this new block in your chain,
         run :meth:`~import_block` with the result block from this method.
 
-        :param transactions: an iterable of transactions to insert to the block
+        :param transactions: an iterable of transactions to insert int the block
         :param parent_header: parent of the new block -- or canonical head if ``None``
+        :param withdrawals: an iterable of withdrawals to insert intto the block
         :return: (new block, receipts, computations)
         """
         ...

eth/chains/base.py

@@ -50,7 +50,7 @@
     ComputationAPI,
     StateAPI,
     SignedTransactionAPI,
-    UnsignedTransactionAPI,
+    UnsignedTransactionAPI, WithdrawalAPI,
 )
 from eth.consensus import (
     ConsensusContext,
@@ -347,16 +347,30 @@ def get_canonical_block_by_number(self, block_number: BlockNumber) -> BlockAPI:
     def get_canonical_block_hash(self, block_number: BlockNumber) -> Hash32:
         return self.chaindb.get_canonical_block_hash(block_number)
 
-    def build_block_with_transactions(
+    def build_block_with_transactions_and_withdrawals(
             self,
             transactions: Sequence[SignedTransactionAPI],
-            parent_header: BlockHeaderAPI = None
+            parent_header: BlockHeaderAPI = None,
+            withdrawals: Sequence[WithdrawalAPI] = None,
     ) -> Tuple[BlockAPI, Tuple[ReceiptAPI, ...], Tuple[ComputationAPI, ...]]:
         base_header = self.ensure_header(parent_header)
         vm = self.get_vm(base_header)
 
-        new_header, receipts, computations = vm.apply_all_transactions(transactions, base_header)
-        new_block = vm.set_block_transactions(vm.get_block(), new_header, transactions, receipts)
+        new_header, receipts, computations = vm.apply_all_transactions(
+            transactions,
+            base_header,
+        )
+
+        if withdrawals:
+            vm.apply_all_withdrawals(withdrawals)
+
+        new_block = vm.set_block_transactions_and_withdrawals(
+            vm.get_block(),
+            new_header,
+            transactions,
+            receipts,
+            withdrawals=withdrawals,
+        )
 
         return new_block, receipts, computations
 
@@ -628,9 +642,10 @@ def __init__(self, base_db: AtomicDatabaseAPI, header: BlockHeaderAPI = None) ->
         super().__init__(base_db)
         self.header = self.ensure_header(header)
 
-    def apply_transaction(self,
-                          transaction: SignedTransactionAPI
-                          ) -> Tuple[BlockAPI, ReceiptAPI, ComputationAPI]:
+    def apply_transaction(
+        self,
+        transaction: SignedTransactionAPI,
+    ) -> Tuple[BlockAPI, ReceiptAPI, ComputationAPI]:
         vm = self.get_vm(self.header)
         base_block = vm.get_block()
 
@@ -644,7 +659,12 @@ def apply_transaction(self,
         transactions = base_block.transactions + (transaction, )
         receipts = base_block.get_receipts(self.chaindb) + (receipt, )
 
-        new_block = vm.set_block_transactions(base_block, new_header, transactions, receipts)
+        new_block = vm.set_block_transactions_and_withdrawals(
+            base_block,
+            new_header,
+            transactions,
+            receipts,
+        )
 
         self.header = new_block.header
 
@@ -674,6 +694,7 @@ def mine_all(
             transactions: Sequence[SignedTransactionAPI],
             *args: Any,
             parent_header: BlockHeaderAPI = None,
+            withdrawals: Sequence[WithdrawalAPI] = None,
             **kwargs: Any,
     ) -> Tuple[BlockImportResult, Tuple[ReceiptAPI, ...], Tuple[ComputationAPI, ...]]:
 
@@ -686,13 +707,26 @@ def mine_all(
         vm = self.get_vm(custom_header)
 
         new_header, receipts, computations = vm.apply_all_transactions(transactions, base_header)
-        filled_block = vm.set_block_transactions(vm.get_block(), new_header, transactions, receipts)
+
+        if withdrawals:
+            vm.apply_all_withdrawals(withdrawals)
+
+        filled_block = vm.set_block_transactions_and_withdrawals(
+            vm.get_block(),
+            new_header,
+            transactions,
+            receipts,
+            withdrawals=withdrawals,
+        )
 
         block_result = vm.mine_block(filled_block, *args, **kwargs)
         imported_block = block_result.block
 
         block_persist_result = self.persist_block(imported_block)
-        block_import_result = BlockImportResult(*block_persist_result, block_result.meta_witness)
+        block_import_result = BlockImportResult(
+            *block_persist_result,
+            block_result.meta_witness,
+        )
 
         self.header = self.create_header_from_parent(imported_block.header)
         return (block_import_result, receipts, computations)

eth/exceptions.py

@@ -65,13 +65,6 @@ class ReceiptNotFound(PyEVMError):
     pass
 
 
-class WithdrawalNotFound(PyEVMError):
-    """
-    Raised when the withdrawal with the given hash or block index does not exist.
-    """
-    pass
-
-
 class ParentNotFound(HeaderNotFound):
     """
     Raised when the parent of a given block does not exist.

eth/typing.py

@@ -26,7 +26,12 @@
 )
 
 if TYPE_CHECKING:
-    from eth.abc import VirtualMachineAPI  # noqa: F401
+    from eth.abc import (  # noqa: F401
+        BlockHeaderAPI,
+        SignedTransactionAPI,
+        VirtualMachineAPI,
+        WithdrawalAPI,
+    )
 
 
 JournalDBCheckpoint = NewType('JournalDBCheckpoint', int)
@@ -38,9 +43,16 @@
                             'storage': Dict[int, int]
                             })
 AccountState = Dict[Address, AccountDetails]
-
 AccountDiff = Iterable[Tuple[Address, str, Union[int, bytes], Union[int, bytes]]]
 
+
+class Block(TypedDict, total=False):
+    header: "BlockHeaderAPI"
+    transactions: Sequence["SignedTransactionAPI"]
+    uncles: Sequence["BlockHeaderAPI"]
+    withdrawals: Sequence["WithdrawalAPI"]
+
+
 BlockRange = Tuple[BlockNumber, BlockNumber]
 
 ChainGaps = Tuple[

eth/vm/base.py

@@ -2,6 +2,7 @@
 import itertools
 import logging
 from typing import (
+    List, TYPE_CHECKING,
     Any,
     ClassVar,
     Iterable,
@@ -85,6 +86,11 @@
     Message,
 )
 
+if TYPE_CHECKING:
+    from eth.typing import (  # noqa: F401
+        Block,
+    )
+
 
 class VM(Configurable, VirtualMachineAPI):
     block_class: Type[BlockAPI] = None
@@ -307,20 +313,24 @@ def apply_withdrawal(
     ) -> None:
         self.state.apply_withdrawal(withdrawal)
 
-    def apply_all_withdrawals(
-        self,
-        withdrawals: Sequence[WithdrawalAPI],
-        base_header: BlockHeaderAPI,
-    ) -> None:
-        vm_header = self.get_header()
-        self._validate_header_before_apply(base_header, vm_header)
+    def apply_all_withdrawals(self, withdrawals: Sequence[WithdrawalAPI]) -> None:
+        touched_addresses: List[Address] = []
 
         for withdrawal in withdrawals:
             # validate withdrawal fields
             withdrawal.validate()
 
             self.apply_withdrawal(withdrawal)
 
+            # collect all touched addresses
+            if withdrawal.address not in touched_addresses:
+                touched_addresses.append(withdrawal.address)
+
+        for address in touched_addresses:
+            # if account is empty after applying all withdrawals, delete it
+            if self.state.account_is_empty(address):
+                self.state.delete_account(address)
+
     #
     # Importing blocks
     #
@@ -367,30 +377,21 @@ def import_block(self, block: BlockAPI) -> BlockAndMetaWitness:
         # run all of the transactions.
         new_header, receipts, _ = self.apply_all_transactions(block.transactions, header)
 
-        block_with_transactions = self.set_block_transactions(
+        withdrawals = block.withdrawals if hasattr(block, "withdrawals") else None
+
+        if withdrawals:
+            # post-shanghai blocks
+            self.apply_all_withdrawals(block.withdrawals)
+
+        filled_block = self.set_block_transactions_and_withdrawals(
             self.get_block(),
             new_header,
             block.transactions,
             receipts,
+            withdrawals=withdrawals,
         )
 
-        if hasattr(block, "withdrawals") and block.withdrawals not in (None, ()):
-            self.apply_all_withdrawals(block_with_transactions.withdrawals, header)
-
-            withdrawals_root_hash, withdrawal_kv_nodes = make_trie_root_and_nodes(
-                block.withdrawals
-            )
-            self.chaindb.persist_trie_data_dict(withdrawal_kv_nodes)
-
-            block_with_withdrawals = block_with_transactions.copy(
-                withdrawals=block.withdrawals,
-                header=block_with_transactions.header.copy(
-                    withdrawals_root=withdrawals_root_hash,
-                ),
-            )
-            return self.mine_block(block_with_withdrawals)
-
-        return self.mine_block(block_with_transactions)
+        return self.mine_block(filled_block)
 
     def mine_block(self, block: BlockAPI, *args: Any, **kwargs: Any) -> BlockAndMetaWitness:
         packed_block = self.pack_block(block, *args, **kwargs)
@@ -401,25 +402,39 @@ def mine_block(self, block: BlockAPI, *args: Any, **kwargs: Any) -> BlockAndMeta
 
         return block_result
 
-    def set_block_transactions(self,
-                               base_block: BlockAPI,
-                               new_header: BlockHeaderAPI,
-                               transactions: Sequence[SignedTransactionAPI],
-                               receipts: Sequence[ReceiptAPI]) -> BlockAPI:
+    def set_block_transactions_and_withdrawals(
+        self,
+        base_block: BlockAPI,
+        new_header: BlockHeaderAPI,
+        transactions: Sequence[SignedTransactionAPI],
+        receipts: Sequence[ReceiptAPI],
+        withdrawals: Sequence[WithdrawalAPI] = None,
+    ) -> BlockAPI:
 
         tx_root_hash, tx_kv_nodes = make_trie_root_and_nodes(transactions)
         self.chaindb.persist_trie_data_dict(tx_kv_nodes)
 
         receipt_root_hash, receipt_kv_nodes = make_trie_root_and_nodes(receipts)
         self.chaindb.persist_trie_data_dict(receipt_kv_nodes)
 
-        return base_block.copy(
-            transactions=transactions,
-            header=new_header.copy(
-                transaction_root=tx_root_hash,
-                receipt_root=receipt_root_hash,
-            ),
-        )
+        block_fields: "Block" = {"transactions": transactions}
+        block_header_fields = {
+            "transaction_root": tx_root_hash,
+            "receipt_root": receipt_root_hash,
+        }
+
+        if withdrawals:
+            withdrawals_root_hash, withdrawal_kv_nodes = make_trie_root_and_nodes(
+                withdrawals,
+            )
+            self.chaindb.persist_trie_data_dict(withdrawal_kv_nodes)
+
+            block_fields["withdrawals"] = withdrawals
+            block_header_fields["withdrawals_root"] = withdrawals_root_hash
+
+        block_fields["header"] = base_block.header.copy(**block_header_fields)
+
+        return base_block.copy(**block_fields)
 
     #
     # Finalization

eth/vm/forks/shanghai/state.py

@@ -21,6 +21,3 @@ def apply_withdrawal(self, withdrawal: WithdrawalAPI) -> None:
         # withdrawal amount is in gwei, convert to wei
         amount_in_wei = withdrawal.amount * 10 ** 9
         self.delta_balance(withdrawal.address, amount_in_wei)
-
-        if self.account_is_empty(withdrawal.address):
-            self.delete_account(withdrawal.address)