CHIA-3179 Optimize rolled back state construction in _reconsider_peak (#19734)

Optimize rolled back state construction in _reconsider_peak.

Author: AmineKhaldi

chia/_tests/core/full_node/stores/test_coin_store.py

@@ -282,15 +282,15 @@ async def test_rollback(db_version: int, bt: BlockTools) -> None:
 
         # The reorg will revert the creation and spend of many coins. It will also revert the spend (but not the
         # creation) of the selected coin.
-        changed_records = await coin_store.rollback_to_block(reorg_index)
-        changed_coin_records = [cr.coin for cr in changed_records]
-        assert selected_coin in changed_records
+        coin_changes = await coin_store.rollback_to_block(reorg_index)
+        changed_coins = {cr.coin for cr in coin_changes.values()}
+        assert selected_coin.coin in changed_coins
         for coin_record in all_records:
             assert coin_record is not None
             if coin_record.confirmed_block_index > reorg_index:
-                assert coin_record.coin in changed_coin_records
+                assert coin_record.coin in changed_coins
             if coin_record.spent_block_index > reorg_index:
-                assert coin_record.coin in changed_coin_records
+                assert coin_record.coin in changed_coins
 
         for block in blocks:
             if not block.is_transaction_block():

chia/consensus/blockchain.py

@@ -510,8 +510,7 @@ async def _reconsider_peak(
                 )
 
             if block_record.prev_hash != peak.header_hash:
-                for coin_record in await self.coin_store.rollback_to_block(fork_info.fork_height):
-                    rolled_back_state[coin_record.name] = coin_record
+                rolled_back_state = await self.coin_store.rollback_to_block(fork_info.fork_height)
                 if self._log_coins and len(rolled_back_state) > 0:
                     log.info(f"rolled back {len(rolled_back_state)} coins, to fork height {fork_info.fork_height}")
                     log.info(

chia/full_node/coin_store.py

@@ -535,10 +535,10 @@ async def batch_coin_states_by_puzzle_hashes(
 
         return coin_states, next_height
 
-    async def rollback_to_block(self, block_index: int) -> list[CoinRecord]:
+    async def rollback_to_block(self, block_index: int) -> dict[bytes32, CoinRecord]:
         """
         Note that block_index can be negative, in which case everything is rolled back
-        Returns the list of coin records that have been modified
+        Returns a map of coin ID to coin record for modified items.
         """
 
         coin_changes: dict[bytes32, CoinRecord] = {}
@@ -572,7 +572,7 @@ async def rollback_to_block(self, block_index: int) -> list[CoinRecord]:
                     coin_changes[coin_name] = record
 
             await conn.execute("UPDATE coin_record SET spent_index=0 WHERE spent_index>?", (block_index,))
-        return list(coin_changes.values())
+        return coin_changes
 
     # Update coin_record to be spent in DB
     async def _set_spent(self, coin_names: list[bytes32], index: uint32) -> None: