From 0c4627e2f2e04523c146eb26e118117346fb4513 Mon Sep 17 00:00:00 2001
From: Yorick <yorick@yorickvanpelt.nl>
Date: Mon, 22 Jul 2024 20:03:30 +0200
Subject: [PATCH] If multiple indices contain the same version, use the first
 index (#5288)

This fixes resolving packages that publish an invalid stub to pypi, such
as tensorrt-llm.

## Summary

In https://github.com/astral-sh/uv/pull/3138 , we implemented
`unsafe-best-match`. However, it seems to not quite work as expected.
When multiple indices contain the same version, it's not clear which
index the current code uses. This PR fixes that to use the first index
the package is in.

## Test Plan

```console
$ echo 'tensorrt-llm==0.11.0' | ./target/debug/uv pip compile - --extra-index-url https://pypi.nvidia.com --python-version=3.10 --index-strategy=unsafe-best-match --annotation-style=line
```
---
 crates/uv-resolver/src/candidate_selector.rs | 32 +++++++++++++++++---
 1 file changed, 28 insertions(+), 4 deletions(-)

diff --git a/crates/uv-resolver/src/candidate_selector.rs b/crates/uv-resolver/src/candidate_selector.rs
index f97023c49..607116237 100644
--- a/crates/uv-resolver/src/candidate_selector.rs
+++ b/crates/uv-resolver/src/candidate_selector.rs
@@ -232,8 +232,20 @@ impl CandidateSelector {
                 Self::select_candidate(
                     version_maps
                         .iter()
-                        .map(|version_map| version_map.iter().rev())
-                        .kmerge_by(|(version1, _), (version2, _)| version1 > version2),
+                        .enumerate()
+                        .map(|(map_index, version_map)| {
+                            version_map.iter().rev().map(move |item| (map_index, item))
+                        })
+                        .kmerge_by(
+                            |(index1, (version1, _)), (index2, (version2, _))| match version1
+                                .cmp(version2)
+                            {
+                                std::cmp::Ordering::Equal => index1 < index2,
+                                std::cmp::Ordering::Less => false,
+                                std::cmp::Ordering::Greater => true,
+                            },
+                        )
+                        .map(|(_, item)| item),
                     package_name,
                     range,
                     allow_prerelease,
@@ -242,8 +254,20 @@ impl CandidateSelector {
                 Self::select_candidate(
                     version_maps
                         .iter()
-                        .map(VersionMap::iter)
-                        .kmerge_by(|(version1, _), (version2, _)| version1 < version2),
+                        .enumerate()
+                        .map(|(map_index, version_map)| {
+                            version_map.iter().map(move |item| (map_index, item))
+                        })
+                        .kmerge_by(
+                            |(index1, (version1, _)), (index2, (version2, _))| match version1
+                                .cmp(version2)
+                            {
+                                std::cmp::Ordering::Equal => index1 < index2,
+                                std::cmp::Ordering::Less => true,
+                                std::cmp::Ordering::Greater => false,
+                            },
+                        )
+                        .map(|(_, item)| item),
                     package_name,
                     range,
                     allow_prerelease,