Faster boto3, second half

crates/distribution-types/src/resolved.rs

@@ -1,4 +1,4 @@
-use std::fmt::Display;
+use std::fmt::{Display, Formatter};
 
 use pep508_rs::PackageName;
 
@@ -42,6 +42,15 @@ impl ResolvedDistRef<'_> {
     }
 }
 
+impl Display for ResolvedDistRef<'_> {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        match self {
+            Self::Installable(dist) => Display::fmt(dist, f),
+            Self::Installed(dist) => Display::fmt(dist, f),
+        }
+    }
+}
+
 impl Name for ResolvedDistRef<'_> {
     fn name(&self) -> &PackageName {
         match self {

crates/uv-resolver/src/candidate_selector.rs

@@ -62,13 +62,9 @@ enum AllowPreRelease {
 }
 
 impl CandidateSelector {
-    /// Select a [`Candidate`] from a set of candidate versions and files.
-    ///
-    /// Unless present in the provided [`Exclusions`], local distributions from the
-    /// [`InstalledPackagesProvider`] are preferred over remote distributions in
-    /// the [`VersionMap`].
-    pub(crate) fn select<'a, InstalledPackages: InstalledPackagesProvider>(
-        &'a self,
+    /// Get a preferred version if one exists. This is the preference from a lockfile or a locally
+    /// installed version.
+    pub(crate) fn get_preference<'a, InstalledPackages: InstalledPackagesProvider>(
         package_name: &'a PackageName,
         range: &Range<Version>,
         version_maps: &'a [VersionMap],
@@ -141,8 +137,12 @@ impl CandidateSelector {
             }
         }
 
-        // Determine the appropriate prerelease strategy for the current package.
-        let allow_prerelease = match &self.prerelease_strategy {
+        None
+    }
+
+    /// Determine the appropriate prerelease strategy for the current package.
+    fn allow_prereleases(&self, package_name: &PackageName) -> AllowPreRelease {
+        match &self.prerelease_strategy {
             PreReleaseStrategy::Disallow => AllowPreRelease::No,
             PreReleaseStrategy::Allow => AllowPreRelease::Yes,
             PreReleaseStrategy::IfNecessary => AllowPreRelease::IfNecessary,
@@ -160,47 +160,62 @@ impl CandidateSelector {
                     AllowPreRelease::IfNecessary
                 }
             }
-        };
+        }
+    }
+
+    /// Select a [`Candidate`] from a set of candidate versions and files.
+    ///
+    /// Unless present in the provided [`Exclusions`], local distributions from the
+    /// [`InstalledPackagesProvider`] are preferred over remote distributions in
+    /// the [`VersionMap`].
+    pub(crate) fn select<'a, InstalledPackages: InstalledPackagesProvider>(
+        &'a self,
+        package_name: &'a PackageName,
+        range: &Range<Version>,
+        version_maps: &'a [VersionMap],
+        preferences: &'a Preferences,
+        installed_packages: &'a InstalledPackages,
+        exclusions: &'a Exclusions,
+    ) -> Option<Candidate<'a>> {
+        if let Some(preference) = Self::get_preference(
+            package_name,
+            range,
+            version_maps,
+            preferences,
+            installed_packages,
+            exclusions,
+        ) {
+            return Some(preference);
+        }
 
         tracing::trace!(
-            "selecting candidate for package {:?} with range {:?} with {} remote versions",
+            "selecting candidate for package {} with range {:?} with {} remote versions",
             package_name,
             range,
             version_maps.iter().map(VersionMap::len).sum::<usize>(),
         );
