From b6ed6412485ed9e1d55799459e60bd68be23c9d0 Mon Sep 17 00:00:00 2001
From: konstin <konstin@mailbox.org>
Date: Thu, 12 Dec 2024 13:45:46 +0100
Subject: [PATCH] Allow backtracking to before a specific package

This allows discarding a previously made decision if it turned out to be a bad decision, even if all options with this decision have not yet been rejected.

We allow attempting to backtrack on packages that were not decided yet to avoid the caller from making the duplicative check.
---
 src/internal/core.rs             | 20 +++++++++++++++--
 src/internal/partial_solution.rs | 38 +++++++++++++++++++++++++++-----
 2 files changed, 50 insertions(+), 8 deletions(-)

diff --git a/src/internal/core.rs b/src/internal/core.rs
index 26a699d9..9e6d2181 100644
--- a/src/internal/core.rs
+++ b/src/internal/core.rs
@@ -82,7 +82,7 @@ impl<DP: DependencyProvider> State<DP> {
     ) {
         let dep_incompats =
             self.add_incompatibility_from_dependencies(package, version.clone(), dependencies);
-        self.partial_solution.add_package_version_incompatibilities(
+        self.partial_solution.add_package_version_dependencies(
             package,
             version.clone(),
             dep_incompats,
@@ -249,7 +249,8 @@ impl<DP: DependencyProvider> State<DP> {
         }
     }
 
-    /// Backtracking.
+    /// After a conflict occurred, backtrack the partial solution to a given decision level, and add
+    /// the incompatibility if it was new.
     fn backtrack(
         &mut self,
         incompat: IncompDpId<DP>,
@@ -265,6 +266,21 @@ impl<DP: DependencyProvider> State<DP> {
         }
     }
 
+    /// Manually backtrack before the given package was selected.
+    ///
+    /// This can be used to switch the order of packages if the previous prioritization was bad.
+    ///
+    /// Returns the number of the decisions that were backtracked, or `None` if the package was not
+    /// decided on yet.
+    pub fn backtrack_package(&mut self, package: Id<DP::P>) -> Option<u32> {
+        let base_decision_level = self.partial_solution.current_decision_level();
+        let new_decision_level = self.partial_solution.backtrack_package(package).ok()?;
+        // Remove contradicted incompatibilities that depend on decisions we just backtracked away.
+        self.contradicted_incompatibilities
+            .retain(|_, dl| *dl <= new_decision_level);
+        Some(base_decision_level.0 - new_decision_level.0)
+    }
+
     /// Add this incompatibility into the set of all incompatibilities.
     ///
     /// PubGrub collapses identical dependencies from adjacent package versions
diff --git a/src/internal/partial_solution.rs b/src/internal/partial_solution.rs
index 9804dcb8..d03506d6 100644
--- a/src/internal/partial_solution.rs
+++ b/src/internal/partial_solution.rs
@@ -6,6 +6,7 @@
 use std::fmt::{Debug, Display};
 use std::hash::BuildHasherDefault;
 
+use log::debug;
 use priority_queue::PriorityQueue;
 use rustc_hash::FxHasher;
 
@@ -377,12 +378,37 @@ impl<DP: DependencyProvider> PartialSolution<DP> {
         self.has_ever_backtracked = true;
     }
 
-    /// We can add the version to the partial solution as a decision
-    /// if it doesn't produce any conflict with the new incompatibilities.
-    /// In practice I think it can only produce a conflict if one of the dependencies
-    /// (which are used to make the new incompatibilities)
-    /// is already in the partial solution with an incompatible version.
-    pub(crate) fn add_package_version_incompatibilities(
+    /// Backtrack the partial solution before a particular package was selected.
+    ///
+    /// This can be used to switch the order of packages if the previous prioritization was bad.
+    ///
+    /// Returns the new decision level on success and an error if the package was not decided on
+    /// yet.
+    pub(crate) fn backtrack_package(&mut self, package: Id<DP::P>) -> Result<DecisionLevel, ()> {
+        let Some(decision_level) = self.package_assignments.get_index_of(&package) else {
+            return Err(());
+        };
+        let decision_level = DecisionLevel(decision_level as u32);
+        if decision_level > self.current_decision_level {
+            return Err(());
+        }
+        debug!(
+            "Package backtracking ot decision level {}",
+            decision_level.0
+        );
+        self.backtrack(decision_level);
+        Ok(decision_level)
+    }
+
+    /// Add a package version as decision if none of its dependencies conflicts with the partial
+    /// solution.
+    ///
+    /// If the resolution never backtracked before, a fast path adds the package version directly
+    /// without checking dependencies.
+    ///
+    /// Returns the incompatibility that caused the current version to be rejected, if a conflict
+    /// in the dependencies was found.
+    pub(crate) fn add_package_version_dependencies(
         &mut self,
         package: Id<DP::P>,
         version: DP::V,