From e0d246210fd28c39d1d641c98e2bce06374fc00a Mon Sep 17 00:00:00 2001
From: Massimiliano Culpo <massimiliano.culpo@gmail.com>
Date: Fri, 1 Nov 2024 21:43:16 +0100
Subject: [PATCH] Fix pickle round-trip of specs propagating variants (#47351)

This changes `Spec` serialization to include information about propagation for abstract specs.
This was previously not included in the JSON representation for abstract specs, and couldn't be
stored.

Now, there is a separate `propagate` dictionary alongside the `parameters` dictionary. This isn't
beautiful, but when we bump the spec version for Spack `v0.24`, we can clean up this and other
aspects of the schema.
---
 lib/spack/spack/spec.py           | 30 ++++++++++++++++++++----------
 lib/spack/spack/test/spec_yaml.py | 24 ++++++++++++++++++++++++
 lib/spack/spack/variant.py        | 11 +++++++----
 3 files changed, 51 insertions(+), 14 deletions(-)

diff --git a/lib/spack/spack/spec.py b/lib/spack/spack/spec.py
index 055be1508b3..0159c33fa90 100644
--- a/lib/spack/spack/spec.py
+++ b/lib/spack/spack/spec.py
@@ -2045,6 +2045,18 @@ def to_node_dict(self, hash=ht.dag_hash):
         if params:
             d["parameters"] = params
 
+        if params and not self.concrete:
+            flag_names = [
+                name
+                for name, flags in self.compiler_flags.items()
+                if any(x.propagate for x in flags)
+            ]
+            d["propagate"] = sorted(
+                itertools.chain(
+                    [v.name for v in self.variants.values() if v.propagate], flag_names
+                )
+            )
+
         if self.external:
             d["external"] = syaml.syaml_dict(
                 [
@@ -2217,16 +2229,10 @@ def node_dict_with_hashes(self, hash=ht.dag_hash):
         spec is concrete, the full hash is added as well.  If 'build' is in
         the hash_type, the build hash is also added."""
         node = self.to_node_dict(hash)
+        # All specs have at least a DAG hash
         node[ht.dag_hash.name] = self.dag_hash()
 
-        # dag_hash is lazily computed -- but if we write a spec out, we want it
-        # to be included. This is effectively the last chance we get to compute
-        # it accurately.
-        if self.concrete:
-            # all specs have at least a DAG hash
-            node[ht.dag_hash.name] = self.dag_hash()
-
-        else:
+        if not self.concrete:
             node["concrete"] = False
 
         # we can also give them other hash types if we want
@@ -4999,13 +5005,17 @@ def from_node_dict(cls, node):
         else:
             spec.compiler = None
 
+        propagated_names = node.get("propagate", [])
         for name, values in node.get("parameters", {}).items():
+            propagate = name in propagated_names
             if name in _valid_compiler_flags:
                 spec.compiler_flags[name] = []
                 for val in values:
-                    spec.compiler_flags.add_flag(name, val, False)
+                    spec.compiler_flags.add_flag(name, val, propagate)
             else:
-                spec.variants[name] = vt.MultiValuedVariant.from_node_dict(name, values)
+                spec.variants[name] = vt.MultiValuedVariant.from_node_dict(
+                    name, values, propagate=propagate
+                )
 
         spec.external_path = None
         spec.external_modules = None
diff --git a/lib/spack/spack/test/spec_yaml.py b/lib/spack/spack/test/spec_yaml.py
index 3f20e5626ee..bdc4e46ad79 100644
--- a/lib/spack/spack/test/spec_yaml.py
+++ b/lib/spack/spack/test/spec_yaml.py
@@ -16,6 +16,7 @@
 import io
 import json
 import os
+import pickle
 
 import pytest
 import ruamel.yaml
@@ -554,3 +555,26 @@ def test_anchorify_2():
 e: *id002
 """
     )
+
+
+@pytest.mark.parametrize(
+    "spec_str",
+    [
+        "hdf5 ++mpi",
+        "hdf5 cflags==-g",
+        "hdf5 foo==bar",
+        "hdf5~~mpi++shared",
+        "hdf5 cflags==-g foo==bar cxxflags==-O3",
+        "hdf5 cflags=-g foo==bar cxxflags==-O3",
+    ],
+)
+def test_pickle_roundtrip_for_abstract_specs(spec_str):
+    """Tests that abstract specs correctly round trip when pickled.
+
+    This test compares both spec objects and their string representation, due to some
+    inconsistencies in how `Spec.__eq__` is implemented.
+    """
+    s = spack.spec.Spec(spec_str)
+    t = pickle.loads(pickle.dumps(s))
+    assert s == t
+    assert str(s) == str(t)
diff --git a/lib/spack/spack/variant.py b/lib/spack/spack/variant.py
index bb4138116fe..7d3468a7ded 100644
--- a/lib/spack/spack/variant.py
+++ b/lib/spack/spack/variant.py
@@ -12,6 +12,7 @@
 import io
 import itertools
 import re
+from typing import List, Union
 
 import llnl.util.lang as lang
 import llnl.util.tty.color
@@ -255,19 +256,21 @@ def __init__(self, name, value, propagate=False):
         self.value = value
 
     @staticmethod
-    def from_node_dict(name, value):
+    def from_node_dict(
+        name: str, value: Union[str, List[str]], *, propagate: bool = False
+    ) -> "AbstractVariant":
         """Reconstruct a variant from a node dict."""
         if isinstance(value, list):
             # read multi-value variants in and be faithful to the YAML
-            mvar = MultiValuedVariant(name, ())
+            mvar = MultiValuedVariant(name, (), propagate=propagate)
             mvar._value = tuple(value)
             mvar._original_value = mvar._value
             return mvar
 
         elif str(value).upper() == "TRUE" or str(value).upper() == "FALSE":
-            return BoolValuedVariant(name, value)
+            return BoolValuedVariant(name, value, propagate=propagate)
 
-        return SingleValuedVariant(name, value)
+        return SingleValuedVariant(name, value, propagate=propagate)
 
     def yaml_entry(self):
         """Returns a key, value tuple suitable to be an entry in a yaml dict.