move augmented to caller

chia/_tests/util/full_sync.py

@@ -28,6 +28,7 @@
 from chia.types.full_block import FullBlock
 from chia.types.peer_info import PeerInfo
 from chia.types.validation_state import ValidationState
+from chia.util.augmented_chain import AugmentedBlockchain
 from chia.util.config import load_config
 from chia.util.ints import uint16
 
@@ -184,6 +185,7 @@ async def run_sync_test(
                 worst_batch_height = None
                 worst_batch_time_per_block = None
                 peer_info = peer.get_peer_logging()
+                blockchain = AugmentedBlockchain(full_node.blockchain)
                 async for r in rows:
                     batch_start_time = time.monotonic()
                     with enable_profiler(profile, height):
@@ -216,6 +218,7 @@ async def run_sync_test(
                                 peer_info,
                                 ForkInfo(fork_height, fork_height, header_hash),
                                 ValidationState(ssi, diff, None),
+                                blockchain,
                             )
                             end_height = block_batch[-1].height
                             full_node.blockchain.clean_block_record(end_height - full_node.constants.BLOCKS_CACHE_SIZE)

chia/_tests/wallet/sync/test_wallet_sync.py

@@ -55,6 +55,7 @@
 from chia.types.peer_info import PeerInfo
 from chia.types.spend_bundle import SpendBundle
 from chia.types.validation_state import ValidationState
+from chia.util.augmented_chain import AugmentedBlockchain
 from chia.util.hash import std_hash
 from chia.util.ints import uint32, uint64, uint128
 from chia.wallet.nft_wallet.nft_wallet import NFTWallet
@@ -444,11 +445,13 @@ async def test_long_sync_wallet(
         full_node.constants, True, block_record, full_node.blockchain
     )
     fork_height = blocks_reorg[-num_blocks - 10].height - 1
+    blockchain = AugmentedBlockchain(full_node.blockchain)
     await full_node.add_block_batch(
         blocks_reorg[-num_blocks - 10 : -1],
         PeerInfo("0.0.0.0", 0),
         ForkInfo(fork_height, fork_height, blocks_reorg[-num_blocks - 10].prev_header_hash),
         ValidationState(sub_slot_iters, difficulty, None),
+        blockchain,
     )
     await full_node.add_block(blocks_reorg[-1])
 
@@ -572,11 +575,13 @@ async def test_wallet_reorg_get_coinbase(
     sub_slot_iters, difficulty = get_next_sub_slot_iters_and_difficulty(
         full_node.constants, True, block_record, full_node.blockchain
     )
+    blockchain = AugmentedBlockchain(full_node.blockchain)
     await full_node.add_block_batch(
         blocks_reorg_2[-44:],
         PeerInfo("0.0.0.0", 0),
         ForkInfo(blocks_reorg_2[-45].height, blocks_reorg_2[-45].height, blocks_reorg_2[-45].header_hash),
         ValidationState(sub_slot_iters, difficulty, None),
+        blockchain,
     )
 
     for wallet_node, wallet_server in wallets:

chia/full_node/full_node.py

@@ -636,8 +636,9 @@ async def short_sync_batch(self, peer: WSChiaConnection, start_height: uint32, t
                         self.constants, new_slot, prev_b, self.blockchain
                     )
                     vs = ValidationState(ssi, diff, None)
+                    blockchain = AugmentedBlockchain(self.blockchain)
                     success, state_change_summary = await self.add_block_batch(
-                        response.blocks, peer_info, fork_info, vs
+                        response.blocks, peer_info, fork_info, vs, blockchain
                     )
                     if not success:
                         raise ValueError(f"Error short batch syncing, failed to validate blocks {height}-{end_height}")
@@ -1477,13 +1478,13 @@ async def add_block_batch(
         peer_info: PeerInfo,
         fork_info: ForkInfo,
         vs: ValidationState,  # in-out parameter
+        blockchain: AugmentedBlockchain,
         wp_summaries: Optional[list[SubEpochSummary]] = None,
     ) -> tuple[bool, Optional[StateChangeSummary]]:
         # Precondition: All blocks must be contiguous blocks, index i+1 must be the parent of index i
         # Returns a bool for success, as well as a StateChangeSummary if the peak was advanced
 
         pre_validate_start = time.monotonic()
-        blockchain = AugmentedBlockchain(self.blockchain)
         blocks_to_validate = await self.skip_blocks(blockchain, all_blocks, fork_info, vs)
 
         if len(blocks_to_validate) == 0:
@@ -1545,7 +1546,7 @@ async def skip_blocks(
             # we have already validated this block once, no need to do it again.
             # however, if this block is not part of the main chain, we need to
             # update the fork context with its additions and removals
-            if blockchain.height_to_hash(block.height) == header_hash:
+            if self.blockchain.height_to_hash(block.height) == header_hash:
                 # we're on the main chain, just fast-forward the fork height
                 fork_info.reset(block.height, header_hash)
             else:

chia/simulator/add_blocks_in_batches.py

@@ -8,6 +8,7 @@
 from chia.types.full_block import FullBlock
 from chia.types.peer_info import PeerInfo
 from chia.types.validation_state import ValidationState
+from chia.util.augmented_chain import AugmentedBlockchain
 from chia.util.batches import to_batches
 from chia.util.ints import uint32
 
@@ -40,8 +41,9 @@ async def add_blocks_in_batches(
         if (b.height % 128) == 0:
             print(f"main chain: {b.height:4} weight: {b.weight}")
         # vs is updated by the call to add_block_batch()
+        blockchain = AugmentedBlockchain(full_node.blockchain)
         success, state_change_summary = await full_node.add_block_batch(
-            block_batch.entries, PeerInfo("0.0.0.0", 0), fork_info, vs
+            block_batch.entries, PeerInfo("0.0.0.0", 0), fork_info, vs, blockchain
         )
         assert success is True
         if state_change_summary is not None:

tools/test_full_sync.py

@@ -21,6 +21,7 @@
 from chia.server.ws_connection import WSChiaConnection
 from chia.types.full_block import FullBlock
 from chia.types.validation_state import ValidationState
+from chia.util.augmented_chain import AugmentedBlockchain
 from chia.util.config import load_config
 
 
@@ -149,6 +150,7 @@ async def run_sync_checkpoint(
 
             block_batch = []
             peer_info = peer.get_peer_logging()
+            blockchain = AugmentedBlockchain(full_node.blockchain)
             async for r in rows:
                 block = FullBlock.from_bytes_unchecked(zstd.decompress(r[0]))
                 block_batch.append(block)
@@ -169,6 +171,7 @@ async def run_sync_checkpoint(
                     peer_info,
                     ForkInfo(fork_height, fork_height, header_hash),
                     ValidationState(ssi, diff, None),
+                    blockchain,
                 )
                 end_height = block_batch[-1].height
                 full_node.blockchain.clean_block_record(end_height - full_node.constants.BLOCKS_CACHE_SIZE)
@@ -192,6 +195,7 @@ async def run_sync_checkpoint(
                     peer_info,
                     ForkInfo(fork_height, fork_height, fork_header_hash),
                     ValidationState(ssi, diff, None),
+                    blockchain,
                 )
                 if not success:
                     raise RuntimeError("failed to ingest block batch")