refactor: rewrite justification operations as pure functions (leanEthereum#165)

* refactor: rewrite justification operations as pure functions

* fix: rename for readability

* fix: linting

* fix: linting

* fix: linting

Author: unnawut

src/lean_spec/subspecs/containers/state/helpers.py

@@ -0,0 +1,88 @@
+"""Helpers for the State container."""
+
+from typing import Dict, List
+
+from lean_spec.subspecs.containers.state.types import (
+    JustificationRoots,
+    JustificationValidators,
+)
+from lean_spec.types import Boolean, Bytes32
+
+
+def get_justifications_map(
+    justifications_roots: JustificationRoots,
+    justifications_validators: JustificationValidators,
+    validator_count: int,
+) -> Dict[Bytes32, List[Boolean]]:
+    """
+    Reconstruct the justifications map from the state's flat data structures.
+
+    Parameters
+    ----------
+    justifications_roots : JustificationRoots
+        The block roots in alphabetical order.
+    justifications_validators : JustificationValidators
+        The list of validator justifications for each block root concatenated in the same order
+        as the list of block roots.
+    validator_count : int
+        The number of validators in the state.
+
+    Returns:
+    -------
+    Dict[Bytes32, List[Boolean]]
+        A mapping from a block root to the list of validator justifications for that root.
+    """
+    # No justified roots means no justifications to reconstruct.
+    if not justifications_roots:
+        return {}
+
+    # Extract the flattened validator justifications.
+    flat_justifications = list(justifications_validators)
+
+    # Reconstruct the map: each root gets its corresponding justification slice.
+    return {
+        root: flat_justifications[i * validator_count : (i + 1) * validator_count]
+        for i, root in enumerate(justifications_roots)
+    }
+
+
+def flatten_justifications_map(
+    justifications_map: Dict[Bytes32, List[Boolean]], validator_count: int
+) -> tuple[JustificationRoots, JustificationValidators]:
+    """
+    Flatten a map of validator justifications into the state's flat data structures
+    for SSZ compatibility.
+
+    Parameters
+    ----------
+    justifications_map : Dict[Bytes32, List[Boolean]]
+        A mapping from a block root to the list of validator justifications for that root.
+    validator_count : int
+        The number of validators in the state.
+
+    Returns:
+    -------
+    JustificationRoots
+        The block roots in alphabetical order.
+    JustificationValidators
+        The list of validator justifications for each block root concatenated in the same order
+        as the list of block roots.
+    """
+    # Build the flattened lists from the map, with sorted keys for deterministic order.
+    roots_list = []
+    justifications_list = []
+
+    for root in sorted(justifications_map.keys()):
+        justifications = justifications_map[root]
+
+        # Validate that the justifications list has the expected length.
+        if len(justifications) != validator_count:
+            raise AssertionError(f"Justifications list for root {root.hex()} has incorrect length")
+
+        # Add the root to the roots list.
+        roots_list.append(root)
+        # Extend the flattened list with the justifications for this root.
+        justifications_list.extend(justifications)
+
+    # Return a new state object with the updated fields.
+    return JustificationRoots(data=roots_list), JustificationValidators(data=justifications_list)

src/lean_spec/subspecs/containers/state/state.py

@@ -6,8 +6,6 @@
 justified and finalized.
 """
 
-from typing import Dict, List
-
 from lean_spec.subspecs.ssz.constants import ZERO_HASH
 from lean_spec.subspecs.ssz.hash import hash_tree_root
 from lean_spec.types import (
@@ -137,77 +135,6 @@ def is_proposer(self, validator_index: ValidatorIndex) -> bool:
             num_validators=Uint64(self.validators.count),
         )
 
-    def get_justifications(self) -> Dict[Bytes32, List[Boolean]]:
-        """
-        Reconstruct a map from justified block roots to validator vote lists.
-
-        This method takes the flat state encoding and rebuilds the associative
-        structure for easier processing.
-
-        Returns:
-        -------
-        Dict[Bytes32, List[Boolean]]
-            A mapping from justified block root to the list of validator votes.
-        """
-        # No justified roots means no justifications to reconstruct.
-        if not self.justifications_roots:
-            return {}
-
-        # Each root has exactly validator_count votes in the flat encoding.
-        validator_count = self.validators.count
-
-        # Extract the flattened validator votes.
-        flat_votes = list(self.justifications_validators)
-
-        # Reconstruct the map: each root gets its corresponding vote slice.
-        return {
-            root: flat_votes[i * validator_count : (i + 1) * validator_count]
-            for i, root in enumerate(self.justifications_roots)
-        }
-
-    def with_justifications(
-        self,
-        justifications: Dict[Bytes32, List[Boolean]],
-    ) -> "State":
-        """
-        Update the state with a new set of justifications.
-
-        This method flattens the justifications map into the state's flat
-        encoding for SSZ compatibility.
-
-        Parameters
-        ----------
-        justifications : Dict[Bytes32, List[Boolean]]
-            A mapping from justified block root to validator vote lists.
-
-        Returns:
-        -------
-        State
-            A new state with updated justification data.
-        """
-        # Build the flattened lists from the map, with sorted keys for deterministic order.
-        roots_list = []
-        votes_list = []
-        for root in sorted(justifications.keys()):
-            votes = justifications[root]
-            # Validate that the vote list has the expected length.
-            expected_len = self.validators.count
-            if len(votes) != expected_len:
-                raise AssertionError(f"Vote list for root {root.hex()} has incorrect length")
-
-            # Add the root to the roots list.
-            roots_list.append(root)
-            # Extend the flattened list with the votes for this root.
-            votes_list.extend(votes)
-
-        # Return a new state object with the updated fields.
-        return self.model_copy(
-            update={
-                "justifications_roots": JustificationRoots(data=roots_list),
-                "justifications_validators": JustificationValidators(data=votes_list),
-            }
-        )
-
     def process_slot(self) -> "State":
         """
         Perform per-slot maintenance tasks.