parse_specs: special case for concretizing lookups quickly (#47556)

We added unification semantics for parsing specs from the CLI, but there are a couple
of special cases in which we can avoid calls to the concretizer for speed when the
specs can all be resolved by lookups.

- [x] special case 1: solving a single spec

- [x] special case 2: all specs are either concrete (come from a file) or have an abstract
      hash. In this case if concretizer:unify:true we need an additional check to confirm
      the specs are compatible.

- [x] add a parameterized test for unifying on the CI

---------

Signed-off-by: Todd Gamblin <tgamblin@llnl.gov>
Co-authored-by: Todd Gamblin <tgamblin@llnl.gov>

4 files changed, 182 insertions, 3 deletions

diff --git a/lib/spack/spack/cmd/__init__.py b/lib/spack/spack/cmd/__init__.py
index c0efd52521..9449face85 100644
--- a/lib/spack/spack/cmd/__init__.py
+++ b/lib/spack/spack/cmd/__init__.py
@@ -8,6 +8,7 @@ import importlib
 import os
 import re
 import sys
+from collections import Counter
 from typing import List, Union
 
 import llnl.string
@@ -189,6 +190,43 @@ def _concretize_spec_pairs(to_concretize, tests=False):
     rules from config."""
     unify = spack.config.get("concretizer:unify", False)
 
+    # Special case for concretizing a single spec
+    if len(to_concretize) == 1:
+        abstract, concrete = to_concretize[0]
+        return [concrete or abstract.concretized()]
+
+    # Special case if every spec is either concrete or has an abstract hash
+    if all(
+        concrete or abstract.concrete or abstract.abstract_hash
+        for abstract, concrete in to_concretize
+    ):
+        # Get all the concrete specs
+        ret = [
+            concrete or (abstract if abstract.concrete else abstract.lookup_hash())
+            for abstract, concrete in to_concretize
+        ]
+
+        # If unify: true, check that specs don't conflict
+        # Since all concrete, "when_possible" is not relevant
+        if unify is True:  # True, "when_possible", False are possible values
+            runtimes = spack.repo.PATH.packages_with_tags("runtime")
+            specs_per_name = Counter(
+                spec.name
+                for spec in traverse.traverse_nodes(
+                    ret, deptype=("link", "run"), key=traverse.by_dag_hash
+                )
+                if spec.name not in runtimes  # runtimes are allowed multiple times
+            )
+
+            conflicts = sorted(name for name, count in specs_per_name.items() if count > 1)
+            if conflicts:
+                raise spack.error.SpecError(
+                    "Specs conflict and `concretizer:unify` is configured true.",
+                    f"    specs depend on multiple versions of {', '.join(conflicts)}",
+                )
+        return ret
+
+    # Standard case
     concretize_method = spack.concretize.concretize_separately  # unify: false
     if unify is True:
         concretize_method = spack.concretize.concretize_together
diff --git a/lib/spack/spack/spec.py b/lib/spack/spack/spec.py
index 3d92879484..d87296a3fb 100644
--- a/lib/spack/spack/spec.py
+++ b/lib/spack/spack/spec.py
@@ -59,7 +59,7 @@ import platform
 import re
 import socket
 import warnings
-from typing import Any, Callable, Dict, List, Match, Optional, Set, Tuple, Union
+from typing import Any, Callable, Dict, Iterable, List, Match, Optional, Set, Tuple, Union
 
 import archspec.cpu
 
@@ -2828,7 +2828,7 @@ class Spec:
             msg += "    For each package listed, choose another spec\n"
             raise SpecDeprecatedError(msg)
 
-    def concretize(self, tests: Union[bool, List[str]] = False) -> None:
+    def concretize(self, tests: Union[bool, Iterable[str]] = False) -> None:
         """Concretize the current spec.
 
         Args:
@@ -2956,7 +2956,7 @@ class Spec:
         for spec in self.traverse():
             spec._cached_hash(ht.dag_hash)
 
-    def concretized(self, tests=False):
+    def concretized(self, tests: Union[bool, Iterable[str]] = False) -> "spack.spec.Spec":
         """This is a non-destructive version of concretize().
 
         First clones, then returns a concrete version of this package
diff --git a/lib/spack/spack/test/cmd/init_py_functions.py b/lib/spack/spack/test/cmd/init_py_functions.py
index 4dc000edb9..deb6222411 100644
--- a/lib/spack/spack/test/cmd/init_py_functions.py
+++ b/lib/spack/spack/test/cmd/init_py_functions.py
@@ -4,10 +4,15 @@
 # SPDX-License-Identifier: (Apache-2.0 OR MIT)
 import pytest
 
+import spack.environment as ev
+import spack.error
+import spack.solver.asp as asp
 from spack.cmd import (
     CommandNameError,
     PythonNameError,
     cmd_name,
+    matching_specs_from_env,
+    parse_specs,
     python_name,
     require_cmd_name,
     require_python_name,
@@ -34,3 +39,99 @@ def test_require_cmd_name():
     with pytest.raises(CommandNameError):
         require_cmd_name("okey_dokey")
     require_cmd_name(cmd_name("okey_dokey"))
+
+
+@pytest.mark.parametrize(
+    "unify,spec_strs,error",
+    [
+        # single spec
+        (True, ["zmpi"], None),
+        (False, ["mpileaks"], None),
+        # multiple specs, some from hash some from file
+        (True, ["zmpi", "mpileaks^zmpi", "libelf"], None),
+        (True, ["mpileaks^zmpi", "mpileaks^mpich", "libelf"], spack.error.SpecError),
+        (False, ["mpileaks^zmpi", "mpileaks^mpich", "libelf"], None),
+    ],
+)
+def test_special_cases_concretization_parse_specs(
+    unify, spec_strs, error, monkeypatch, mutable_config, mutable_database, tmpdir
+):
+    """Test that special cases in parse_specs(concretize=True) bypass solver"""
+
+    # monkeypatch to ensure we do not call the actual concretizer
+    def _fail(*args, **kwargs):
+        assert False
+
+    monkeypatch.setattr(asp.SpackSolverSetup, "setup", _fail)
+
+    spack.config.set("concretizer:unify", unify)
+
+    args = [f"/{spack.store.STORE.db.query(s)[0].dag_hash()}" for s in spec_strs]
+    if len(args) > 1:
+        # We convert the last one to a specfile input
+        filename = tmpdir.join("spec.json")
+        spec = parse_specs(args[-1], concretize=True)[0]
+        with open(filename, "w") as f:
+            spec.to_json(f)
+        args[-1] = str(filename)
+
+    if error:
+        with pytest.raises(error):
+            parse_specs(args, concretize=True)
+    else:
+        # assertion error from monkeypatch above if test fails
+        parse_specs(args, concretize=True)
+
+
+@pytest.mark.parametrize(
+    "unify,spec_strs,error",
+    [
+        # single spec
+        (True, ["zmpi"], None),
+        (False, ["mpileaks"], None),
+        # multiple specs, some from hash some from file
+        (True, ["zmpi", "mpileaks^zmpi", "libelf"], None),
+        (True, ["mpileaks^zmpi", "mpileaks^mpich", "libelf"], spack.error.SpecError),
+        (False, ["mpileaks^zmpi", "mpileaks^mpich", "libelf"], None),
+    ],
+)
+def test_special_cases_concretization_matching_specs_from_env(
+    unify,
+    spec_strs,
+    error,
+    monkeypatch,
+    mutable_config,
+    mutable_database,
+    tmpdir,
+    mutable_mock_env_path,
+):
+    """Test that special cases in parse_specs(concretize=True) bypass solver"""
+
+    # monkeypatch to ensure we do not call the actual concretizer
+    def _fail(*args, **kwargs):
+        assert False
+
+    monkeypatch.setattr(asp.SpackSolverSetup, "setup", _fail)
+
+    spack.config.set("concretizer:unify", unify)
+
+    ev.create("test")
+    env = ev.read("test")
+
+    args = [f"/{spack.store.STORE.db.query(s)[0].dag_hash()}" for s in spec_strs]
+    if len(args) > 1:
+        # We convert the last one to a specfile input
+        filename = tmpdir.join("spec.json")
+        spec = parse_specs(args[-1], concretize=True)[0]
+        with open(filename, "w") as f:
+            spec.to_json(f)
+        args[-1] = str(filename)
+
+    with env:
+        specs = parse_specs(args, concretize=False)
+        if error:
+            with pytest.raises(error):
+                matching_specs_from_env(specs)
+        else:
+            # assertion error from monkeypatch above if test fails
+            matching_specs_from_env(specs)
diff --git a/lib/spack/spack/test/cmd/spec.py b/lib/spack/spack/test/cmd/spec.py
index a57c40ec92..1d0d08f494 100644
--- a/lib/spack/spack/test/cmd/spec.py
+++ b/lib/spack/spack/test/cmd/spec.py
@@ -179,3 +179,43 @@ def test_spec_version_assigned_git_ref_as_version(name, version, error):
     else:
         output = spec(name + "@" + version)
         assert version in output
+
+
+@pytest.mark.parametrize(
+    "unify, spec_hash_args, match, error",
+    [
+        # success cases with unfiy:true
+        (True, ["mpileaks_mpich"], "mpich", None),
+        (True, ["mpileaks_zmpi"], "zmpi", None),
+        (True, ["mpileaks_mpich", "dyninst"], "mpich", None),
+        (True, ["mpileaks_zmpi", "dyninst"], "zmpi", None),
+        # same success cases with unfiy:false
+        (False, ["mpileaks_mpich"], "mpich", None),
+        (False, ["mpileaks_zmpi"], "zmpi", None),
+        (False, ["mpileaks_mpich", "dyninst"], "mpich", None),
+        (False, ["mpileaks_zmpi", "dyninst"], "zmpi", None),
+        # cases with unfiy:false
+        (True, ["mpileaks_mpich", "mpileaks_zmpi"], "callpath, mpileaks", spack.error.SpecError),
+        (False, ["mpileaks_mpich", "mpileaks_zmpi"], "zmpi", None),
+    ],
+)
+def test_spec_unification_from_cli(
+    install_mockery, mutable_config, mutable_database, unify, spec_hash_args, match, error
+):
+    """Ensure specs grouped together on the CLI are concretized together when unify:true."""
+    spack.config.set("concretizer:unify", unify)
+
+    db = spack.store.STORE.db
+    spec_lookup = {
+        "mpileaks_mpich": db.query_one("mpileaks ^mpich").dag_hash(),
+        "mpileaks_zmpi": db.query_one("mpileaks ^zmpi").dag_hash(),
+        "dyninst": db.query_one("dyninst").dag_hash(),
+    }
+
+    hashes = [f"/{spec_lookup[name]}" for name in spec_hash_args]
+    if error:
+        with pytest.raises(error, match=match):
+            output = spec(*hashes)
+    else:
+        output = spec(*hashes)
+        assert match in output