index: avoid quadratic complexity through bulk update (spack#48771)

Refactors `ProviderIndex` and `TagIndex` to support bulk updates.

Previously, updating the index with N packages involved iterating
through the entire index for each package to remove stale entries,
resulting in O(N * IndexSize) complexity.

This change splits the update process into two steps:

1. Bulk removal of stale entries for all updated packages.
2. Bulk insertion of new entries.

This reduces the complexity of full index generation from quadratic to linear.

Signed-off-by: Harmen Stoppels <me@harmenstoppels.nl>

Author: haampie

lib/spack/spack/patch.py

@@ -6,7 +6,7 @@
 import os
 import pathlib
 import sys
-from typing import TYPE_CHECKING, Any, Dict, Optional, Tuple, Type, Union
+from typing import TYPE_CHECKING, Any, Dict, Optional, Set, Tuple, Type, Union
 
 import spack
 import spack.error
@@ -455,36 +455,38 @@ def patch_for_package(self, sha256: str, pkg: "spack.package_base.PackageBase")
         patch_dict["sha256"] = sha256
         return from_dict(patch_dict, repository=self.repository)
 
-    def update_package(self, pkg_fullname: str) -> None:
+    def update_packages(self, pkgs_fullname: Set[str]) -> None:
         """Update the patch cache.
 
         Args:
             pkg_fullname: package to update.
         """
         # remove this package from any patch entries that reference it.
-        empty = []
-        for sha256, package_to_patch in self.index.items():
-            remove = []
-            for fullname, patch_dict in package_to_patch.items():
-                if patch_dict["owner"] == pkg_fullname:
-                    remove.append(fullname)
+        if self.index:
+            empty = []
+            for sha256, package_to_patch in self.index.items():
+                remove = []
+                for fullname, patch_dict in package_to_patch.items():
+                    if patch_dict["owner"] in pkgs_fullname:
+                        remove.append(fullname)
 
-            for fullname in remove:
-                package_to_patch.pop(fullname)
+                for fullname in remove:
+                    package_to_patch.pop(fullname)
 
-            if not package_to_patch:
-                empty.append(sha256)
+                if not package_to_patch:
+                    empty.append(sha256)
 
-        # remove any entries that are now empty
-        for sha256 in empty:
-            del self.index[sha256]
+            # remove any entries that are now empty
+            for sha256 in empty:
+                del self.index[sha256]
 
         # update the index with per-package patch indexes
-        pkg_cls = self.repository.get_pkg_class(pkg_fullname)
-        partial_index = self._index_patches(pkg_cls, self.repository)
-        for sha256, package_to_patch in partial_index.items():
-            p2p = self.index.setdefault(sha256, {})
-            p2p.update(package_to_patch)
+        for pkg_fullname in pkgs_fullname:
+            pkg_cls = self.repository.get_pkg_class(pkg_fullname)
+            partial_index = self._index_patches(pkg_cls, self.repository)
+            for sha256, package_to_patch in partial_index.items():
+                p2p = self.index.setdefault(sha256, {})
+                p2p.update(package_to_patch)
 
     def update(self, other: "PatchCache") -> None:
         """Update this cache with the contents of another.

lib/spack/spack/provider_index.py

@@ -53,18 +53,8 @@ def __init__(
         self.repository = repository
         self.restrict = restrict
         self.providers = {}
-
-        specs = specs or []
-        for spec in specs:
-            if isinstance(spec, str):
-                from spack.spec import Spec
-
-                spec = Spec(spec)
-
-            if self.repository.is_virtual_safe(spec.name):
-                continue
-
-            self.update(spec)
+        if specs:
+            self.update_packages(specs)
 
     def providers_for(self, virtual: Union[str, "spack.spec.Spec"]) -> List["spack.spec.Spec"]:
         """Return a list of specs of all packages that provide virtual packages with the supplied
@@ -122,57 +112,55 @@ def __str__(self):
     def __repr__(self):
         return repr(self.providers)
 
-    def update(self, spec: Union[str, "spack.spec.Spec"]) -> None:
+    def update_packages(self, specs: Iterable[Union[str, "spack.spec.Spec"]]):
         """Update the provider index with additional virtual specs.
 
         Args:
             spec: spec potentially providing additional virtual specs
         """
-        if isinstance(spec, str):
-            from spack.spec import Spec
-
-            spec = Spec(spec)
+        from spack.spec import Spec
 
-        if not spec.name:
-            # Empty specs do not have a package
-            return
+        for spec in specs:
+            if not isinstance(spec, Spec):
+                spec = Spec(spec)
 
-        msg = "cannot update an index passing the virtual spec '{}'".format(spec.name)
-        assert not self.repository.is_virtual_safe(spec.name), msg
+            if not spec.name or self.repository.is_virtual_safe(spec.name):
+                # Only non-virtual packages with name can provide virtual specs.
+                continue
 
-        pkg_cls = self.repository.get_pkg_class(spec.name)
-        for provider_spec_readonly, provided_specs in pkg_cls.provided.items():
-            for provided_spec in provided_specs:
-                # TODO: fix this comment.
-                # We want satisfaction other than flags
-                provider_spec = provider_spec_readonly.copy()
-                provider_spec.compiler_flags = spec.compiler_flags.copy()
+            pkg_cls = self.repository.get_pkg_class(spec.name)
+            for provider_spec_readonly, provided_specs in pkg_cls.provided.items():
+                for provided_spec in provided_specs:
+                    # TODO: fix this comment.
+                    # We want satisfaction other than flags
+                    provider_spec = provider_spec_readonly.copy()
+                    provider_spec.compiler_flags = spec.compiler_flags.copy()
 
-                if spec.intersects(provider_spec, deps=False):
-                    provided_name = provided_spec.name
+                    if spec.intersects(provider_spec, deps=False):
+                        provided_name = provided_spec.name
 
-                    provider_map = self.providers.setdefault(provided_name, {})
-                    if provided_spec not in provider_map:
-                        provider_map[provided_spec] = set()
+                        provider_map = self.providers.setdefault(provided_name, {})
+                        if provided_spec not in provider_map:
+                            provider_map[provided_spec] = set()
 
-                    if self.restrict:
-                        provider_set = provider_map[provided_spec]
+                        if self.restrict:
+                            provider_set = provider_map[provided_spec]
 
-                        # If this package existed in the index before,
-                        # need to take the old versions out, as they're
-                        # now more constrained.
-                        old = set([s for s in provider_set if s.name == spec.name])
-                        provider_set.difference_update(old)
+                            # If this package existed in the index before,
+                            # need to take the old versions out, as they're
+                            # now more constrained.
+                            old = {s for s in provider_set if s.name == spec.name}
+                            provider_set.difference_update(old)
 
-                        # Now add the new version.
-                        provider_set.add(spec)
+                            # Now add the new version.
+                            provider_set.add(spec)
 
-                    else:
-                        # Before putting the spec in the map, constrain
-                        # it so that it provides what was asked for.
-                        constrained = spec.copy()
-                        constrained.constrain(provider_spec)
-                        provider_map[provided_spec].add(constrained)
+                        else:
+                            # Before putting the spec in the map, constrain
+                            # it so that it provides what was asked for.
+                            constrained = spec.copy()
+                            constrained.constrain(provider_spec)
+                            provider_map[provided_spec].add(constrained)
 
     def to_json(self, stream=None):
         """Dump a JSON representation of this object.
@@ -207,14 +195,14 @@ def merge(self, other):
 
                 spdict[provided_spec] = spdict[provided_spec].union(opdict[provided_spec])
 
-    def remove_provider(self, pkg_name):
+    def remove_providers(self, pkg_names: Set[str]):
         """Remove a provider from the ProviderIndex."""
         empty_pkg_dict = []
         for pkg, pkg_dict in self.providers.items():
             empty_pset = []
             for provided, pset in pkg_dict.items():
-                same_name = set(p for p in pset if p.fullname == pkg_name)
-                pset.difference_update(same_name)
+                to_remove = {spec for spec in pset if spec.name in pkg_names}
+                pset.difference_update(to_remove)
 
                 if not pset:
                     empty_pset.append(provided)

lib/spack/spack/repo.py

@@ -485,7 +485,7 @@ def read(self, stream):
         """Read this index from a provided file object."""
 
     @abc.abstractmethod
-    def update(self, pkg_fullname):
+    def update(self, pkgs_fullname: Set[str]):
         """Update the index in memory with information about a package."""
 
     @abc.abstractmethod
@@ -502,8 +502,8 @@ def _create(self) -> spack.tag.TagIndex:
     def read(self, stream):
         self.index = spack.tag.TagIndex.from_json(stream)
 
-    def update(self, pkg_fullname):
-        self.index.update_package(pkg_fullname.split(".")[-1], self.repository)
+    def update(self, pkgs_fullname: Set[str]):
+        self.index.update_packages({p.split(".")[-1] for p in pkgs_fullname}, self.repository)
 
     def write(self, stream):
         self.index.to_json(stream)
@@ -518,15 +518,15 @@ def _create(self) -> "spack.provider_index.ProviderIndex":
     def read(self, stream):
         self.index = spack.provider_index.ProviderIndex.from_json(stream, self.repository)
 
-    def update(self, pkg_fullname):
-        name = pkg_fullname.split(".")[-1]
+    def update(self, pkgs_fullname: Set[str]):
         is_virtual = (
-            not self.repository.exists(name) or self.repository.get_pkg_class(name).virtual
+            lambda name: not self.repository.exists(name)
+            or self.repository.get_pkg_class(name).virtual
         )
-        if is_virtual:
-            return
-        self.index.remove_provider(pkg_fullname)
-        self.index.update(pkg_fullname)
+        non_virtual_pkgs_fullname = {p for p in pkgs_fullname if not is_virtual(p.split(".")[-1])}
+        non_virtual_pkgs_names = {p.split(".")[-1] for p in non_virtual_pkgs_fullname}
+        self.index.remove_providers(non_virtual_pkgs_names)
+        self.index.update_packages(non_virtual_pkgs_fullname)
 
     def write(self, stream):
         self.index.to_json(stream)
@@ -551,8 +551,8 @@ def read(self, stream):
     def write(self, stream):
         self.index.to_json(stream)
 
-    def update(self, pkg_fullname):
-        self.index.update_package(pkg_fullname)
+    def update(self, pkgs_fullname: Set[str]):
+        self.index.update_packages(pkgs_fullname)
 
 
 class RepoIndex:
@@ -644,9 +644,7 @@ def _build_index(self, name: str, indexer: Indexer):
                 if new_index_mtime != index_mtime:
                     needs_update = self.checker.modified_since(new_index_mtime)
 
-                for pkg_name in needs_update:
-                    indexer.update(f"{self.namespace}.{pkg_name}")
-
+                indexer.update({f"{self.namespace}.{pkg_name}" for pkg_name in needs_update})
                 indexer.write(new)
 
         return indexer.index

lib/spack/spack/tag.py

@@ -2,7 +2,7 @@
 #
 # SPDX-License-Identifier: (Apache-2.0 OR MIT)
 """Classes and functions to manage package tags"""
-from typing import TYPE_CHECKING, Dict, List
+from typing import TYPE_CHECKING, Dict, List, Set
 
 import spack.error
 import spack.util.spack_json as sjson
@@ -51,26 +51,28 @@ def merge(self, other: "TagIndex") -> None:
             else:
                 self.tags[tag] = sorted({*self.tags[tag], *pkgs})
 
-    def update_package(self, pkg_name: str, repo: "spack.repo.Repo") -> None:
-        """Updates a package in the tag index.
+    def update_packages(self, pkg_names: Set[str], repo: "spack.repo.Repo") -> None:
+        """Updates packages in the tag index.
 
         Args:
-            pkg_name: name of the package to be updated
+            pkg_names: names of the packages to be updated
+            repo: the repository to get package classes from
         """
-        pkg_cls = repo.get_pkg_class(pkg_name)
-
-        # Remove the package from the list of packages, if present
+        # Remove the packages from the list of packages, if present
         for pkg_list in self.tags.values():
-            if pkg_name in pkg_list:
-                pkg_list.remove(pkg_name)
-
-        # Add it again under the appropriate tags
-        for tag in getattr(pkg_cls, "tags", []):
-            tag = tag.lower()
-            if tag not in self.tags:
-                self.tags[tag] = [pkg_cls.name]
-            else:
-                self.tags[tag].append(pkg_cls.name)
+            if pkg_names.isdisjoint(pkg_list):
+                continue
+            pkg_list[:] = [pkg for pkg in pkg_list if pkg not in pkg_names]
+
+        # Add them again under the appropriate tags
+        for pkg_name in pkg_names:
+            pkg_cls = repo.get_pkg_class(pkg_name)
+            for tag in getattr(pkg_cls, "tags", []):
+                tag = tag.lower()
+                if tag not in self.tags:
+                    self.tags[tag] = [pkg_cls.name]
+                else:
+                    self.tags[tag].append(pkg_cls.name)
 
 
 class TagIndexError(spack.error.SpackError):

lib/spack/spack/test/provider_index.py

@@ -72,3 +72,15 @@ def test_copy(mock_packages):
     p = ProviderIndex(specs=spack.repo.all_package_names(), repository=spack.repo.PATH)
     q = p.copy()
     assert p == q
+
+
+def test_remove_providers(mock_packages):
+    """Test removing providers from the index."""
+    p = ProviderIndex(specs=["mpich"], repository=spack.repo.PATH)
+    # Check that mpich is a provider for mpi
+    providers = p.providers_for("mpi")
+    assert any(spec.name == "mpich" for spec in providers)
+    p.remove_providers({"mpich"})
+    # After removal, mpich should no longer be a provider for mpi
+    providers = p.providers_for("mpi")
+    assert not any(spec.name == "mpich" for spec in providers)

lib/spack/spack/test/tag.py

@@ -146,7 +146,6 @@ def test_tag_no_tags(mock_packages):
 def test_tag_update_package(mock_packages):
     mock_index = mock_packages.tag_index
     index = spack.tag.TagIndex()
-    for name in spack.repo.all_package_names():
-        index.update_package(name, repo=mock_packages)
+    index.update_packages(set(spack.repo.all_package_names()), repo=mock_packages)
 
     ensure_tags_results_equal(mock_index.tags, index.tags)