perf: cache candidates seperate from their sorting

crates/libsolv_rs/src/conda_util.rs

@@ -4,17 +4,21 @@ use crate::mapping::Mapping;
 use crate::solvable::Solvable;
 use crate::MatchSpecId;
 use rattler_conda_types::{MatchSpec, Version};
+use std::cell::OnceCell;
 use std::cmp::Ordering;
 use std::collections::HashMap;
 
 /// Returns the order of two candidates based on the order used by conda.
+#[allow(clippy::too_many_arguments)]
 pub(crate) fn compare_candidates(
     a: SolvableId,
     b: SolvableId,
     solvables: &Arena<SolvableId, Solvable>,
     interned_strings: &HashMap<String, NameId>,
     packages_by_name: &Mapping<NameId, Vec<SolvableId>>,
     match_specs: &Arena<MatchSpecId, MatchSpec>,
+    match_spec_to_candidates: &Mapping<MatchSpecId, OnceCell<Vec<SolvableId>>>,
+    match_spec_highest_version: &Mapping<MatchSpecId, OnceCell<Option<(Version, bool)>>>,
 ) -> Ordering {
     let a_solvable = solvables[a].package();
     let b_solvable = solvables[b].package();
@@ -48,28 +52,48 @@ pub(crate) fn compare_candidates(
 
     // Otherwise, compare the dependencies of the variants. If there are similar
     // dependencies select the variant that selects the highest version of the dependency.
-    let a_match_specs = a_solvable.dependencies.iter().map(|id| &match_specs[*id]);
-    let b_match_specs = b_solvable.dependencies.iter().map(|id| &match_specs[*id]);
+    let a_match_specs = a_solvable
+        .dependencies
+        .iter()
+        .map(|id| (*id, &match_specs[*id]));
+    let b_match_specs = b_solvable
+        .dependencies
+        .iter()
+        .map(|id| (*id, &match_specs[*id]));
 
     let b_specs_by_name: HashMap<_, _> = b_match_specs
-        .filter_map(|spec| spec.name.as_ref().map(|name| (name, spec)))
+        .filter_map(|(spec_id, spec)| spec.name.as_ref().map(|name| (name, (spec_id))))
         .collect();
 
-    let a_specs_by_name =
-        a_match_specs.filter_map(|spec| spec.name.as_ref().map(|name| (name, spec)));
+    let a_specs_by_name = a_match_specs
+        .filter_map(|(spec_id, spec)| spec.name.as_ref().map(|name| (name, (spec_id))));
 
     let mut total_score = 0;
-    for (a_dep_name, a_spec) in a_specs_by_name {
-        if let Some(b_spec) = b_specs_by_name.get(&a_dep_name) {
-            if &a_spec == b_spec {
+    for (a_dep_name, a_spec_id) in a_specs_by_name {
+        if let Some(b_spec_id) = b_specs_by_name.get(&a_dep_name) {
+            if &a_spec_id == b_spec_id {
                 continue;
             }
 
             // Find which of the two specs selects the highest version
-            let highest_a =
-                find_highest_version(solvables, interned_strings, packages_by_name, a_spec);
-            let highest_b =
-                find_highest_version(solvables, interned_strings, packages_by_name, b_spec);
+            let highest_a = find_highest_version(
+                a_spec_id,
+                solvables,
+                interned_strings,
+                packages_by_name,
+                match_specs,
+                match_spec_to_candidates,
+                match_spec_highest_version,
+            );
+            let highest_b = find_highest_version(
+                *b_spec_id,
+                solvables,
+                interned_strings,
+                packages_by_name,
+                match_specs,
+                match_spec_to_candidates,
+                match_spec_highest_version,
+            );
 
             // Skip version if no package is selected by either spec
             let (a_version, a_tracked_features, b_version, b_tracked_features) = if let (
@@ -114,34 +138,71 @@ pub(crate) fn compare_candidates(
 }
 
 pub(crate) fn find_highest_version(
+    match_spec_id: MatchSpecId,
     solvables: &Arena<SolvableId, Solvable>,
     interned_strings: &HashMap<String, NameId>,
     packages_by_name: &Mapping<NameId, Vec<SolvableId>>,
-    match_spec: &MatchSpec,
+    match_specs: &Arena<MatchSpecId, MatchSpec>,
+    match_spec_to_candidates: &Mapping<MatchSpecId, OnceCell<Vec<SolvableId>>>,
+    match_spec_highest_version: &Mapping<MatchSpecId, OnceCell<Option<(Version, bool)>>>,
 ) -> Option<(Version, bool)> {
-    let name = match_spec.name.as_deref().unwrap();
-    let name_id = interned_strings[name];
+    match_spec_highest_version[match_spec_id]
+        .get_or_init(|| {
+            let candidates = find_candidates(
+                match_spec_id,
+                match_specs,
+                interned_strings,
+                packages_by_name,
+                solvables,
+                match_spec_to_candidates,
+            );
+
+            candidates
+                .iter()
+                .map(|id| &solvables[*id].package().record)
+                .fold(None, |init, record| {
+                    Some(init.map_or_else(
+                        || {
+                            (
+                                record.version.version().clone(),
+                                !record.track_features.is_empty(),
+                            )
+                        },
+                        |(version, has_tracked_features)| {
+                            (
+                                version.max(record.version.version().clone()),
+                                has_tracked_features && record.track_features.is_empty(),
+                            )
+                        },
+                    ))
+                })
+        })
+        .as_ref()
+        .map(|(version, has_tracked_features)| (version.clone(), *has_tracked_features))
+}
 
-    // For each record that matches the spec
-    let candidates = packages_by_name[name_id]
-        .iter()
-        .map(|&s| solvables[s].package().record)
-        .filter(|s| match_spec.matches(s));
-
-    candidates.fold(None, |init, record| {
-        Some(init.map_or_else(
-            || {
-                (
-                    record.version.version().clone(),
-                    !record.track_features.is_empty(),
-                )
-            },
-            |(version, has_tracked_features)| {
-                (
-                    version.max(record.version.version().clone()),
-                    has_tracked_features && record.track_features.is_empty(),
-                )
-            },
-        ))
+pub(crate) fn find_candidates<'b>(
+    match_spec_id: MatchSpecId,
+    match_specs: &Arena<MatchSpecId, MatchSpec>,
+    names_to_ids: &HashMap<String, NameId>,
+    packages_by_name: &Mapping<NameId, Vec<SolvableId>>,
+    solvables: &Arena<SolvableId, Solvable>,
+    match_spec_to_candidates: &'b Mapping<MatchSpecId, OnceCell<Vec<SolvableId>>>,
+) -> &'b Vec<SolvableId> {
+    match_spec_to_candidates[match_spec_id].get_or_init(|| {
+        let match_spec = &match_specs[match_spec_id];
+        let match_spec_name = match_spec
+            .name
+            .as_deref()
+            .expect("match spec without name!");
+        let name_id = names_to_ids
+            .get(match_spec_name)
+            .expect("cannot find name in name lookup");
+
+        packages_by_name[*name_id]
+            .iter()
+            .cloned()
+            .filter(|&solvable| match_spec.matches(solvables[solvable].package().record))
+            .collect()
     })
 }

crates/libsolv_rs/src/id.rs

@@ -74,6 +74,12 @@ impl ArenaId for SolvableId {
     }
 }
 
+impl From<SolvableId> for u32 {
+    fn from(value: SolvableId) -> Self {
+        value.0
+    }
+}
+
 #[repr(transparent)]
 #[derive(Copy, Clone, PartialOrd, Ord, Eq, PartialEq, Debug, Hash)]
 pub(crate) struct ClauseId(u32);

crates/libsolv_rs/src/pool.rs

@@ -3,7 +3,9 @@ use crate::conda_util;
 use crate::id::{MatchSpecId, NameId, RepoId, SolvableId};
 use crate::mapping::Mapping;
 use crate::solvable::{PackageSolvable, Solvable};
-use rattler_conda_types::{MatchSpec, PackageRecord};
+use rattler_conda_types::{MatchSpec, PackageRecord, Version};
+use std::cell::OnceCell;
+use std::cmp::Ordering;
 use std::collections::hash_map::Entry;
 use std::collections::HashMap;
 use std::str::FromStr;
@@ -35,7 +37,7 @@ pub struct Pool<'a> {
     match_specs_to_ids: HashMap<String, MatchSpecId>,
 
     /// Cached candidates for each match spec
-    pub(crate) match_spec_to_candidates: Mapping<MatchSpecId, Vec<SolvableId>>,
+    pub(crate) match_spec_to_sorted_candidates: Mapping<MatchSpecId, Vec<SolvableId>>,
 
     /// Cached forbidden solvables for each match spec
     pub(crate) match_spec_to_forbidden: Mapping<MatchSpecId, Vec<SolvableId>>,
@@ -50,13 +52,13 @@ impl<'a> Default for Pool<'a> {
             solvables,
             total_repos: 0,
 
-            names_to_ids: HashMap::new(),
+            names_to_ids: Default::default(),
             package_names: Arena::new(),
             packages_by_name: Mapping::empty(),
 
-            match_specs_to_ids: HashMap::default(),
+            match_specs_to_ids: Default::default(),
             match_specs: Arena::new(),
-            match_spec_to_candidates: Mapping::empty(),
+            match_spec_to_sorted_candidates: Mapping::empty(),
             match_spec_to_forbidden: Mapping::empty(),
         }
     }
@@ -127,7 +129,10 @@ impl<'a> Pool<'a> {
         &self,
         match_spec_id: MatchSpecId,
         favored_map: &HashMap<NameId, SolvableId>,
-        match_spec_to_candidates: &mut Mapping<MatchSpecId, Vec<SolvableId>>,
+        match_spec_to_sorted_candidates: &mut Mapping<MatchSpecId, Vec<SolvableId>>,
+        match_spec_to_candidates: &Mapping<MatchSpecId, OnceCell<Vec<SolvableId>>>,
+        match_spec_highest_version: &Mapping<MatchSpecId, OnceCell<Option<(Version, bool)>>>,
+        solvable_order: &mut HashMap<u64, Ordering>,
     ) {
         let match_spec = &self.match_specs[match_spec_id];
         let match_spec_name = match_spec
@@ -139,21 +144,30 @@ impl<'a> Pool<'a> {
             Some(&name_id) => name_id,
         };
 
-        let mut pkgs: Vec<_> = self.packages_by_name[name_id]
-            .iter()
-            .cloned()
-            .filter(|&solvable| match_spec.matches(self.solvables[solvable].package().record))
-            .collect();
+        let mut pkgs = conda_util::find_candidates(
+            match_spec_id,
+            &self.match_specs,
+            &self.names_to_ids,
+            &self.packages_by_name,
+            &self.solvables,
+            match_spec_to_candidates,
+        )
+        .clone();
 
         pkgs.sort_by(|&p1, &p2| {
-            conda_util::compare_candidates(
-                p1,
-                p2,
-                &self.solvables,
-                &self.names_to_ids,
-                &self.packages_by_name,
-                &self.match_specs,
-            )
+            let key = u32::from(p1) as u64 | ((u32::from(p2) as u64) << 32);
+            *solvable_order.entry(key).or_insert_with(|| {
+                conda_util::compare_candidates(
+                    p1,
+                    p2,
+                    &self.solvables,
+                    &self.names_to_ids,
+                    &self.packages_by_name,
+                    &self.match_specs,
+                    match_spec_to_candidates,
+                    match_spec_highest_version,
+                )
+            })
         });
 
         if let Some(&favored_id) = favored_map.get(&name_id) {
@@ -163,7 +177,7 @@ impl<'a> Pool<'a> {
             }
         }
 
-        match_spec_to_candidates[match_spec_id] = pkgs;
+        match_spec_to_sorted_candidates[match_spec_id] = pkgs;
     }
 
     /// Populates the list of forbidden packages for the provided match spec

crates/libsolv_rs/src/problem.rs

@@ -52,7 +52,7 @@ impl Problem {
                 Clause::Requires(package_id, match_spec_id) => {
                     let package_node = Self::add_node(&mut graph, &mut nodes, package_id);
 
-                    let candidates = &solver.pool().match_spec_to_candidates[match_spec_id];
+                    let candidates = &solver.pool().match_spec_to_sorted_candidates[match_spec_id];
                     if candidates.is_empty() {
                         tracing::info!(
                             "{package_id:?} requires {match_spec_id:?}, which has no candidates"

crates/libsolv_rs/src/solver/clause.rs

@@ -130,7 +130,7 @@ impl Clause {
                     negate: true,
                 });
 
-                for &solvable_id in &pool.match_spec_to_candidates[match_spec_id] {
+                for &solvable_id in &pool.match_spec_to_sorted_candidates[match_spec_id] {
                     visit(Literal {
                         solvable_id,
                         negate: false,
@@ -230,6 +230,7 @@ impl ClauseState {
         }
     }
 
+    #[inline]
     pub fn next_watched_clause(&self, solvable_id: SolvableId) -> ClauseId {
         if solvable_id == self.watched_literals[0] {
             self.next_watches[0]
@@ -345,7 +346,7 @@ impl ClauseState {
                 }
 
                 // The available candidates
-                for &candidate in &pool.match_spec_to_candidates[match_spec_id] {
+                for &candidate in &pool.match_spec_to_sorted_candidates[match_spec_id] {
                     let lit = Literal {
                         solvable_id: candidate,
                         negate: false,