Add DerivationTree.packages() -> HashSet<&P> (pubgrub-rs#219)

* Add `DerivationTree.packages() -> HashSet<&P>` * Review * Add `.packages()` to a test --------- Co-authored-by: Zanie Blue <[email protected]>
astral-sh · May 10, 2024 · 89f9dff · 89f9dff
1 parent fce8e02
commit 89f9dff
Show file tree

Hide file tree

Showing 2 changed files with 41 additions and 3 deletions.
diff --git a/src/report.rs b/src/report.rs
@@ -9,7 +9,7 @@ use std::sync::Arc;
 
 use crate::package::Package;
 use crate::term::Term;
-use crate::type_aliases::Map;
+use crate::type_aliases::{Map, Set};
 use crate::version_set::VersionSet;
 
 /// Reporter trait.
@@ -71,6 +71,32 @@ pub struct Derived<P: Package, VS: VersionSet, M: Eq + Clone + Debug + Display>
 }
 
 impl<P: Package, VS: VersionSet, M: Eq + Clone + Debug + Display> DerivationTree<P, VS, M> {
+    /// Get all packages referred to in the derivation tree.
+    pub fn packages(&self) -> Set<&P> {
+        let mut packages = Set::default();
+        match self {
+            Self::External(external) => match external {
+                External::FromDependencyOf(p, _, p2, _) => {
+                    packages.insert(p);
+                    packages.insert(p2);
+                }
+                External::NoVersions(p, _)
+                | External::NotRoot(p, _)
+                | External::Custom(p, _, _) => {
+                    packages.insert(p);
+                }
+            },
+            Self::Derived(derived) => {
+                // Less efficient than recursing with a `&mut Set<&P>`, but it's sufficient for
+                // small to medium-sized inputs such as a single `DerivationTree`.
+                packages.extend(derived.terms.keys());
+                packages.extend(derived.cause1.packages().iter());
+                packages.extend(derived.cause2.packages().iter());
+            }
+        }
+        packages
+    }
+
     /// Merge the [NoVersions](External::NoVersions) external incompatibilities
     /// with the other one they are matched with
     /// in a derived incompatibility.

diff --git a/tests/examples.rs b/tests/examples.rs
@@ -4,7 +4,7 @@ use pubgrub::error::PubGrubError;
 use pubgrub::range::Range;
 use pubgrub::report::{DefaultStringReporter, Reporter as _};
 use pubgrub::solver::{resolve, OfflineDependencyProvider};
-use pubgrub::type_aliases::Map;
+use pubgrub::type_aliases::{Map, Set};
 use pubgrub::version::SemanticVersion;
 
 type NumVS = Range<u32>;
@@ -217,13 +217,20 @@ fn confusing_with_lots_of_holes() {
     let mut dependency_provider = OfflineDependencyProvider::<&str, NumVS>::new();
 
     // root depends on foo...
-    dependency_provider.add_dependencies("root", 1u32, vec![("foo", Range::full())]);
+    dependency_provider.add_dependencies(
+        "root",
+        1u32,
+        vec![("foo", Range::full()), ("baz", Range::full())],
+    );
 
     for i in 1..6 {
         // foo depends on bar...
         dependency_provider.add_dependencies("foo", i as u32, vec![("bar", Range::full())]);
     }
 
+    // This package is part of the dependency tree, but it's not part of the conflict
+    dependency_provider.add_dependencies("baz", 1u32, vec![]);
+
     let Err(PubGrubError::NoSolution(mut derivation_tree)) =
         resolve(&dependency_provider, "root", 1u32)
     else {
@@ -239,4 +246,9 @@ And because there is no version of foo in <1 | >1, <2 | >2, <3 | >3, <4 | >4, <5
         &DefaultStringReporter::report(&derivation_tree),
         "Because foo depends on bar and root 1 depends on foo, root 1 is forbidden."
     );
+    assert_eq!(
+        derivation_tree.packages(),
+        // baz isn't shown.
+        Set::from_iter(&["root", "foo", "bar"])
+    );
 }