From a504c478d5395e7f8a769fa59388c5cc8f03c6e7 Mon Sep 17 00:00:00 2001
From: nichmor <nmorkotilo@gmail.com>
Date: Fri, 14 Jun 2024 15:23:49 +0300
Subject: [PATCH] experiment: use cached packages

---
 crates/rattler_solve/src/lib.rs               |  8 +++++--
 .../rattler_solve/src/resolvo/conda_util.rs   | 21 ++++++++++++++++++-
 crates/rattler_solve/src/resolvo/mod.rs       | 13 ++++++++++--
 3 files changed, 37 insertions(+), 5 deletions(-)

diff --git a/crates/rattler_solve/src/lib.rs b/crates/rattler_solve/src/lib.rs
index 7e4f7c9ec..36201cef5 100644
--- a/crates/rattler_solve/src/lib.rs
+++ b/crates/rattler_solve/src/lib.rs
@@ -9,10 +9,10 @@ pub mod libsolv_c;
 #[cfg(feature = "resolvo")]
 pub mod resolvo;
 
-use std::fmt;
+use std::{collections::HashMap, fmt};
 
 use chrono::{DateTime, Utc};
-use rattler_conda_types::{GenericVirtualPackage, MatchSpec, RepoDataRecord};
+use rattler_conda_types::{GenericVirtualPackage, MatchSpec, PackageName, RepoDataRecord, Version};
 
 /// Represents a solver implementation, capable of solving [`SolverTask`]s
 pub trait SolverImpl {
@@ -149,6 +149,9 @@ pub struct SolverTask<TAvailablePackagesIterator> {
 
     /// The solve strategy.
     pub strategy: SolveStrategy,
+
+    /// Cached packages.
+    pub cached_packages: HashMap<PackageName, Version>,
 }
 
 impl<'r, I: IntoIterator<Item = &'r RepoDataRecord>> FromIterator<I>
@@ -166,6 +169,7 @@ impl<'r, I: IntoIterator<Item = &'r RepoDataRecord>> FromIterator<I>
             channel_priority: ChannelPriority::default(),
             exclude_newer: None,
             strategy: SolveStrategy::default(),
+            cached_packages: Default::default(),
         }
     }
 }
diff --git a/crates/rattler_solve/src/resolvo/conda_util.rs b/crates/rattler_solve/src/resolvo/conda_util.rs
index c6962e389..6316cfda9 100644
--- a/crates/rattler_solve/src/resolvo/conda_util.rs
+++ b/crates/rattler_solve/src/resolvo/conda_util.rs
@@ -1,7 +1,7 @@
 use std::{cmp::Ordering, collections::HashMap};
 
 use futures::future::FutureExt;
-use rattler_conda_types::Version;
+use rattler_conda_types::{PackageName, Version};
 use resolvo::{Dependencies, SolvableId, SolverCache, VersionSetId};
 
 use crate::resolvo::CondaDependencyProvider;
@@ -23,6 +23,7 @@ pub(super) fn compare_candidates(
         Option<(rattler_conda_types::Version, bool)>,
     >,
     strategy: CompareStrategy,
+    cached_versions: HashMap<PackageName, Version>
 ) -> Ordering {
     let pool = &solver.provider().pool;
 
@@ -42,6 +43,24 @@ pub(super) fn compare_candidates(
         Ordering::Equal => {}
     };
 
+    // Check if either version is in the cached versions set
+    let some_cached_version = cached_versions.get(a_record.name());
+    tracing::error!("cached_version: {:?}", cached_versions);
+    if let Some(cc) = some_cached_version{
+        tracing::error!("cached_version: {:?}", cc);
+    }
+    if let Some(cached_version) = some_cached_version{
+        let a_is_cached = cached_version == a_record.version();
+        let b_is_cached = cached_version == b_record.version();
+        match a_is_cached.cmp(&b_is_cached) {
+            Ordering::Greater => return Ordering::Less,
+            Ordering::Less => return Ordering::Greater,
+            Ordering::Equal => {}
+        }
+
+    } 
+    
+
     // Otherwise, select the variant with the highest version
     match (strategy, a_record.version().cmp(b_record.version())) {
         (CompareStrategy::Default, Ordering::Greater)
diff --git a/crates/rattler_solve/src/resolvo/mod.rs b/crates/rattler_solve/src/resolvo/mod.rs
index ed53cd7d8..86a577556 100644
--- a/crates/rattler_solve/src/resolvo/mod.rs
+++ b/crates/rattler_solve/src/resolvo/mod.rs
@@ -6,7 +6,7 @@ use std::{
     collections::{HashMap, HashSet},
     fmt::{Display, Formatter},
     marker::PhantomData,
-    ops::Deref,
+    ops::Deref, str::FromStr,
 };
 
 use chrono::{DateTime, Utc};
@@ -174,6 +174,8 @@ pub struct CondaDependencyProvider<'a> {
     strategy: SolveStrategy,
 
     direct_dependencies: HashSet<NameId>,
+
+    cached_versions: HashMap<PackageName, rattler_conda_types::Version>
 }
 
 impl<'a> CondaDependencyProvider<'a> {
@@ -189,6 +191,7 @@ impl<'a> CondaDependencyProvider<'a> {
         channel_priority: ChannelPriority,
         exclude_newer: Option<DateTime<Utc>>,
         strategy: SolveStrategy,
+        cached_versions: HashMap<PackageName, rattler_conda_types::Version>,
     ) -> Result<Self, SolveError> {
         let pool = Pool::default();
         let mut records: HashMap<NameId, Candidates> = HashMap::default();
@@ -390,6 +393,10 @@ impl<'a> CondaDependencyProvider<'a> {
             candidates.locked = Some(solvable);
         }
 
+        let cached_versions = HashMap::from([
+            (PackageName::try_from("boltons").unwrap(), rattler_conda_types::Version::from_str("21.0.0").unwrap()),
+        ]);
+
         Ok(Self {
             pool,
             records,
@@ -398,6 +405,7 @@ impl<'a> CondaDependencyProvider<'a> {
             stop_time,
             strategy,
             direct_dependencies,
+            cached_versions,
         })
     }
 
@@ -478,7 +486,7 @@ impl<'a> DependencyProvider for CondaDependencyProvider<'a> {
             }
         };
         solvables.sort_by(|&p1, &p2| {
-            conda_util::compare_candidates(p1, p2, solver, &mut highest_version_spec, strategy)
+            conda_util::compare_candidates(p1, p2, solver, &mut highest_version_spec, strategy, self.cached_versions.clone())
         });
     }
 
@@ -608,6 +616,7 @@ impl super::SolverImpl for Solver {
             task.channel_priority,
             task.exclude_newer,
             task.strategy,
+            task.cached_packages.clone(),
         )?;
 
         // Construct the requirements that the solver needs to satisfy.