From 9b9877a3be6bdf3c733ffba399fe47b663145c3d Mon Sep 17 00:00:00 2001
From: Thomas Coratger <thomas.coratger@gmail.com>
Date: Sat, 21 Mar 2026 23:54:10 +0100
Subject: [PATCH] store: more compact prune_stale_attestation_data

---
 src/lean_spec/subspecs/forkchoice/store.py | 38 +++++++++-------------
 1 file changed, 15 insertions(+), 23 deletions(-)

diff --git a/src/lean_spec/subspecs/forkchoice/store.py b/src/lean_spec/subspecs/forkchoice/store.py
index 4b4b18a3..fbeb1130 100644
--- a/src/lean_spec/subspecs/forkchoice/store.py
+++ b/src/lean_spec/subspecs/forkchoice/store.py
@@ -244,36 +244,28 @@ def prune_stale_attestation_data(self) -> "Store":
         Returns:
             New Store with stale attestation data removed.
         """
-        finalized_slot = self.latest_finalized.slot
-
         # Filter out stale entries from all attestation-related mappings.
         #
         # Each mapping is keyed by attestation data, so we check membership by slot
         # against the finalized slot.
 
-        new_attestation_sigs = {
-            attestation_data: sigs
-            for attestation_data, sigs in self.attestation_signatures.items()
-            if attestation_data.target.slot > finalized_slot
-        }
-
-        new_aggregated_new = {
-            attestation_data: proofs
-            for attestation_data, proofs in self.latest_new_aggregated_payloads.items()
-            if attestation_data.target.slot > finalized_slot
-        }
-
-        new_aggregated_known = {
-            attestation_data: proofs
-            for attestation_data, proofs in self.latest_known_aggregated_payloads.items()
-            if attestation_data.target.slot > finalized_slot
-        }
-
         return self.model_copy(
             update={
-                "attestation_signatures": new_attestation_sigs,
-                "latest_new_aggregated_payloads": new_aggregated_new,
-                "latest_known_aggregated_payloads": new_aggregated_known,
+                "attestation_signatures": {
+                    attestation_data: sigs
+                    for attestation_data, sigs in self.attestation_signatures.items()
+                    if attestation_data.target.slot > self.latest_finalized.slot
+                },
+                "latest_new_aggregated_payloads": {
+                    attestation_data: proofs
+                    for attestation_data, proofs in self.latest_new_aggregated_payloads.items()
+                    if attestation_data.target.slot > self.latest_finalized.slot
+                },
+                "latest_known_aggregated_payloads": {
+                    attestation_data: proofs
+                    for attestation_data, proofs in self.latest_known_aggregated_payloads.items()
+                    if attestation_data.target.slot > self.latest_finalized.slot
+                },
             }
         )