From c110da40345f70220a8a0ab95aa46fbf99deefd7 Mon Sep 17 00:00:00 2001
From: Mac L <mjladson@pm.me>
Date: Fri, 28 Jun 2024 16:09:27 +0400
Subject: [PATCH] Use specific_variant_attributes

---
 consensus/types/src/attestation.rs            |  5 +--
 consensus/types/src/beacon_state.rs           | 34 +++++++++++--------
 consensus/types/src/execution_payload.rs      |  5 ++-
 .../types/src/execution_payload_header.rs     |  5 ++-
 4 files changed, 31 insertions(+), 18 deletions(-)

diff --git a/consensus/types/src/attestation.rs b/consensus/types/src/attestation.rs
index d54ce72f660..9e8612e9f3a 100644
--- a/consensus/types/src/attestation.rs
+++ b/consensus/types/src/attestation.rs
@@ -78,10 +78,11 @@ pub struct Attestation<E: EthSpec> {
     pub aggregation_bits: BitList<E::MaxValidatorsPerSlot>,
     #[tree_hash(stable_index = 1)]
     pub data: AttestationData,
-    #[superstruct(only(Electra))]
     #[tree_hash(stable_index = 2)]
-    pub committee_bits: BitVector<E::MaxCommitteesPerSlot>,
     pub signature: AggregateSignature,
+    #[superstruct(only(Electra))]
+    #[tree_hash(stable_index = 3)]
+    pub committee_bits: BitVector<E::MaxCommitteesPerSlot>,
 }
 
 impl<E: EthSpec> Hash for Attestation<E> {
diff --git a/consensus/types/src/beacon_state.rs b/consensus/types/src/beacon_state.rs
index 2e3f4e62e59..4a786eb9928 100644
--- a/consensus/types/src/beacon_state.rs
+++ b/consensus/types/src/beacon_state.rs
@@ -234,7 +234,6 @@ impl From<BeaconStateHash> for Hash256 {
         serde(bound = "E: EthSpec", deny_unknown_fields),
         arbitrary(bound = "E: EthSpec"),
         derivative(Clone),
-        tree_hash(struct_behaviour = "profile", max_fields = "typenum::U128"),
     ),
     specific_variant_attributes(
         Base(metastruct(
@@ -307,20 +306,27 @@ impl From<BeaconStateHash> for Hash256 {
             )),
             num_fields(all()),
         )),
-        Electra(metastruct(
-            mappings(
-                map_beacon_state_electra_fields(),
-                map_beacon_state_electra_tree_list_fields(mutable, fallible, groups(tree_lists)),
-                map_beacon_state_electra_tree_list_fields_immutable(groups(tree_lists)),
+        Electra(
+            metastruct(
+                mappings(
+                    map_beacon_state_electra_fields(),
+                    map_beacon_state_electra_tree_list_fields(
+                        mutable,
+                        fallible,
+                        groups(tree_lists)
+                    ),
+                    map_beacon_state_electra_tree_list_fields_immutable(groups(tree_lists)),
+                ),
+                bimappings(bimap_beacon_state_electra_tree_list_fields(
+                    other_type = "BeaconStateElectra",
+                    self_mutable,
+                    fallible,
+                    groups(tree_lists)
+                )),
+                num_fields(all()),
             ),
-            bimappings(bimap_beacon_state_electra_tree_list_fields(
-                other_type = "BeaconStateElectra",
-                self_mutable,
-                fallible,
-                groups(tree_lists)
-            )),
-            num_fields(all()),
-        ))
+            tree_hash(struct_behaviour = "profile", max_fields = "typenum::U128")
+        ),
     ),
     cast_error(ty = "Error", expr = "Error::IncorrectStateVariant"),
     partial_getter_error(ty = "Error", expr = "Error::IncorrectStateVariant"),
diff --git a/consensus/types/src/execution_payload.rs b/consensus/types/src/execution_payload.rs
index 87f45311fbf..a4b85d28764 100644
--- a/consensus/types/src/execution_payload.rs
+++ b/consensus/types/src/execution_payload.rs
@@ -33,8 +33,11 @@ pub type Withdrawals<E> = VariableList<Withdrawal, <E as EthSpec>::MaxWithdrawal
         derivative(PartialEq, Hash(bound = "E: EthSpec")),
         serde(bound = "E: EthSpec", deny_unknown_fields),
         arbitrary(bound = "E: EthSpec"),
-        tree_hash(max_fields = "typenum::U64")
     ),
+    specific_variant_attributes(Electra(tree_hash(
+        struct_behaviour = "profile",
+        max_fields = "typenum::U64"
+    ))),
     cast_error(ty = "Error", expr = "BeaconStateError::IncorrectStateVariant"),
     partial_getter_error(ty = "Error", expr = "BeaconStateError::IncorrectStateVariant"),
     map_into(FullPayload, BlindedPayload),
diff --git a/consensus/types/src/execution_payload_header.rs b/consensus/types/src/execution_payload_header.rs
index b86dd5a62f8..cd25e689a25 100644
--- a/consensus/types/src/execution_payload_header.rs
+++ b/consensus/types/src/execution_payload_header.rs
@@ -26,8 +26,11 @@ use tree_hash_derive::TreeHash;
         derivative(PartialEq, Hash(bound = "E: EthSpec")),
         serde(bound = "E: EthSpec", deny_unknown_fields),
         arbitrary(bound = "E: EthSpec"),
-        tree_hash(max_fields = "typenum::U64")
     ),
+    specific_variant_attributes(Electra(tree_hash(
+        struct_behaviour = "profile",
+        max_fields = "typenum::U64"
+    ))),
     ref_attributes(
         derive(PartialEq, TreeHash, Debug),
         tree_hash(enum_behaviour = "transparent")