Add PubGrub's priority queue (#221)

Pulls in pubgrub-rs/pubgrub#104.

crates/puffin-resolver/src/pubgrub/mod.rs

@@ -10,6 +10,7 @@ use crate::pubgrub::specifier::PubGrubSpecifier;
 use crate::pubgrub::version::{PubGrubVersion, MAX_VERSION};
 
 pub(crate) mod package;
+pub(crate) mod priority;
 mod specifier;
 pub(crate) mod version;
 

crates/puffin-resolver/src/pubgrub/priority.rs

@@ -0,0 +1,3 @@
+use std::cmp::Reverse;
+
+pub(crate) type PubGrubPriority = Reverse<usize>;

crates/puffin-resolver/src/resolution.rs

@@ -1,8 +1,8 @@
+use std::hash::BuildHasherDefault;
+
 use colored::Colorize;
 use fxhash::FxHashMap;
 use petgraph::visit::EdgeRef;
-use std::hash::BuildHasherDefault;
-
 use pubgrub::range::Range;
 use pubgrub::solver::{Kind, State};
 use pubgrub::type_aliases::SelectedDependencies;
@@ -13,6 +13,7 @@ use puffin_client::File;
 use puffin_package::package_name::PackageName;
 
 use crate::pubgrub::package::PubGrubPackage;
+use crate::pubgrub::priority::PubGrubPriority;
 use crate::pubgrub::version::PubGrubVersion;
 
 /// A package pinned at a specific version.