-        match &self.resolution_strategy {
-            ResolutionStrategy::Highest => version_maps.iter().find_map(|version_map| {
+        let highest = match &self.resolution_strategy {
+            ResolutionStrategy::Highest => true,
+            ResolutionStrategy::Lowest => false,
+            ResolutionStrategy::LowestDirect(direct_dependencies) => {
+                !direct_dependencies.contains(package_name)
+            }
+        };
+        let allow_prerelease = self.allow_prereleases(package_name);
+
+        if highest {
+            version_maps.iter().find_map(|version_map| {
                 Self::select_candidate(
                     version_map.iter().rev(),
                     package_name,
                     range,
                     allow_prerelease,
                 )
-            }),
-            ResolutionStrategy::Lowest => version_maps.iter().find_map(|version_map| {
+            })
+        } else {
+            version_maps.iter().find_map(|version_map| {
                 Self::select_candidate(version_map.iter(), package_name, range, allow_prerelease)
-            }),
-            ResolutionStrategy::LowestDirect(direct_dependencies) => {
-                if direct_dependencies.contains(package_name) {
-                    version_maps.iter().find_map(|version_map| {
-                        Self::select_candidate(
-                            version_map.iter(),
-                            package_name,
-                            range,
-                            allow_prerelease,
-                        )
-                    })
-                } else {
-                    version_maps.iter().find_map(|version_map| {
-                        Self::select_candidate(
-                            version_map.iter().rev(),
-                            package_name,
-                            range,
-                            allow_prerelease,
-                        )
-                    })
-                }
-            }
+            })
         }
     }
 

crates/uv-resolver/src/resolver/mod.rs

@@ -530,7 +530,14 @@ impl<
         let mut last_version = version.clone();
         let mut range = current_range.clone();
         while counter > 0 {
-            let Some(candidate) = self.selector.select(package_name, &range, version_map) else {
+            let Some(candidate) = self.selector.select(
+                package_name,
+                &range,
+                version_map,
+                &self.preferences,
+                self.installed_packages,
+                &self.exclusions,
+            ) else {
                 break;
             };
             range = range.intersection(&Range::singleton(candidate.version().clone()).complement());
@@ -545,7 +552,51 @@ impl<
             // Emit a request to fetch the metadata for this version.
             trace!("Prefetching {}", dist);
             if self.index.distributions.register(candidate.package_id()) {
-                request_sink.send(Request::Dist(dist.clone())).await?;
+                let request = match dist {
+                    ResolvedDistRef::Installable(dist) => Request::Dist(dist.clone()),
+                    ResolvedDistRef::Installed(dist) => Request::Installed(dist.clone()),
+                };
+                request_sink.send(request).await?;
+            }
+        }
+
+        // We encounter cases (botocore) where the above doesn't work: Say we previously selected
+        // a==x.y.z, which depends on b==x.y.z. a==x.y.z is incompatible, but we don't know that
+        // yet. We just select b==x.y.z and want to prefetch, since for all versions of a we try,
+        // we have to wait for the matching version of b. The selector gives us only one version of
+        // b, so we're now here 0 versions prefetched. Instead, we guess that the next version of b
+        // will be x.y.(z-1) and so forth.
+        if counter > 0 {
+            let remaining_range = Range::strictly_lower_than(&last_version);
+
+            'outer: for version_map_ in version_map {
+                // TODO(konstin): comment
+                for (pre_version, maybe_dist) in version_map_
+                    .iter()
+                    .rev()
+                    .filter(|(version, _)| remaining_range.contains(version))
+                    .take(counter)
+                {
+                    let Some(dist) = maybe_dist.prioritized_dist() else {
+                        break 'outer;
+                    };
+                    let candidate = Candidate::new(package_name, pre_version, dist);
+
+                    let CandidateDist::Compatible(dist) = candidate.dist() else {
+                        break 'outer;
+                    };
+                    let dist = dist.for_resolution();
+
+                    // Emit a request to fetch the metadata for this version.
+                    trace!("Prefetching {}", dist);
+                    if self.index.distributions.register(candidate.package_id()) {
+                        let request = match dist {
+                            ResolvedDistRef::Installable(dist) => Request::Dist(dist.clone()),
+                            ResolvedDistRef::Installed(dist) => Request::Installed(dist.clone()),
+                        };
+                        request_sink.send(request).await?;
+                    }
+                }
             }
         }