From e46b2bc69c25983c09bea434013110342327a8c6 Mon Sep 17 00:00:00 2001
From: almogdepaz <almogdepaz@gmail.com>
Date: Mon, 9 Dec 2024 11:06:45 +0200
Subject: [PATCH 01/22] add assertions

---
 chia/consensus/blockchain.py | 9 ++++++++-
 1 file changed, 8 insertions(+), 1 deletion(-)

diff --git a/chia/consensus/blockchain.py b/chia/consensus/blockchain.py
index f989bfe96344..28df0e32b0d9 100644
--- a/chia/consensus/blockchain.py
+++ b/chia/consensus/blockchain.py
@@ -774,7 +774,14 @@ def contains_block(self, header_hash: bytes32) -> bool:
         True if we have already added this block to the chain. This may return false for orphan blocks
         that we have added but no longer keep in memory.
         """
-        return header_hash in self.__block_records
+        # check block records
+        block_rec = self.__block_records[header_hash]
+        if block_rec is None:
+            return False
+        block_rec = self.height_to_block_record(block_rec.height)
+        assert block_rec is not None
+        assert block_rec.header_hash == header_hash
+        return True
 
     def block_record(self, header_hash: bytes32) -> BlockRecord:
         return self.__block_records[header_hash]

From e06555f2c9a9f65dff62ddfd4dd75e9501665917 Mon Sep 17 00:00:00 2001
From: almogdepaz <almogdepaz@gmail.com>
Date: Mon, 9 Dec 2024 11:28:37 +0200
Subject: [PATCH 02/22] use_height_if_available

---
 chia/consensus/blockchain.py | 17 +++++++++--------
 1 file changed, 9 insertions(+), 8 deletions(-)

diff --git a/chia/consensus/blockchain.py b/chia/consensus/blockchain.py
index 28df0e32b0d9..b86f82ae24bc 100644
--- a/chia/consensus/blockchain.py
+++ b/chia/consensus/blockchain.py
@@ -769,18 +769,19 @@ async def validate_unfinished_block(
 
         return PreValidationResult(None, required_iters, conds, uint32(0))
 
-    def contains_block(self, header_hash: bytes32) -> bool:
+    def contains_block(self, header_hash: bytes32, height: Optional[uint32] = None) -> bool:
         """
         True if we have already added this block to the chain. This may return false for orphan blocks
         that we have added but no longer keep in memory.
         """
-        # check block records
-        block_rec = self.__block_records[header_hash]
-        if block_rec is None:
-            return False
-        block_rec = self.height_to_block_record(block_rec.height)
-        assert block_rec is not None
-        assert block_rec.header_hash == header_hash
+        if height is None:
+            block_rec = self.__block_records[header_hash]
+            if block_rec is None:
+                return False
+            height = block_rec.height
+        block_hash_from_hh = self.height_to_hash(height)
+        assert block_hash_from_hh is not None
+        assert block_hash_from_hh == header_hash
         return True
 
     def block_record(self, header_hash: bytes32) -> BlockRecord:

From 0fcbac1d6be0d00bd7155e4bec6725568cb9109e Mon Sep 17 00:00:00 2001
From: almogdepaz <almogdepaz@gmail.com>
Date: Mon, 9 Dec 2024 12:18:52 +0200
Subject: [PATCH 03/22] use try_block_record

---
 chia/_tests/blockchain/test_augmented_chain.py |  2 +-
 chia/_tests/util/blockchain_mock.py            | 16 ++++++++++++++--
 chia/consensus/blockchain.py                   |  2 +-
 chia/consensus/blockchain_interface.py         |  2 +-
 chia/full_node/full_node.py                    | 13 ++++++++-----
 chia/util/augmented_chain.py                   | 12 ++++++++++--
 chia/util/block_cache.py                       | 16 ++++++++++++++--
 chia/wallet/wallet_blockchain.py               | 16 ++++++++++++++--
 8 files changed, 63 insertions(+), 16 deletions(-)

diff --git a/chia/_tests/blockchain/test_augmented_chain.py b/chia/_tests/blockchain/test_augmented_chain.py
index 14437b9c56ce..1dcf51ebe8ec 100644
--- a/chia/_tests/blockchain/test_augmented_chain.py
+++ b/chia/_tests/blockchain/test_augmented_chain.py
@@ -46,7 +46,7 @@ def height_to_block_record(self, height: uint32) -> BlockRecord:
     def height_to_hash(self, height: uint32) -> Optional[bytes32]:
         return self.heights.get(height)
 
-    def contains_block(self, header_hash: bytes32) -> bool:
+    def contains_block(self, header_hash: bytes32, height: Optional[uint32] = None) -> bool:
         return False  # pragma: no cover
 
     def contains_height(self, height: uint32) -> bool:
diff --git a/chia/_tests/util/blockchain_mock.py b/chia/_tests/util/blockchain_mock.py
index cad13a5eeba5..11fde9585c9d 100644
--- a/chia/_tests/util/blockchain_mock.py
+++ b/chia/_tests/util/blockchain_mock.py
@@ -66,8 +66,20 @@ def height_to_hash(self, height: uint32) -> Optional[bytes32]:
         assert height in self._height_to_hash
         return self._height_to_hash[height]
 
-    def contains_block(self, header_hash: bytes32) -> bool:
-        return header_hash in self._block_records
+    def contains_block(self, header_hash: bytes32, height: Optional[uint32] = None) -> bool:
+        """
+        True if we have already added this block to the chain. This may return false for orphan blocks
+        that we have added but no longer keep in memory.
+        """
+        if height is None:
+            block_rec = self.try_block_record(header_hash)
+            if block_rec is None:
+                return False
+            height = block_rec.height
+        block_hash_from_hh = self.height_to_hash(height)
+        assert block_hash_from_hh is not None
+        assert block_hash_from_hh == header_hash
+        return True
 
     async def contains_block_from_db(self, header_hash: bytes32) -> bool:
         return header_hash in self._block_records
diff --git a/chia/consensus/blockchain.py b/chia/consensus/blockchain.py
index b86f82ae24bc..eff03500542f 100644
--- a/chia/consensus/blockchain.py
+++ b/chia/consensus/blockchain.py
@@ -775,7 +775,7 @@ def contains_block(self, header_hash: bytes32, height: Optional[uint32] = None)
         that we have added but no longer keep in memory.
         """
         if height is None:
-            block_rec = self.__block_records[header_hash]
+            block_rec = self.try_block_record(header_hash)
             if block_rec is None:
                 return False
             height = block_rec.height
diff --git a/chia/consensus/blockchain_interface.py b/chia/consensus/blockchain_interface.py
index 66a2e9dbd222..e3550f6966ca 100644
--- a/chia/consensus/blockchain_interface.py
+++ b/chia/consensus/blockchain_interface.py
@@ -14,7 +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 contains_block(self, header_hash: bytes32) -> bool: ...
+    def contains_block(self, header_hash: bytes32, height: Optional[uint32] = None) -> bool: ...
     def height_to_hash(self, height: uint32) -> Optional[bytes32]: ...
     def height_to_block_record(self, height: uint32) -> BlockRecord: ...
 
