use a priority-queue

Cargo.toml

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

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>> {
-        self.remote_dependencies.choose_package_version(packages)
-    }
-
     // Caches dependencies if they were already queried
     fn get_dependencies(
         &self,
@@ -66,6 +59,16 @@ 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>> {
+        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() {

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();

src/internal/partial_solution.rs

@@ -4,6 +4,10 @@
 //! where terms are regrouped by package in a [Map](crate::type_aliases::Map).
 
 use std::fmt::Display;
+use std::hash::BuildHasherDefault;
+
+use priority_queue::PriorityQueue;
+use rustc_hash::FxHasher;
 
 use crate::internal::arena::Arena;
 use crate::internal::incompatibility::{IncompId, Incompatibility, Relation};
@@ -15,8 +19,6 @@ use crate::version_set::VersionSet;
 
 use super::small_vec::SmallVec;
 
-use std::hash::BuildHasherDefault;
-
 type FnvIndexMap<K, V> = indexmap::IndexMap<K, V, BuildHasherDefault<rustc_hash::FxHasher>>;
 
 #[derive(Debug, Copy, Clone, Ord, PartialOrd, Eq, PartialEq)]
@@ -31,13 +33,17 @@ impl DecisionLevel {
 /// The partial solution contains all package assignments,
 /// organized by package and historically ordered.
 #[derive(Clone, Debug)]
-pub struct PartialSolution<P: Package, VS: VersionSet> {
+pub struct PartialSolution<P: Package, VS: VersionSet, Priority: Ord + Clone> {
     next_global_index: u32,
     current_decision_level: DecisionLevel,
     package_assignments: FnvIndexMap<P, PackageAssignments<P, VS>>,
+    prioritized_potential_packages: PriorityQueue<P, Priority, BuildHasherDefault<FxHasher>>,
+    just_backtracked: bool,
 }
 
-impl<P: Package, VS: VersionSet> Display for PartialSolution<P, VS> {
+impl<P: Package, VS: VersionSet, Priority: Ord + Clone> Display
+    for PartialSolution<P, VS, Priority>
+{
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         let mut assignments: Vec<_> = self
             .package_assignments
@@ -124,13 +130,15 @@ pub enum SatisfierSearch<P: Package, VS: VersionSet> {
     },
 }
 
-impl<P: Package, VS: VersionSet> PartialSolution<P, VS> {
+impl<P: Package, VS: VersionSet, Priority: Ord + Clone> PartialSolution<P, VS, Priority> {
     /// Initialize an empty PartialSolution.
     pub fn empty() -> Self {
         Self {
             next_global_index: 0,
             current_decision_level: DecisionLevel(0),
             package_assignments: FnvIndexMap::default(),
+            prioritized_potential_packages: PriorityQueue::default(),
+            just_backtracked: false,
         }
     }
 
@@ -209,22 +217,21 @@ impl<P: Package, VS: VersionSet> PartialSolution<P, VS> {
         }
     }
 
-    /// Extract potential packages for the next iteration of unit propagation.
-    /// Return `None` if there is no suitable package anymore, which stops the algorithm.
-    /// A package is a potential pick if there isn't an already
-    /// selected version (no "decision")
-    /// and if it contains at least one positive derivation term
-    /// in the partial solution.
-    pub fn potential_packages(&self) -> Option<impl Iterator<Item = (&P, &VS)>> {
-        let mut iter = (self.current_decision_level.0 as usize..self.package_assignments.len())
-            .map(move |i| self.package_assignments.get_index(i).unwrap())
+    pub fn prioritize(&mut self, prioritizer: impl Fn(&P, &VS) -> Priority) -> Option<P> {
+        let check_all = self.just_backtracked;
+        self.just_backtracked = false;
+        let current_decision_level = self.current_decision_level;
+        let package_assignments = &self.package_assignments;
+        let prioritized_potential_packages = &mut self.prioritized_potential_packages;
+        (self.current_decision_level.0 as usize..package_assignments.len())
+            .map(|i| package_assignments.get_index(i).unwrap())
+            .filter(|(_, pa)| check_all || pa.highest_decision_level == current_decision_level)
             .filter_map(|(p, pa)| pa.assignments_intersection.potential_package_filter(p))
-            .peekable();
-        if iter.peek().is_some() {
-            Some(iter)
-        } else {
-            None
-        }
+            .for_each(|(p, r)| {
+                let priority = prioritizer(&p, r);
+                prioritized_potential_packages.push(p.clone(), priority);
+            });
+        prioritized_potential_packages.pop().map(|(p, _)| p)
     }
 
     /// If a partial solution has, for every positive derivation,
@@ -287,6 +294,8 @@ impl<P: Package, VS: VersionSet> PartialSolution<P, VS> {
                 true
             }
         });
+        self.prioritized_potential_packages.clear();
+        self.just_backtracked = true;
     }
 
     /// We can add the version to the partial solution as a decision

src/lib.rs

@@ -75,7 +75,7 @@
 //! trait for our own type.
 //! Let's say that we will use [String] for packages,
 //! and [SemanticVersion](version::SemanticVersion) for versions.
-//! This may be done quite easily by implementing the two following functions.
+//! This may be done quite easily by implementing the three following functions.
 //! ```
 //! # use pubgrub::solver::{DependencyProvider, Dependencies};
 //! # use pubgrub::version::SemanticVersion;
@@ -89,7 +89,12 @@
 //! type SemVS = Range<SemanticVersion>;
 //!
 //! impl DependencyProvider<String, SemVS> for MyDependencyProvider {
-//!     fn choose_package_version<T: Borrow<String>, U: Borrow<SemVS>>(&self,packages: impl Iterator<Item=(T, U)>) -> Result<(T, Option<SemanticVersion>), Box<dyn Error>> {
+//!     fn choose_version(&self, package: &String, range: &SemVS) -> Result<Option<SemanticVersion>, Box<dyn Error>> {
+//!         unimplemented!()
+//!     }
+//!
+//!     type Priority = usize;
+//!     fn prioritize(&self, package: &String, range: &SemVS) -> Self::Priority {
 //!         unimplemented!()
 //!     }
 //!

src/solver.rs

@@ -68,7 +68,7 @@
 //! to satisfy the dependencies of that package and version pair.
 //! If there is no solution, the reason will be provided as clear as possible.
 
-use std::borrow::Borrow;
+use std::cmp::Reverse;
 use std::collections::{BTreeMap, BTreeSet as Set};
 use std::error::Error;
 
@@ -101,26 +101,29 @@ pub fn resolve<P: Package, VS: VersionSet>(
             state.partial_solution
         );
 
-        let potential_packages = state.partial_solution.potential_packages();
+        let potential_packages = state
+            .partial_solution
+            .prioritize(|p, r| dependency_provider.prioritize(p, r));
         if potential_packages.is_none() {
             drop(potential_packages);
             // The borrow checker did not like using a match on potential_packages.
             // This `if ... is_none ... drop` is a workaround.
             // I believe this is a case where Polonius could help, when and if it lands in rustc.
             return Ok(state.partial_solution.extract_solution());
         }
-        let decision = dependency_provider
-            .choose_package_version(potential_packages.unwrap())
-            .map_err(PubGrubError::ErrorChoosingPackageVersion)?;
-        log::info!("DP chose: {} @ {:?}", decision.0, decision.1);
-        next = decision.0.clone();
+        next = potential_packages.unwrap();
 
-        // Pick the next compatible version.
         let term_intersection = state
             .partial_solution
             .term_intersection_for_package(&next)
             .expect("a package was chosen but we don't have a term.");
-        let v = match decision.1 {
+        let decision = dependency_provider
+            .choose_version(&next, term_intersection.unwrap_positive())
+            .map_err(PubGrubError::ErrorChoosingPackageVersion)?;
+        log::info!("DP chose: {} @ {:?}", next, decision);
+
+        // Pick the next compatible version.
+        let v = match decision {
             None => {
                 let inc = Incompatibility::no_versions(next.clone(), term_intersection.clone());
                 state.add_incompatibility(inc);
@@ -217,36 +220,10 @@ pub enum Dependencies<P: Package, VS: VersionSet> {
 /// Trait that allows the algorithm to retrieve available packages and their dependencies.
 /// An implementor needs to be supplied to the [resolve] function.
 pub trait DependencyProvider<P: Package, VS: VersionSet> {
-    /// [Decision making](https://github.com/dart-lang/pub/blob/master/doc/solver.md#decision-making)
-    /// is the process of choosing the next package
-    /// and version that will be appended to the partial solution.
-    /// Every time such a decision must be made,
-    /// potential valid packages and sets of versions are preselected by the resolver,
-    /// and the dependency provider must choose.
-    ///
-    /// The strategy employed to choose such package and version
-    /// cannot change the existence of a solution or not,
-    /// but can drastically change the performances of the solver,
-    /// or the properties of the solution.
-    /// The documentation of Pub (PubGrub implementation for the dart programming language)
-    /// states the following:
-    ///
-    /// > Pub chooses the latest matching version of the package
-    /// > with the fewest versions that match the outstanding constraint.
-    /// > This tends to find conflicts earlier if any exist,
-    /// > since these packages will run out of versions to try more quickly.
-    /// > But there's likely room for improvement in these heuristics.
-    ///
-    /// A helper function [choose_package_with_fewest_versions] is provided to ease
-    /// implementations of this method if you can produce an iterator
-    /// of the available versions in preference order for any package.
-    ///
-    /// Note: the type `T` ensures that this returns an item from the `packages` argument.
-    #[allow(clippy::type_complexity)]
-    fn choose_package_version<T: Borrow<P>, U: Borrow<VS>>(
-        &self,
-        potential_packages: impl Iterator<Item = (T, U)>,
-    ) -> Result<(T, Option<VS::V>), Box<dyn Error>>;
+    fn choose_version(&self, package: &P, range: &VS) -> Result<Option<VS::V>, Box<dyn Error>>;
+
+    type Priority: Ord + Clone;
+    fn prioritize(&self, package: &P, range: &VS) -> Self::Priority;
 
     /// Retrieves the package dependencies.
     /// Return [Dependencies::Unknown] if its dependencies are unknown.
@@ -266,35 +243,6 @@ pub trait DependencyProvider<P: Package, VS: VersionSet> {
     }
 }
 
-/// This is a helper function to make it easy to implement
-/// [DependencyProvider::choose_package_version].
-/// It takes a function `list_available_versions` that takes a package and returns an iterator
-/// of the available versions in preference order.
-/// The helper finds the package from the `packages` argument with the fewest versions from
-/// `list_available_versions` contained in the constraints. Then takes that package and finds the
-/// first version contained in the constraints.
-pub fn choose_package_with_fewest_versions<P: Package, VS: VersionSet, T, U, I, F>(
-    list_available_versions: F,
-    potential_packages: impl Iterator<Item = (T, U)>,
-) -> (T, Option<VS::V>)
-where
-    T: Borrow<P>,
-    U: Borrow<VS>,
-    I: Iterator<Item = VS::V>,
-    F: Fn(&P) -> I,
-{
-    let count_valid = |(p, set): &(T, U)| {
-        list_available_versions(p.borrow())
-            .filter(|v| set.borrow().contains(v.borrow()))
-            .count()
-    };
-    let (pkg, set) = potential_packages
-        .min_by_key(count_valid)
-        .expect("potential_packages gave us an empty iterator");
-    let version = list_available_versions(pkg.borrow()).find(|v| set.borrow().contains(v.borrow()));
-    (pkg, version)
-}
-
 /// A basic implementation of [DependencyProvider].
 #[derive(Debug, Clone, Default)]
 #[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
@@ -367,22 +315,21 @@ impl<P: Package, VS: VersionSet> OfflineDependencyProvider<P, VS> {
 /// Packages are picked with the fewest versions contained in the constraints first.
 /// Versions are picked with the newest versions first.
 impl<P: Package, VS: VersionSet> DependencyProvider<P, VS> for OfflineDependencyProvider<P, VS> {
-    #[allow(clippy::type_complexity)]
-    fn choose_package_version<T: Borrow<P>, U: Borrow<VS>>(
-        &self,
-        potential_packages: impl Iterator<Item = (T, U)>,
-    ) -> Result<(T, Option<VS::V>), Box<dyn Error>> {
-        Ok(choose_package_with_fewest_versions(
-            |p| {
-                self.dependencies
-                    .get(p)
-                    .into_iter()
-                    .flat_map(|k| k.keys())
-                    .rev()
-                    .cloned()
-            },
-            potential_packages,
-        ))
+    fn choose_version(&self, package: &P, range: &VS) -> Result<Option<VS::V>, Box<dyn Error>> {
+        Ok(self
+            .dependencies
+            .get(package)
+            .and_then(|versions| versions.keys().rev().find(|v| range.contains(v)).cloned()))
+    }
+
+    type Priority = Reverse<usize>;
+    fn prioritize(&self, package: &P, range: &VS) -> Self::Priority {
+        Reverse(
+            self.dependencies
+                .get(package)
+                .map(|versions| versions.keys().filter(|v| range.contains(v)).count())
+                .unwrap_or(0),
+        )
     }
 
     fn get_dependencies(