@@ -95,7 +96,7 @@ impl Graph {
     pub fn from_state(
         selection: &SelectedDependencies<PubGrubPackage, PubGrubVersion>,
         pins: &FxHashMap<PackageName, FxHashMap<Version, File>>,
-        state: &State<PubGrubPackage, Range<PubGrubVersion>>,
+        state: &State<PubGrubPackage, Range<PubGrubVersion>, PubGrubPriority>,
     ) -> Self {
         // TODO(charlie): petgraph is a really heavy and unnecessary dependency here. We should
         // write our own graph, given that our requirements are so simple.

crates/puffin-resolver/src/resolver.rs

@@ -1,6 +1,6 @@
 //! Given a set of requirements, find a set of compatible packages.
 
-use std::borrow::Borrow;
+use std::cmp::Reverse;
 use std::collections::hash_map::Entry;
 use std::collections::BTreeMap;
 use std::future::Future;
@@ -34,6 +34,7 @@ use crate::distribution::{DistributionFile, SdistFile, WheelFile};
 use crate::error::ResolveError;
 use crate::manifest::Manifest;
 use crate::pubgrub::package::PubGrubPackage;
+use crate::pubgrub::priority::PubGrubPriority;
 use crate::pubgrub::version::{PubGrubVersion, MIN_VERSION};
 use crate::pubgrub::{iter_requirements, version_range};
 use crate::resolution::Graph;
@@ -130,32 +131,41 @@ impl<'a, Context: BuildContext + Sync> Resolver<'a, Context> {
             // Run unit propagation.
             state.unit_propagation(next)?;
 
-            // Fetch the list of candidates.
-            let Some(potential_packages) = state.partial_solution.potential_packages() else {
-                let Some(selection) = state.partial_solution.extract_solution() else {
-                    return Err(PubGrubError::Failure(
-                        "How did we end up with no package to choose but no solution?".into(),
-                    )
-                    .into());
-                };
+            // Pre-visit all candidate packages, to allow metadata to be fetched in parallel.
+            self.pre_visit(
+                state.partial_solution.prioritized_packages(),
+                &mut requested_versions,
+                request_sink,
+            )?;
 
+            // Choose a package version.
+            let Some(highest_priority_pkg) = state
+                .partial_solution
+                .pick_highest_priority_pkg(|package, range| self.prioritize(package, range))
+            else {
+                let selection = state.partial_solution.extract_solution();
                 return Ok(Graph::from_state(&selection, &pins, &state));
             };
+            next = highest_priority_pkg;
 
-            // Choose a package version.
-            let potential_packages = potential_packages.collect::<Vec<_>>();
+            let term_intersection = state
+                .partial_solution
+                .term_intersection_for_package(&next)
+                .ok_or_else(|| {
+                    PubGrubError::Failure("a package was chosen but we don't have a term.".into())
+                })?;
             let decision = self
-                .choose_package_version(
-                    potential_packages,
+                .choose_version(
+                    &next,
+                    term_intersection.unwrap_positive(),
                     &mut pins,
                     &mut requested_versions,
                     request_sink,
                 )
                 .await?;
-            next = decision.0.clone();
 
             // Pick the next compatible version.
-            let version = match decision.1 {
+            let version = match decision {
                 None => {
                     debug!("No compatible version found for: {}", next);
 
@@ -226,19 +236,28 @@ impl<'a, Context: BuildContext + Sync> Resolver<'a, Context> {
         }
     }
 
-    /// Given a set of candidate packages, choose the next package (and version) to add to the
-    /// partial solution.
-    async fn choose_package_version<T: Borrow<PubGrubPackage>, U: Borrow<Range<PubGrubVersion>>>(
+    #[allow(clippy::unused_self)]
+    fn prioritize(
         &self,
-        mut potential_packages: Vec<(T, U)>,
-        pins: &mut FxHashMap<PackageName, FxHashMap<pep440_rs::Version, File>>,
+        _package: &PubGrubPackage,
+        _range: &Range<PubGrubVersion>,
+    ) -> PubGrubPriority {
+        // TODO(charlie): Define a priority function.
+        Reverse(0)
+    }
+
+    /// Visit the set of candidate packages prior to selection. This allows us to fetch metadata for
+    /// all of the packages in parallel.
+    fn pre_visit(
+        &self,
+        packages: impl Iterator<Item = (&'a PubGrubPackage, &'a Range<PubGrubVersion>)>,
         in_flight: &mut FxHashSet<String>,
         request_sink: &futures::channel::mpsc::UnboundedSender<Request>,
-    ) -> Result<(T, Option<PubGrubVersion>), ResolveError> {
+    ) -> Result<(), ResolveError> {
         // Iterate over the potential packages, and fetch file metadata for any of them. These
         // represent our current best guesses for the versions that we _might_ select.
-        for (index, (package, range)) in potential_packages.iter().enumerate() {
-            let PubGrubPackage::Package(package_name, _) = package.borrow() else {
+        for (package, range) in packages {
+            let PubGrubPackage::Package(package_name, _) = package else {
                 continue;
             };
 
@@ -250,13 +269,9 @@ impl<'a, Context: BuildContext + Sync> Resolver<'a, Context> {
 
             // Try to find a compatible version. If there aren't any compatible versions,
             // short-circuit and return `None`.
-            let Some(candidate) = self
-                .selector
-                .select(package_name, range.borrow(), version_map)
-            else {
-                // Short circuit: we couldn't find _any_ compatible versions for a package.
-                let (package, _range) = potential_packages.swap_remove(index);
-                return Ok((package, None));
+            let Some(candidate) = self.selector.select(package_name, range, version_map) else {
+                // Short-circuit: we couldn't find _any_ compatible versions for a package.
+                return Ok(());
             };
 
             // Emit a request to fetch the metadata for this version.
@@ -277,32 +292,35 @@ impl<'a, Context: BuildContext + Sync> Resolver<'a, Context> {
                 }
             }
         }
+        Ok(())
+    }
 
-        // Always choose the first package.
-        // TODO(charlie): Devise a better strategy here (for example: always choose the package with
-        // the fewest versions).
-        let (package, range) = potential_packages.swap_remove(0);
-
-        return match package.borrow() {
-            PubGrubPackage::Root => Ok((package, Some(MIN_VERSION.clone()))),
+    /// Given a set of candidate packages, choose the next package (and version) to add to the
+    /// partial solution.
+    async fn choose_version(
+        &self,
+        package: &PubGrubPackage,
+        range: &Range<PubGrubVersion>,
+        pins: &mut FxHashMap<PackageName, FxHashMap<pep440_rs::Version, File>>,
+        in_flight: &mut FxHashSet<String>,
+        request_sink: &futures::channel::mpsc::UnboundedSender<Request>,
+    ) -> Result<Option<PubGrubVersion>, ResolveError> {
+        return match package {
+            PubGrubPackage::Root => Ok(Some(MIN_VERSION.clone())),
             PubGrubPackage::Package(package_name, _) => {
                 // Wait for the metadata to be available.
                 let entry = self.index.packages.wait(package_name).await.unwrap();
                 let version_map = entry.value();
 
                 debug!(
                     "Searching for a compatible version of {} ({})",
-                    package_name,
-                    range.borrow(),
+                    package_name, range,
                 );
 
                 // Find a compatible version.
-                let Some(candidate) =
-                    self.selector
-                        .select(package_name, range.borrow(), version_map)
-                else {
+                let Some(candidate) = self.selector.select(package_name, range, version_map) else {
                     // Short circuit: we couldn't find _any_ compatible versions for a package.
-                    return Ok((package, None));
+                    return Ok(None);
                 };
 
                 debug!(
@@ -340,7 +358,7 @@ impl<'a, Context: BuildContext + Sync> Resolver<'a, Context> {
                 }
 
                 let version = candidate.version.clone();
-                Ok((package, Some(version)))
+                Ok(Some(version))
             }
         };
     }

scripts/benchmarks/requirements.in

@@ -6,25 +6,3 @@ pygls>=1.0.1
 lsprotocol>=2023.0.0a1
 ruff>=0.0.274
 typing_extensions
-scipy
-numpy
-pandas<2.0.0
-matplotlib>=3.0.0
-scikit-learn
-rich
-textual
-jupyter>=1.0.0,<2.0.0
-transformers[torch]
-django<4.0.0
-sqlalchemy
-psycopg2-binary
-trio<0.20
-trio-websocket
-trio-asyncio
-trio-typing
-trio-protocol
-fastapi
-typer
-pydantic
-uvicorn
-traitlets

vendor/pubgrub/Cargo.toml

@@ -22,6 +22,8 @@ include = ["Cargo.toml", "LICENSE", "README.md", "src/**", "tests/**", "examples
 [dependencies]
 thiserror = "1.0"
 rustc-hash = "1.1.0"
+indexmap = "2.0.2"
+priority-queue = "1.1.1"
 serde = { version = "1.0", features = ["derive"], optional = true }
 log = "0.4.14" # for debug logs in tests
 

vendor/pubgrub/examples/caching_dependency_provider.rs

@@ -32,13 +32,6 @@ impl<P: Package, VS: VersionSet, DP: DependencyProvider<P, VS>>
 impl<P: Package, VS: VersionSet, DP: DependencyProvider<P, VS>> DependencyProvider<P, VS>
     for CachingDependencyProvider<P, VS, DP>
 {
-    fn choose_package_version<T: std::borrow::Borrow<P>, U: std::borrow::Borrow<VS>>(
-        &self,
-        packages: impl Iterator<Item = (T, U)>,
-    ) -> Result<(T, Option<VS::V>), Box<dyn Error + Send + Sync>> {
-        self.remote_dependencies.choose_package_version(packages)
-    }
-
     // Caches dependencies if they were already queried
     fn get_dependencies(
         &self,
@@ -66,6 +59,20 @@ impl<P: Package, VS: VersionSet, DP: DependencyProvider<P, VS>> DependencyProvid
             error @ Err(_) => error,
         }
     }
+
+    fn choose_version(
+        &self,
+        package: &P,
+        range: &VS,
+    ) -> Result<Option<VS::V>, Box<dyn Error + Send + Sync>> {
+        self.remote_dependencies.choose_version(package, range)
+    }
+
+    type Priority = DP::Priority;
+
+    fn prioritize(&self, package: &P, range: &VS) -> Self::Priority {
+        self.remote_dependencies.prioritize(package, range)
+    }
 }
 
 fn main() {

vendor/pubgrub/src/error.rs

@@ -46,7 +46,7 @@ pub enum PubGrubError<P: Package, VS: VersionSet> {
     /// Error arising when the implementer of
     /// [DependencyProvider](crate::solver::DependencyProvider)
     /// returned an error in the method
-    /// [choose_package_version](crate::solver::DependencyProvider::choose_package_version).
+    /// [choose_version](crate::solver::DependencyProvider::choose_version).
     #[error("Decision making failed")]
     ErrorChoosingPackageVersion(Box<dyn std::error::Error + Send + Sync>),
 

vendor/pubgrub/src/internal/core.rs

@@ -20,7 +20,7 @@ use crate::version_set::VersionSet;
 
 /// Current state of the PubGrub algorithm.
 #[derive(Clone)]
-pub struct State<P: Package, VS: VersionSet> {
+pub struct State<P: Package, VS: VersionSet, Priority: Ord + Clone> {
     root_package: P,
     root_version: VS::V,
 
@@ -32,7 +32,7 @@ pub struct State<P: Package, VS: VersionSet> {
 
     /// Partial solution.
     /// TODO: remove pub.
-    pub partial_solution: PartialSolution<P, VS>,
+    pub partial_solution: PartialSolution<P, VS, Priority>,
 
     /// The store is the reference storage for all incompatibilities.
     pub incompatibility_store: Arena<Incompatibility<P, VS>>,
@@ -43,7 +43,7 @@ pub struct State<P: Package, VS: VersionSet> {
     unit_propagation_buffer: SmallVec<P>,
 }
 
-impl<P: Package, VS: VersionSet> State<P, VS> {
+impl<P: Package, VS: VersionSet, Priority: Ord + Clone> State<P, VS, Priority> {
     /// Initialization of PubGrub state.
     pub fn init(root_package: P, root_version: VS::V) -> Self {
         let mut incompatibility_store = Arena::new();