diff --git a/chia/full_node/full_node.py b/chia/full_node/full_node.py
index e644ecb8b4c7..def377fc5ef5 100644
--- a/chia/full_node/full_node.py
+++ b/chia/full_node/full_node.py
@@ -595,7 +595,7 @@ async def short_sync_batch(self, peer: WSChiaConnection, start_height: uint32, t
                 self.sync_store.batch_syncing.remove(peer.peer_node_id)
                 self.log.error(f"Error short batch syncing, could not fetch block at height {start_height}")
                 return False
-            if not self.blockchain.contains_block(first.block.prev_header_hash):
+            if not self.blockchain.contains_block(first.block.prev_header_hash, first.block.height - 1):
                 self.log.info("Batch syncing stopped, this is a deep chain")
                 self.sync_store.batch_syncing.remove(peer.peer_node_id)
                 # First sb not connected to our blockchain, do a long sync instead
@@ -699,7 +699,10 @@ async def short_sync_backtrack(
                         f"Failed to fetch block {curr_height} from {peer.get_peer_logging()}, wrong type {type(curr)}"
                     )
                 blocks.append(curr.block)
-                if self.blockchain.contains_block(curr.block.prev_header_hash) or curr_height == 0:
+                if (
+                    self.blockchain.contains_block(curr.block.prev_header_hash, curr.block.height - 1)
+                    or curr_height == 0
+                ):
                     found_fork_point = True
                     break
                 curr_height -= 1
@@ -748,7 +751,7 @@ async def new_peak(self, request: full_node_protocol.NewPeak, peer: WSChiaConnec
         # Store this peak/peer combination in case we want to sync to it, and to keep track of peers
         self.sync_store.peer_has_block(request.header_hash, peer.peer_node_id, request.weight, request.height, True)
 
-        if self.blockchain.contains_block(request.header_hash):
+        if self.blockchain.contains_block(request.header_hash, request.height):
             return None
 
         # Not interested in less heavy peaks
@@ -2006,7 +2009,7 @@ async def add_block(
 
         # Adds the block to seen, and check if it's seen before (which means header is in memory)
         header_hash = block.header_hash
-        if self.blockchain.contains_block(header_hash):
+        if self.blockchain.contains_block(header_hash, block.height):
             return None
 
         pre_validation_result: Optional[PreValidationResult] = None
@@ -2078,7 +2081,7 @@ async def add_block(
             enable_profiler(self.profile_block_validation) as pr,
         ):
             # After acquiring the lock, check again, because another asyncio thread might have added it
-            if self.blockchain.contains_block(header_hash):
+            if self.blockchain.contains_block(header_hash, block.height):
                 return None
             validation_start = time.monotonic()
             # Tries to add the block to the blockchain, if we already validated transactions, don't do it again
diff --git a/chia/util/augmented_chain.py b/chia/util/augmented_chain.py
index 75478a2e1656..5b7489486e49 100644
--- a/chia/util/augmented_chain.py
+++ b/chia/util/augmented_chain.py
@@ -117,8 +117,16 @@ def height_to_hash(self, height: uint32) -> Optional[bytes32]:
             return ret
         return self._underlying.height_to_hash(height)
 
-    def contains_block(self, header_hash: bytes32) -> bool:
-        return (header_hash in self._extra_blocks) or self._underlying.contains_block(header_hash)
+    def contains_block(self, header_hash: bytes32, height: Optional[uint32] = None) -> bool:
+        if height is None:
+            block_rec = self.try_block_record(header_hash)
+            if block_rec is None:
+                return False
+            height = block_rec.height
+        block_hash_from_hh = self.height_to_hash(height)
+        assert block_hash_from_hh is not None
+        assert block_hash_from_hh == header_hash
+        return True
 
     def contains_height(self, height: uint32) -> bool:
         return (height in self._height_to_hash) or self._underlying.contains_height(height)
diff --git a/chia/util/block_cache.py b/chia/util/block_cache.py
index ed0f7845d3a1..0a705f3ac2c8 100644
--- a/chia/util/block_cache.py
+++ b/chia/util/block_cache.py
@@ -43,8 +43,20 @@ def height_to_hash(self, height: uint32) -> Optional[bytes32]:
             return None
         return self._height_to_hash[height]
 
-    def contains_block(self, header_hash: bytes32) -> bool:
-        return header_hash in self._block_records
+    def contains_block(self, header_hash: bytes32, height: Optional[uint32] = None) -> bool:
+        """
+        True if we have already added this block to the chain. This may return false for orphan blocks
+        that we have added but no longer keep in memory.
+        """
+        if height is None:
+            block_rec = self.try_block_record(header_hash)
+            if block_rec is None:
+                return False
+            height = block_rec.height
+        block_hash_from_hh = self.height_to_hash(height)
+        assert block_hash_from_hh is not None
+        assert block_hash_from_hh == header_hash
+        return True
 
     def contains_height(self, height: uint32) -> bool:
         return height in self._height_to_hash
diff --git a/chia/wallet/wallet_blockchain.py b/chia/wallet/wallet_blockchain.py
index 882b3c1a5583..548657deb50c 100644
--- a/chia/wallet/wallet_blockchain.py
+++ b/chia/wallet/wallet_blockchain.py
@@ -192,8 +192,20 @@ async def get_finished_sync_up_to(self) -> uint32:
     def get_latest_timestamp(self) -> uint64:
         return self._latest_timestamp
 
-    def contains_block(self, header_hash: bytes32) -> bool:
-        return header_hash in self._block_records
+    def contains_block(self, header_hash: bytes32, height: Optional[uint32] = None) -> bool:
+        """
+        True if we have already added this block to the chain. This may return false for orphan blocks
+        that we have added but no longer keep in memory.
+        """
+        if height is None:
+            block_rec = self.block_record(header_hash)
+            if block_rec is None:
+                return False
+            height = block_rec.height
+        block_hash_from_hh = self.height_to_hash(height)
+        assert block_hash_from_hh is not None
+        assert block_hash_from_hh == header_hash
+        return True
 
     def contains_height(self, height: uint32) -> bool:
         return height in self._height_to_hash

From 74b27f44121bcfbb00ce026f962ff6aeed5c4b62 Mon Sep 17 00:00:00 2001
From: almogdepaz <almogdepaz@gmail.com>
Date: Mon, 9 Dec 2024 12:40:39 +0200
Subject: [PATCH 04/22] fix recursion

---
 chia/consensus/blockchain.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/chia/consensus/blockchain.py b/chia/consensus/blockchain.py
index eff03500542f..773132585080 100644
--- a/chia/consensus/blockchain.py
+++ b/chia/consensus/blockchain.py
@@ -940,7 +940,7 @@ async def get_block_records_at(self, heights: list[uint32], batch_size: int = 90
         return records
 
     def try_block_record(self, header_hash: bytes32) -> Optional[BlockRecord]:
-        if self.contains_block(header_hash):
+        if header_hash in self.__block_records:
             return self.block_record(header_hash)
         return None
 

From b229e7e54c0bd9d594aab8ad7f32f64e988f19f6 Mon Sep 17 00:00:00 2001
From: almogdepaz <almogdepaz@gmail.com>
Date: Mon, 9 Dec 2024 12:53:06 +0200
Subject: [PATCH 05/22] handle none

---
 chia/_tests/util/blockchain_mock.py | 3 ++-
 chia/consensus/blockchain.py        | 3 ++-
 chia/util/augmented_chain.py        | 5 ++++-
 chia/util/block_cache.py            | 3 ++-
 chia/wallet/wallet_blockchain.py    | 3 ++-
 5 files changed, 12 insertions(+), 5 deletions(-)

diff --git a/chia/_tests/util/blockchain_mock.py b/chia/_tests/util/blockchain_mock.py
index 11fde9585c9d..80a9dca1bf4e 100644
--- a/chia/_tests/util/blockchain_mock.py
+++ b/chia/_tests/util/blockchain_mock.py
@@ -77,7 +77,8 @@ def contains_block(self, header_hash: bytes32, height: Optional[uint32] = None)
                 return False
             height = block_rec.height
         block_hash_from_hh = self.height_to_hash(height)
-        assert block_hash_from_hh is not None
+        if block_hash_from_hh is None:
+            return False
         assert block_hash_from_hh == header_hash
         return True
 
diff --git a/chia/consensus/blockchain.py b/chia/consensus/blockchain.py
index 773132585080..316043e22b24 100644
--- a/chia/consensus/blockchain.py
+++ b/chia/consensus/blockchain.py
@@ -780,7 +780,8 @@ def contains_block(self, header_hash: bytes32, height: Optional[uint32] = None)
                 return False
             height = block_rec.height
         block_hash_from_hh = self.height_to_hash(height)
-        assert block_hash_from_hh is not None
+        if block_hash_from_hh is None:
+            return False
         assert block_hash_from_hh == header_hash
         return True
 
diff --git a/chia/util/augmented_chain.py b/chia/util/augmented_chain.py
index 5b7489486e49..f101250cd726 100644
--- a/chia/util/augmented_chain.py
+++ b/chia/util/augmented_chain.py
@@ -123,8 +123,11 @@ def contains_block(self, header_hash: bytes32, height: Optional[uint32] = None)
             if block_rec is None:
                 return False
             height = block_rec.height
+        if not self.contains_height(height):
+            return False
         block_hash_from_hh = self.height_to_hash(height)
-        assert block_hash_from_hh is not None
+        if block_hash_from_hh is None:
+            return False
         assert block_hash_from_hh == header_hash
         return True
 
diff --git a/chia/util/block_cache.py b/chia/util/block_cache.py
index 0a705f3ac2c8..d4c43b20df43 100644
--- a/chia/util/block_cache.py
+++ b/chia/util/block_cache.py
@@ -54,7 +54,8 @@ def contains_block(self, header_hash: bytes32, height: Optional[uint32] = None)
                 return False
             height = block_rec.height
         block_hash_from_hh = self.height_to_hash(height)
-        assert block_hash_from_hh is not None
+        if block_hash_from_hh is None:
+            return False
         assert block_hash_from_hh == header_hash
         return True
 
diff --git a/chia/wallet/wallet_blockchain.py b/chia/wallet/wallet_blockchain.py
index 548657deb50c..45275190e28c 100644
--- a/chia/wallet/wallet_blockchain.py
+++ b/chia/wallet/wallet_blockchain.py
@@ -203,7 +203,8 @@ def contains_block(self, header_hash: bytes32, height: Optional[uint32] = None)
                 return False
             height = block_rec.height
         block_hash_from_hh = self.height_to_hash(height)
-        assert block_hash_from_hh is not None
+        if block_hash_from_hh is None:
+            return False
         assert block_hash_from_hh == header_hash
         return True
 

From 6feaf131720b290520499b1b7fe9cbe1830354a3 Mon Sep 17 00:00:00 2001
From: almogdepaz <almogdepaz@gmail.com>
Date: Mon, 9 Dec 2024 15:00:22 +0200
Subject: [PATCH 06/22] remove assert

---
 chia/_tests/util/blockchain_mock.py | 5 ++---
 chia/consensus/blockchain.py        | 5 ++---
 chia/util/augmented_chain.py        | 5 ++---
 chia/util/block_cache.py            | 5 ++---
 chia/wallet/wallet_blockchain.py    | 5 ++---
 5 files changed, 10 insertions(+), 15 deletions(-)

diff --git a/chia/_tests/util/blockchain_mock.py b/chia/_tests/util/blockchain_mock.py
index 80a9dca1bf4e..805686dacf46 100644
--- a/chia/_tests/util/blockchain_mock.py
+++ b/chia/_tests/util/blockchain_mock.py
@@ -77,9 +77,8 @@ def contains_block(self, header_hash: bytes32, height: Optional[uint32] = None)
                 return False
             height = block_rec.height
         block_hash_from_hh = self.height_to_hash(height)
-        if block_hash_from_hh is None:
-            return False
-        assert block_hash_from_hh == header_hash
+        if block_hash_from_hh is None or block_hash_from_hh != header_hash:
+            return False        
         return True
 
     async def contains_block_from_db(self, header_hash: bytes32) -> bool:
diff --git a/chia/consensus/blockchain.py b/chia/consensus/blockchain.py
index 316043e22b24..b59ab38494b9 100644
--- a/chia/consensus/blockchain.py
+++ b/chia/consensus/blockchain.py
@@ -780,9 +780,8 @@ def contains_block(self, header_hash: bytes32, height: Optional[uint32] = None)
                 return False
             height = block_rec.height
         block_hash_from_hh = self.height_to_hash(height)
-        if block_hash_from_hh is None:
-            return False
-        assert block_hash_from_hh == header_hash
+        if block_hash_from_hh is None or block_hash_from_hh != header_hash:
+            return False        
         return True
 
     def block_record(self, header_hash: bytes32) -> BlockRecord:
diff --git a/chia/util/augmented_chain.py b/chia/util/augmented_chain.py
index f101250cd726..56e56a73cae0 100644
--- a/chia/util/augmented_chain.py
+++ b/chia/util/augmented_chain.py
@@ -126,9 +126,8 @@ def contains_block(self, header_hash: bytes32, height: Optional[uint32] = None)
         if not self.contains_height(height):
             return False
         block_hash_from_hh = self.height_to_hash(height)
-        if block_hash_from_hh is None:
-            return False
-        assert block_hash_from_hh == header_hash
+        if block_hash_from_hh is None or block_hash_from_hh != header_hash:
+            return False        
         return True
 
     def contains_height(self, height: uint32) -> bool:
diff --git a/chia/util/block_cache.py b/chia/util/block_cache.py
index d4c43b20df43..917e3063e141 100644
--- a/chia/util/block_cache.py
+++ b/chia/util/block_cache.py
@@ -54,9 +54,8 @@ def contains_block(self, header_hash: bytes32, height: Optional[uint32] = None)
                 return False
             height = block_rec.height
         block_hash_from_hh = self.height_to_hash(height)
-        if block_hash_from_hh is None:
-            return False
-        assert block_hash_from_hh == header_hash
+        if block_hash_from_hh is None or block_hash_from_hh != header_hash:
+            return False        
         return True
 
     def contains_height(self, height: uint32) -> bool:
diff --git a/chia/wallet/wallet_blockchain.py b/chia/wallet/wallet_blockchain.py
index 45275190e28c..1458beb5574f 100644
--- a/chia/wallet/wallet_blockchain.py
+++ b/chia/wallet/wallet_blockchain.py
@@ -203,9 +203,8 @@ def contains_block(self, header_hash: bytes32, height: Optional[uint32] = None)
                 return False
             height = block_rec.height
         block_hash_from_hh = self.height_to_hash(height)
-        if block_hash_from_hh is None:
-            return False
-        assert block_hash_from_hh == header_hash
+        if block_hash_from_hh is None or block_hash_from_hh != header_hash:
+            return False        
         return True
 
     def contains_height(self, height: uint32) -> bool:

From 9c1eab13900e88f04fbadd979de9d456829ad423 Mon Sep 17 00:00:00 2001
From: almogdepaz <almogdepaz@gmail.com>
Date: Mon, 9 Dec 2024 15:01:00 +0200
Subject: [PATCH 07/22] remove assert

---
 chia/_tests/util/blockchain_mock.py | 2 +-
 chia/consensus/blockchain.py        | 2 +-
 chia/util/augmented_chain.py        | 2 +-
 chia/util/block_cache.py            | 2 +-
 chia/wallet/wallet_blockchain.py    | 2 +-
 5 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/chia/_tests/util/blockchain_mock.py b/chia/_tests/util/blockchain_mock.py
index 805686dacf46..d9c761a3fc7f 100644
--- a/chia/_tests/util/blockchain_mock.py
+++ b/chia/_tests/util/blockchain_mock.py
@@ -78,7 +78,7 @@ def contains_block(self, header_hash: bytes32, height: Optional[uint32] = None)
             height = block_rec.height
         block_hash_from_hh = self.height_to_hash(height)
         if block_hash_from_hh is None or block_hash_from_hh != header_hash:
-            return False        
+            return False
         return True
 
     async def contains_block_from_db(self, header_hash: bytes32) -> bool:
diff --git a/chia/consensus/blockchain.py b/chia/consensus/blockchain.py
index b59ab38494b9..41c4094045bb 100644
--- a/chia/consensus/blockchain.py
+++ b/chia/consensus/blockchain.py
@@ -781,7 +781,7 @@ def contains_block(self, header_hash: bytes32, height: Optional[uint32] = None)
             height = block_rec.height
         block_hash_from_hh = self.height_to_hash(height)
         if block_hash_from_hh is None or block_hash_from_hh != header_hash:
-            return False        
+            return False
         return True
 
     def block_record(self, header_hash: bytes32) -> BlockRecord:
diff --git a/chia/util/augmented_chain.py b/chia/util/augmented_chain.py
index 56e56a73cae0..8ceabfe434ac 100644
--- a/chia/util/augmented_chain.py
+++ b/chia/util/augmented_chain.py
@@ -127,7 +127,7 @@ def contains_block(self, header_hash: bytes32, height: Optional[uint32] = None)
             return False
         block_hash_from_hh = self.height_to_hash(height)
         if block_hash_from_hh is None or block_hash_from_hh != header_hash:
-            return False        
+            return False
         return True
 
     def contains_height(self, height: uint32) -> bool:
diff --git a/chia/util/block_cache.py b/chia/util/block_cache.py
index 917e3063e141..3ec8af4eb87d 100644
--- a/chia/util/block_cache.py
+++ b/chia/util/block_cache.py
@@ -55,7 +55,7 @@ def contains_block(self, header_hash: bytes32, height: Optional[uint32] = None)
             height = block_rec.height
         block_hash_from_hh = self.height_to_hash(height)
         if block_hash_from_hh is None or block_hash_from_hh != header_hash:
-            return False        
+            return False
         return True
 
     def contains_height(self, height: uint32) -> bool:
diff --git a/chia/wallet/wallet_blockchain.py b/chia/wallet/wallet_blockchain.py
index 1458beb5574f..5ffb3a83aefb 100644
--- a/chia/wallet/wallet_blockchain.py
+++ b/chia/wallet/wallet_blockchain.py
@@ -204,7 +204,7 @@ def contains_block(self, header_hash: bytes32, height: Optional[uint32] = None)
             height = block_rec.height
         block_hash_from_hh = self.height_to_hash(height)
         if block_hash_from_hh is None or block_hash_from_hh != header_hash:
-            return False        
+            return False
         return True
 
     def contains_height(self, height: uint32) -> bool:

From 14c3339947db3772fc049876d0ad20ba507fdd51 Mon Sep 17 00:00:00 2001
From: almogdepaz <almogdepaz@gmail.com>
Date: Mon, 9 Dec 2024 15:11:06 +0200
Subject: [PATCH 08/22] use try in wallet blockchain

---
 chia/wallet/wallet_blockchain.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/chia/wallet/wallet_blockchain.py b/chia/wallet/wallet_blockchain.py
index 5ffb3a83aefb..5d48d213007f 100644
--- a/chia/wallet/wallet_blockchain.py
+++ b/chia/wallet/wallet_blockchain.py
@@ -198,7 +198,7 @@ def contains_block(self, header_hash: bytes32, height: Optional[uint32] = None)
         that we have added but no longer keep in memory.
         """
         if height is None:
-            block_rec = self.block_record(header_hash)
+            block_rec = self.try_block_record(header_hash)
             if block_rec is None:
                 return False
             height = block_rec.height

From e0cd0cd02d8abbf3da8dfd6d4e3b7d1308ddd515 Mon Sep 17 00:00:00 2001
From: almogdepaz <almogdepaz@gmail.com>
Date: Mon, 23 Dec 2024 16:13:36 +0200
Subject: [PATCH 09/22] fix height o and height to hash in augmented

---
 chia/full_node/full_node.py  | 8 ++++----
 chia/util/augmented_chain.py | 1 -
 2 files changed, 4 insertions(+), 5 deletions(-)

diff --git a/chia/full_node/full_node.py b/chia/full_node/full_node.py
index def377fc5ef5..c2d2a487d598 100644
--- a/chia/full_node/full_node.py
+++ b/chia/full_node/full_node.py
@@ -699,9 +699,8 @@ async def short_sync_backtrack(
                         f"Failed to fetch block {curr_height} from {peer.get_peer_logging()}, wrong type {type(curr)}"
                     )
                 blocks.append(curr.block)
-                if (
-                    self.blockchain.contains_block(curr.block.prev_header_hash, curr.block.height - 1)
-                    or curr_height == 0
+                if curr_height == 0 or self.blockchain.contains_block(
+                    curr.block.prev_header_hash, curr.block.height - 1
                 ):
                     found_fork_point = True
                     break
@@ -738,6 +737,7 @@ async def new_peak(self, request: full_node_protocol.NewPeak, peer: WSChiaConnec
         """
 
         try:
+            self.log.info(f"adding {request.height}")
             seen_header_hash = self.sync_store.seen_header_hash(request.header_hash)
             # Updates heights in the UI. Sleeps 1.5s before, so other peers have time to update their peaks as well.
             # Limit to 3 refreshes.
@@ -1619,7 +1619,7 @@ async def add_prevalidated_blocks(
                 cc_sub_slot = block.finished_sub_slots[0].challenge_chain
                 if cc_sub_slot.new_sub_slot_iters is not None or cc_sub_slot.new_difficulty is not None:
                     expected_sub_slot_iters, expected_difficulty = get_next_sub_slot_iters_and_difficulty(
-                        self.constants, True, block_record, self.blockchain
+                        self.constants, True, block_record, blockchain
                     )
                     assert cc_sub_slot.new_sub_slot_iters is not None
                     vs.ssi = cc_sub_slot.new_sub_slot_iters
diff --git a/chia/util/augmented_chain.py b/chia/util/augmented_chain.py
index 8ceabfe434ac..68e6475998ef 100644
--- a/chia/util/augmented_chain.py
+++ b/chia/util/augmented_chain.py
@@ -87,7 +87,6 @@ def add_block_record(self, block_record: BlockRecord) -> None:
         # need to keep the extra block around anymore
         hh = block_record.header_hash
         if hh in self._extra_blocks:
-            del self._height_to_hash[block_record.height]
             del self._extra_blocks[hh]
 
     # BlockRecordsProtocol

From 31bc5b0b3d62261c6a1f770ac82a46eba7ed09a6 Mon Sep 17 00:00:00 2001
From: almogdepaz <almogdepaz@gmail.com>
Date: Tue, 14 Jan 2025 11:50:39 +0200
Subject: [PATCH 10/22] get_correct_block_for_fork

---
 chia/util/augmented_chain.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/chia/util/augmented_chain.py b/chia/util/augmented_chain.py
index 68e6475998ef..c034adf2b846 100644
--- a/chia/util/augmented_chain.py
+++ b/chia/util/augmented_chain.py
@@ -48,7 +48,8 @@ def add_extra_block(self, block: FullBlock, block_record: BlockRecord) -> None:
     def remove_extra_block(self, hh: bytes32) -> None:
         if hh in self._extra_blocks:
             block_record = self._extra_blocks.pop(hh)[1]
-            del self._height_to_hash[block_record.height]
+            block_rec_from_underlyning = self._underlying.block_record(block_record.header_hash)
+            assert block_rec_from_underlyning is not None
 
     # BlocksProtocol
     async def lookup_block_generators(self, header_hash: bytes32, generator_refs: set[uint32]) -> dict[uint32, bytes]:
@@ -108,6 +109,7 @@ def height_to_block_record(self, height: uint32) -> BlockRecord:
             ret = self._get_block_record(header_hash)
             if ret is not None:
                 return ret
+            return self._underlying.block_record(header_hash)
         return self._underlying.height_to_block_record(height)
 
     def height_to_hash(self, height: uint32) -> Optional[bytes32]:

From 1437e4ea1eb90e083d44bbc005140f191c067a20 Mon Sep 17 00:00:00 2001
From: almogdepaz <almogdepaz@gmail.com>
Date: Tue, 14 Jan 2025 12:00:56 +0200
Subject: [PATCH 11/22] remove debug code

---
 chia/full_node/full_node.py  | 1 -
 chia/util/augmented_chain.py | 4 +---
 2 files changed, 1 insertion(+), 4 deletions(-)

diff --git a/chia/full_node/full_node.py b/chia/full_node/full_node.py
index 988986c0de97..3396bfe341c4 100644
--- a/chia/full_node/full_node.py
+++ b/chia/full_node/full_node.py
@@ -742,7 +742,6 @@ async def new_peak(self, request: full_node_protocol.NewPeak, peer: WSChiaConnec
         """
 
         try:
-            self.log.info(f"adding {request.height}")
             seen_header_hash = self.sync_store.seen_header_hash(request.header_hash)
             # Updates heights in the UI. Sleeps 1.5s before, so other peers have time to update their peaks as well.
             # Limit to 3 refreshes.
diff --git a/chia/util/augmented_chain.py b/chia/util/augmented_chain.py
index c034adf2b846..45f6a32c3edf 100644
--- a/chia/util/augmented_chain.py
+++ b/chia/util/augmented_chain.py
@@ -47,9 +47,7 @@ def add_extra_block(self, block: FullBlock, block_record: BlockRecord) -> None:
 
     def remove_extra_block(self, hh: bytes32) -> None:
         if hh in self._extra_blocks:
-            block_record = self._extra_blocks.pop(hh)[1]
-            block_rec_from_underlyning = self._underlying.block_record(block_record.header_hash)
-            assert block_rec_from_underlyning is not None
+            self._extra_blocks.pop(hh)[1]
 
     # BlocksProtocol
     async def lookup_block_generators(self, header_hash: bytes32, generator_refs: set[uint32]) -> dict[uint32, bytes]:

From 4f144a6092440138865a90a715f239691b7846d6 Mon Sep 17 00:00:00 2001
From: almogdepaz <almogdepaz@gmail.com>
Date: Tue, 14 Jan 2025 13:16:45 +0200
Subject: [PATCH 12/22] use try_block_record

---
 chia/consensus/difficulty_adjustment.py | 10 ++++------
 1 file changed, 4 insertions(+), 6 deletions(-)

diff --git a/chia/consensus/difficulty_adjustment.py b/chia/consensus/difficulty_adjustment.py
index 63a85d27fb04..5fda123735c4 100644
--- a/chia/consensus/difficulty_adjustment.py
+++ b/chia/consensus/difficulty_adjustment.py
@@ -224,11 +224,10 @@ def _get_next_sub_slot_iters(
     if next_height < constants.EPOCH_BLOCKS:
         return uint64(constants.SUB_SLOT_ITERS_STARTING)
 
-    if not blocks.contains_block(prev_header_hash):
+    prev_b = blocks.try_block_record(prev_header_hash)
+    if prev_b is None:
         raise ValueError(f"Header hash {prev_header_hash} not in blocks")
 
-    prev_b: BlockRecord = blocks.block_record(prev_header_hash)
-
     # If we are in the same epoch, return same ssi
     if not skip_epoch_check:
         _, can_finish_epoch = can_finish_sub_and_full_epoch(
@@ -304,11 +303,10 @@ def _get_next_difficulty(
         # We are in the first epoch
         return uint64(constants.DIFFICULTY_STARTING)
 
-    if not blocks.contains_block(prev_header_hash):
+    prev_b = blocks.try_block_record(prev_header_hash)
+    if prev_b is None:
         raise ValueError(f"Header hash {prev_header_hash} not in blocks")
 
-    prev_b: BlockRecord = blocks.block_record(prev_header_hash)
-
     # If we are in the same slot as previous block, return same difficulty
     if not skip_epoch_check:
         _, can_finish_epoch = can_finish_sub_and_full_epoch(

From 5a918d7a569693e7fe0ec78e681f3d142e27c245 Mon Sep 17 00:00:00 2001
From: almogdepaz <almogdepaz@gmail.com>
Date: Mon, 20 Jan 2025 11:18:55 +0200
Subject: [PATCH 13/22] check main chain againt undelying blockchain

---
 chia/full_node/full_node.py             | 9 ++++++---
 chia/simulator/add_blocks_in_batches.py | 5 +++--
 chia/util/augmented_chain.py            | 2 +-
 3 files changed, 10 insertions(+), 6 deletions(-)

diff --git a/chia/full_node/full_node.py b/chia/full_node/full_node.py
index 3396bfe341c4..8f5266c149c2 100644
--- a/chia/full_node/full_node.py
+++ b/chia/full_node/full_node.py
@@ -620,6 +620,7 @@ async def short_sync_batch(self, peer: WSChiaConnection, start_height: uint32, t
                 fork_hash = self.constants.GENESIS_CHALLENGE
             assert fork_hash
             fork_info = ForkInfo(start_height - 1, start_height - 1, fork_hash)
+            blockchain = AugmentedBlockchain(self.blockchain)
             for height in range(start_height, target_height, batch_size):
                 end_height = min(target_height, height + batch_size)
                 request = RequestBlocks(uint32(height), uint32(end_height), True)
@@ -638,7 +639,7 @@ async def short_sync_batch(self, peer: WSChiaConnection, start_height: uint32, t
                     )
                     vs = ValidationState(ssi, diff, None)
                     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}")
@@ -1478,13 +1479,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:
@@ -1547,7 +1548,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:
@@ -1628,6 +1629,8 @@ async def add_prevalidated_blocks(
                     vs.ssi = cc_sub_slot.new_sub_slot_iters
                     assert cc_sub_slot.new_difficulty is not None
                     vs.difficulty = cc_sub_slot.new_difficulty
+                    if expected_sub_slot_iters != vs.ssi:
+                        self.log.info(f"block {block.height} fails")
                     assert expected_sub_slot_iters == vs.ssi
                     assert expected_difficulty == vs.difficulty
             result, error, state_change_summary = await self.blockchain.add_block(
diff --git a/chia/simulator/add_blocks_in_batches.py b/chia/simulator/add_blocks_in_batches.py
index dc0a1910060b..319bd3c02f97 100644
--- a/chia/simulator/add_blocks_in_batches.py
+++ b/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
 
@@ -34,14 +35,14 @@ async def add_blocks_in_batches(
     fork_info = ForkInfo(fork_height, blocks[0].height - 1, peak_hash)
 
     vs = ValidationState(ssi, diff, None)
-
+    blockchain = AugmentedBlockchain(full_node.blockchain)
     for block_batch in to_batches(blocks, 64):
         b = block_batch.entries[0]
         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()
         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:
diff --git a/chia/util/augmented_chain.py b/chia/util/augmented_chain.py
index 45f6a32c3edf..104fae86a1e1 100644
--- a/chia/util/augmented_chain.py
+++ b/chia/util/augmented_chain.py
@@ -81,7 +81,7 @@ async def get_block_record_from_db(self, header_hash: bytes32) -> Optional[Block
 
     def add_block_record(self, block_record: BlockRecord) -> None:
         self._underlying.add_block_record(block_record)
-
+        self._height_to_hash[block_record.height] = block_record.header_hash
         # now that we're adding the block to the underlying blockchain, we don't
         # need to keep the extra block around anymore
         hh = block_record.header_hash

From 5cf019ff5ecf7757e6670cf9a9a3698f519dac4c Mon Sep 17 00:00:00 2001
From: almogdepaz <almogdepaz@gmail.com>
Date: Mon, 20 Jan 2025 11:37:23 +0200
Subject: [PATCH 14/22] pass augmented in tests

---
 chia/_tests/util/full_sync.py               | 3 +++
 chia/_tests/wallet/sync/test_wallet_sync.py | 5 +++++
 tools/test_full_sync.py                     | 4 ++++
 3 files changed, 12 insertions(+)

diff --git a/chia/_tests/util/full_sync.py b/chia/_tests/util/full_sync.py
index 362afbbbfb0f..4d07f41d386b 100644
--- a/chia/_tests/util/full_sync.py
+++ b/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)
diff --git a/chia/_tests/wallet/sync/test_wallet_sync.py b/chia/_tests/wallet/sync/test_wallet_sync.py
index c9737b165fc6..92dea1984dda 100644
--- a/chia/_tests/wallet/sync/test_wallet_sync.py
+++ b/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
@@ -431,11 +432,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])
 
@@ -559,11 +562,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:
diff --git a/tools/test_full_sync.py b/tools/test_full_sync.py
index 0abb7db4efd8..b19e7cbfcc3c 100755
--- a/tools/test_full_sync.py
+++ b/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")

From 4215992a2290e9d02426d5fcc7c635c769710e65 Mon Sep 17 00:00:00 2001
From: almogdepaz <almogdepaz@gmail.com>
Date: Mon, 20 Jan 2025 11:37:57 +0200
Subject: [PATCH 15/22] lint

---
 chia/_tests/util/full_sync.py | 2 +-
 tools/test_full_sync.py       | 4 ++--
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/chia/_tests/util/full_sync.py b/chia/_tests/util/full_sync.py
index 4d07f41d386b..8db74afe7baa 100644
--- a/chia/_tests/util/full_sync.py
+++ b/chia/_tests/util/full_sync.py
@@ -218,7 +218,7 @@ async def run_sync_test(
                                 peer_info,
                                 ForkInfo(fork_height, fork_height, header_hash),
                                 ValidationState(ssi, diff, None),
-                                blockchain
+                                blockchain,
                             )
                             end_height = block_batch[-1].height
                             full_node.blockchain.clean_block_record(end_height - full_node.constants.BLOCKS_CACHE_SIZE)
diff --git a/tools/test_full_sync.py b/tools/test_full_sync.py
index b19e7cbfcc3c..9eb75ef222e5 100755
--- a/tools/test_full_sync.py
+++ b/tools/test_full_sync.py
@@ -171,7 +171,7 @@ async def run_sync_checkpoint(
                     peer_info,
                     ForkInfo(fork_height, fork_height, header_hash),
                     ValidationState(ssi, diff, None),
-                    blockchain
+                    blockchain,
                 )
                 end_height = block_batch[-1].height
                 full_node.blockchain.clean_block_record(end_height - full_node.constants.BLOCKS_CACHE_SIZE)
@@ -195,7 +195,7 @@ async def run_sync_checkpoint(
                     peer_info,
                     ForkInfo(fork_height, fork_height, fork_header_hash),
                     ValidationState(ssi, diff, None),
-                    blockchain
+                    blockchain,
                 )
                 if not success:
                     raise RuntimeError("failed to ingest block batch")

From 0ccadb56613150bf6fe036c89ab2f024963a9ffc Mon Sep 17 00:00:00 2001
From: almogdepaz <almogdepaz@gmail.com>
Date: Mon, 27 Jan 2025 15:23:00 +0200
Subject: [PATCH 16/22] revert changes to wallet blockchain

---
 chia/wallet/wallet_blockchain.py | 10 +---------
 1 file changed, 1 insertion(+), 9 deletions(-)

diff --git a/chia/wallet/wallet_blockchain.py b/chia/wallet/wallet_blockchain.py
index 5d48d213007f..7efffabe73c2 100644
--- a/chia/wallet/wallet_blockchain.py
+++ b/chia/wallet/wallet_blockchain.py
@@ -197,15 +197,7 @@ def contains_block(self, header_hash: bytes32, height: Optional[uint32] = None)
         True if we have already added this block to the chain. This may return false for orphan blocks
         that we have added but no longer keep in memory.
         """
-        if height is None:
-            block_rec = self.try_block_record(header_hash)
-            if block_rec is None:
-                return False
-            height = block_rec.height
-        block_hash_from_hh = self.height_to_hash(height)
-        if block_hash_from_hh is None or block_hash_from_hh != header_hash:
-            return False
-        return True
+        return header_hash in self._block_records
 
     def contains_height(self, height: uint32) -> bool:
         return height in self._height_to_hash

From a48ce05b0528236947c6376d66005be409c232fc Mon Sep 17 00:00:00 2001
From: almogdepaz <almogdepaz@gmail.com>
Date: Mon, 27 Jan 2025 19:32:55 +0200
Subject: [PATCH 17/22] remove redundant comment

---
 chia/_tests/util/blockchain_mock.py | 6 +-----
 chia/consensus/blockchain.py        | 6 +-----
 chia/util/block_cache.py            | 6 +-----
 3 files changed, 3 insertions(+), 15 deletions(-)

diff --git a/chia/_tests/util/blockchain_mock.py b/chia/_tests/util/blockchain_mock.py
index d9c761a3fc7f..a2f8a7d3812a 100644
--- a/chia/_tests/util/blockchain_mock.py
+++ b/chia/_tests/util/blockchain_mock.py
@@ -66,11 +66,7 @@ def height_to_hash(self, height: uint32) -> Optional[bytes32]:
         assert height in self._height_to_hash
         return self._height_to_hash[height]
 
-    def contains_block(self, header_hash: bytes32, height: Optional[uint32] = None) -> bool:
-        """
-        True if we have already added this block to the chain. This may return false for orphan blocks
-        that we have added but no longer keep in memory.
-        """
+    def contains_block(self, header_hash: bytes32, height: Optional[uint32] = None) -> bool:        
         if height is None:
             block_rec = self.try_block_record(header_hash)
             if block_rec is None:
diff --git a/chia/consensus/blockchain.py b/chia/consensus/blockchain.py
index 7c21bb715603..8fd4dd7c7ae5 100644
--- a/chia/consensus/blockchain.py
+++ b/chia/consensus/blockchain.py
@@ -784,11 +784,7 @@ async def validate_unfinished_block(
 
         return PreValidationResult(None, required_iters, conds, uint32(0))
 
-    def contains_block(self, header_hash: bytes32, height: Optional[uint32] = None) -> bool:
-        """
-        True if we have already added this block to the chain. This may return false for orphan blocks
-        that we have added but no longer keep in memory.
-        """
+    def contains_block(self, header_hash: bytes32, height: Optional[uint32] = None) -> bool:        
         if height is None:
             block_rec = self.try_block_record(header_hash)
             if block_rec is None:
diff --git a/chia/util/block_cache.py b/chia/util/block_cache.py
index 3ec8af4eb87d..8218c9e9d4f3 100644
--- a/chia/util/block_cache.py
+++ b/chia/util/block_cache.py
@@ -43,11 +43,7 @@ def height_to_hash(self, height: uint32) -> Optional[bytes32]:
             return None
         return self._height_to_hash[height]
 
-    def contains_block(self, header_hash: bytes32, height: Optional[uint32] = None) -> bool:
-        """
-        True if we have already added this block to the chain. This may return false for orphan blocks
-        that we have added but no longer keep in memory.
-        """
+    def contains_block(self, header_hash: bytes32, height: Optional[uint32] = None) -> bool:        
         if height is None:
             block_rec = self.try_block_record(header_hash)
             if block_rec is None:

From bf03d6d3029a39cedc68390bdb3499c9533cc973 Mon Sep 17 00:00:00 2001
From: almogdepaz <almogdepaz@gmail.com>
Date: Mon, 27 Jan 2025 19:33:29 +0200
Subject: [PATCH 18/22] lint

---
 chia/_tests/util/blockchain_mock.py | 2 +-
 chia/consensus/blockchain.py        | 2 +-
 chia/util/block_cache.py            | 2 +-
 3 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/chia/_tests/util/blockchain_mock.py b/chia/_tests/util/blockchain_mock.py
index a2f8a7d3812a..f924e422b1cf 100644
--- a/chia/_tests/util/blockchain_mock.py
+++ b/chia/_tests/util/blockchain_mock.py
@@ -66,7 +66,7 @@ def height_to_hash(self, height: uint32) -> Optional[bytes32]:
         assert height in self._height_to_hash
         return self._height_to_hash[height]
 
-    def contains_block(self, header_hash: bytes32, height: Optional[uint32] = None) -> bool:        
+    def contains_block(self, header_hash: bytes32, height: Optional[uint32] = None) -> bool:
         if height is None:
             block_rec = self.try_block_record(header_hash)
             if block_rec is None:
diff --git a/chia/consensus/blockchain.py b/chia/consensus/blockchain.py
index 8fd4dd7c7ae5..6a98498fba6b 100644
--- a/chia/consensus/blockchain.py
+++ b/chia/consensus/blockchain.py
@@ -784,7 +784,7 @@ async def validate_unfinished_block(
 
         return PreValidationResult(None, required_iters, conds, uint32(0))
 
-    def contains_block(self, header_hash: bytes32, height: Optional[uint32] = None) -> bool:        
+    def contains_block(self, header_hash: bytes32, height: Optional[uint32] = None) -> bool:
         if height is None:
             block_rec = self.try_block_record(header_hash)
             if block_rec is None:
diff --git a/chia/util/block_cache.py b/chia/util/block_cache.py
index 8218c9e9d4f3..0326fc2b4511 100644
--- a/chia/util/block_cache.py
+++ b/chia/util/block_cache.py
@@ -43,7 +43,7 @@ def height_to_hash(self, height: uint32) -> Optional[bytes32]:
             return None
         return self._height_to_hash[height]
 
-    def contains_block(self, header_hash: bytes32, height: Optional[uint32] = None) -> bool:        
+    def contains_block(self, header_hash: bytes32, height: Optional[uint32] = None) -> bool:
         if height is None:
             block_rec = self.try_block_record(header_hash)
             if block_rec is None:

From 8c3271d7d7c0a220da05636d8e23f83e84bc7acb Mon Sep 17 00:00:00 2001
From: almogdepaz <almogdepaz@gmail.com>
Date: Mon, 27 Jan 2025 22:16:22 +0200
Subject: [PATCH 19/22] remove debug message

---
 chia/full_node/full_node.py | 2 --
 1 file changed, 2 deletions(-)

diff --git a/chia/full_node/full_node.py b/chia/full_node/full_node.py
index 7d41ffe5ff98..c63bab554337 100644
--- a/chia/full_node/full_node.py
+++ b/chia/full_node/full_node.py
@@ -1626,8 +1626,6 @@ async def add_prevalidated_blocks(
                     vs.ssi = cc_sub_slot.new_sub_slot_iters
                     assert cc_sub_slot.new_difficulty is not None
                     vs.difficulty = cc_sub_slot.new_difficulty
-                    if expected_sub_slot_iters != vs.ssi:
-                        self.log.info(f"block {block.height} fails")
                     assert expected_sub_slot_iters == vs.ssi
                     assert expected_difficulty == vs.difficulty
             block_rec = blockchain.block_record(block.header_hash)

From 45337112dcb48521df0c28383215b2ccde3ac2c0 Mon Sep 17 00:00:00 2001
From: almogdepaz <almogdepaz@gmail.com>
Date: Mon, 3 Feb 2025 12:30:59 +0200
Subject: [PATCH 20/22] revert bad merge

---
 chia/simulator/add_blocks_in_batches.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/chia/simulator/add_blocks_in_batches.py b/chia/simulator/add_blocks_in_batches.py
index b5a8508d434b..319bd3c02f97 100644
--- a/chia/simulator/add_blocks_in_batches.py
+++ b/chia/simulator/add_blocks_in_batches.py
@@ -41,7 +41,6 @@ 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, blockchain
         )

From ec8ef125b1a47ccc4fcf3d9da0195dba802988f1 Mon Sep 17 00:00:00 2001
From: almogdepaz <almogdepaz@gmail.com>
Date: Wed, 5 Feb 2025 18:45:36 +0200
Subject: [PATCH 21/22] make height mandatory

---
 chia/_tests/blockchain/test_augmented_chain.py |  4 ++--
 chia/consensus/blockchain.py                   |  6 ++----
 chia/consensus/blockchain_interface.py         |  2 +-
 chia/full_node/full_node.py                    | 10 +++++-----
 chia/full_node/full_node_api.py                |  2 +-
 chia/util/augmented_chain.py                   |  2 +-
 6 files changed, 12 insertions(+), 14 deletions(-)

diff --git a/chia/_tests/blockchain/test_augmented_chain.py b/chia/_tests/blockchain/test_augmented_chain.py
index 1dcf51ebe8ec..4a54bf6f7c04 100644
--- a/chia/_tests/blockchain/test_augmented_chain.py
+++ b/chia/_tests/blockchain/test_augmented_chain.py
@@ -122,13 +122,13 @@ async def test_augmented_chain(default_10000_blocks: list[FullBlock]) -> None:
         assert abc.try_block_record(blocks[i].header_hash) == block_records[i]
         assert abc.height_to_hash(uint32(i)) == blocks[i].header_hash
         assert await abc.prev_block_hash([blocks[i].header_hash]) == [blocks[i].prev_header_hash]
-        assert abc.contains_block(blocks[i].header_hash) is True
+        assert abc.block_record(blocks[i].header_hash) is not None
         assert await abc.get_block_record_from_db(blocks[i].header_hash) == block_records[i]
         assert abc.contains_height(uint32(i))
 
     for i in range(5, 10):
         assert abc.height_to_hash(uint32(i)) is None
-        assert not abc.contains_block(blocks[i].header_hash)
+        assert abc.block_record(blocks[i].header_hash) is None
         assert not await abc.get_block_record_from_db(blocks[i].header_hash)
         assert not abc.contains_height(uint32(i))
 
diff --git a/chia/consensus/blockchain.py b/chia/consensus/blockchain.py
index 910b833dbfec..c7ba22d6b6c7 100644
--- a/chia/consensus/blockchain.py
+++ b/chia/consensus/blockchain.py
@@ -609,7 +609,6 @@ async def _reconsider_peak(
         )
 
     def get_next_difficulty(self, header_hash: bytes32, new_slot: bool) -> uint64:
-        assert self.contains_block(header_hash)
         curr = self.block_record(header_hash)
         if curr.height <= 2:
             return self.constants.DIFFICULTY_STARTING
@@ -617,7 +616,6 @@ def get_next_difficulty(self, header_hash: bytes32, new_slot: bool) -> uint64:
         return get_next_sub_slot_iters_and_difficulty(self.constants, new_slot, curr, self)[1]
 
     def get_next_slot_iters(self, header_hash: bytes32, new_slot: bool) -> uint64:
-        assert self.contains_block(header_hash)
         curr = self.block_record(header_hash)
         if curr.height <= 2:
             return self.constants.SUB_SLOT_ITERS_STARTING
@@ -704,7 +702,7 @@ async def validate_unfinished_block_header(
             return None, Err.TOO_MANY_GENERATOR_REFS
 
         if (
-            not self.contains_block(block.prev_header_hash)
+            self.block_record(block.prev_header_hash) is None
             and block.prev_header_hash != self.constants.GENESIS_CHALLENGE
         ):
             return None, Err.INVALID_PREV_BLOCK_HASH
@@ -788,7 +786,7 @@ async def validate_unfinished_block(
 
         return PreValidationResult(None, required_iters, conds, uint32(0))
 
-    def contains_block(self, header_hash: bytes32, height: Optional[uint32] = None) -> bool:
+    def contains_block(self, header_hash: bytes32, height: uint32) -> bool:
         if height is None:
             block_rec = self.try_block_record(header_hash)
             if block_rec is None:
diff --git a/chia/consensus/blockchain_interface.py b/chia/consensus/blockchain_interface.py
index e3550f6966ca..d0bb85ca74a4 100644
--- a/chia/consensus/blockchain_interface.py
+++ b/chia/consensus/blockchain_interface.py
@@ -14,7 +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 contains_block(self, header_hash: bytes32, height: Optional[uint32] = None) -> bool: ...
+    def contains_block(self, header_hash: bytes32, height: uint32) -> bool: ...
     def height_to_hash(self, height: uint32) -> Optional[bytes32]: ...
     def height_to_block_record(self, height: uint32) -> BlockRecord: ...
 
diff --git a/chia/full_node/full_node.py b/chia/full_node/full_node.py
index 778c863263f6..19117a029ae8 100644
--- a/chia/full_node/full_node.py
+++ b/chia/full_node/full_node.py
@@ -637,7 +637,6 @@ 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, blockchain
                     )
@@ -2019,7 +2018,7 @@ async def add_block(
 
         # Adds the block to seen, and check if it's seen before (which means header is in memory)
         header_hash = block.header_hash
-        if self.blockchain.contains_block(header_hash):
+        if self.blockchain.contains_block(header_hash, block.height):
             if fork_info is not None:
                 await self.blockchain.run_single_block(block, fork_info)
             return None
@@ -2093,7 +2092,7 @@ async def add_block(
             enable_profiler(self.profile_block_validation) as pr,
         ):
             # After acquiring the lock, check again, because another asyncio thread might have added it
-            if self.blockchain.contains_block(header_hash):
+            if self.blockchain.contains_block(header_hash, block.height):
                 if fork_info is not None:
                     await self.blockchain.run_single_block(block, fork_info)
                 return None
@@ -2281,8 +2280,9 @@ async def add_unfinished_block(
         """
         receive_time = time.time()
 
-        if block.prev_header_hash != self.constants.GENESIS_CHALLENGE and not self.blockchain.contains_block(
-            block.prev_header_hash
+        if (
+            block.prev_header_hash != self.constants.GENESIS_CHALLENGE
+            and self.blockchain.block_record(block.prev_header_hash) is None
         ):
             # No need to request the parent, since the peer will send it to us anyway, via NewPeak
             self.log.debug("Received a disconnected unfinished block")
diff --git a/chia/full_node/full_node_api.py b/chia/full_node/full_node_api.py
index f6bcc802172d..53b1e55cb296 100644
--- a/chia/full_node/full_node_api.py
+++ b/chia/full_node/full_node_api.py
@@ -285,7 +285,7 @@ async def respond_transaction(
     async def request_proof_of_weight(self, request: full_node_protocol.RequestProofOfWeight) -> Optional[Message]:
         if self.full_node.weight_proof_handler is None:
             return None
-        if not self.full_node.blockchain.contains_block(request.tip):
+        if self.full_node.blockchain.block_record(request.tip) is None:
             self.log.error(f"got weight proof request for unknown peak {request.tip}")
             return None
         if request.tip in self.full_node.pow_creation:
diff --git a/chia/util/augmented_chain.py b/chia/util/augmented_chain.py
index 104fae86a1e1..0e9f78aeed13 100644
--- a/chia/util/augmented_chain.py
+++ b/chia/util/augmented_chain.py
@@ -116,7 +116,7 @@ def height_to_hash(self, height: uint32) -> Optional[bytes32]:
             return ret
         return self._underlying.height_to_hash(height)
 
-    def contains_block(self, header_hash: bytes32, height: Optional[uint32] = None) -> bool:
+    def contains_block(self, header_hash: bytes32, height: uint32) -> bool:
         if height is None:
             block_rec = self.try_block_record(header_hash)
             if block_rec is None:

From 30d0323a604c1a1da3a12cd0876f1599a438fe55 Mon Sep 17 00:00:00 2001
From: almogdepaz <almogdepaz@gmail.com>
Date: Wed, 5 Feb 2025 18:53:19 +0200
Subject: [PATCH 22/22] make height mandatory

---
 chia/_tests/blockchain/test_augmented_chain.py | 2 +-
 chia/consensus/blockchain.py                   | 5 -----
 chia/util/augmented_chain.py                   | 7 -------
 3 files changed, 1 insertion(+), 13 deletions(-)

diff --git a/chia/_tests/blockchain/test_augmented_chain.py b/chia/_tests/blockchain/test_augmented_chain.py
index 4a54bf6f7c04..113f75e6fed7 100644
--- a/chia/_tests/blockchain/test_augmented_chain.py
+++ b/chia/_tests/blockchain/test_augmented_chain.py
@@ -46,7 +46,7 @@ def height_to_block_record(self, height: uint32) -> BlockRecord:
     def height_to_hash(self, height: uint32) -> Optional[bytes32]:
         return self.heights.get(height)
 
-    def contains_block(self, header_hash: bytes32, height: Optional[uint32] = None) -> bool:
+    def contains_block(self, header_hash: bytes32, height: uint32) -> bool:
         return False  # pragma: no cover
 
     def contains_height(self, height: uint32) -> bool:
diff --git a/chia/consensus/blockchain.py b/chia/consensus/blockchain.py
index c7ba22d6b6c7..bc57734904cc 100644
--- a/chia/consensus/blockchain.py
+++ b/chia/consensus/blockchain.py
@@ -787,11 +787,6 @@ async def validate_unfinished_block(
         return PreValidationResult(None, required_iters, conds, uint32(0))
 
     def contains_block(self, header_hash: bytes32, height: uint32) -> bool:
-        if height is None:
-            block_rec = self.try_block_record(header_hash)
-            if block_rec is None:
-                return False
-            height = block_rec.height
         block_hash_from_hh = self.height_to_hash(height)
         if block_hash_from_hh is None or block_hash_from_hh != header_hash:
             return False
diff --git a/chia/util/augmented_chain.py b/chia/util/augmented_chain.py
index 0e9f78aeed13..dbe375160a27 100644
--- a/chia/util/augmented_chain.py
+++ b/chia/util/augmented_chain.py
@@ -117,13 +117,6 @@ def height_to_hash(self, height: uint32) -> Optional[bytes32]:
         return self._underlying.height_to_hash(height)
 
     def contains_block(self, header_hash: bytes32, height: uint32) -> bool:
-        if height is None:
-            block_rec = self.try_block_record(header_hash)
-            if block_rec is None:
-                return False
-            height = block_rec.height
-        if not self.contains_height(height):
-            return False
         block_hash_from_hh = self.height_to_hash(height)
         if block_hash_from_hh is None or block_hash_from_hh != header_hash:
             return False