uv/crates/uv-distribution/src/index/registry_wheel_index.rs

use std::collections::hash_map::Entry;
use std::collections::BTreeMap;

use rustc_hash::FxHashMap;

use distribution_types::{CachedRegistryDist, FlatIndexLocation, Hashed, IndexLocations, IndexUrl};
use pep440_rs::Version;
use pep508_rs::VerbatimUrl;
use platform_tags::Tags;
use uv_cache::{Cache, CacheBucket, WheelCache};
use uv_fs::{directories, files, symlinks};
use uv_normalize::PackageName;
use uv_types::RequiredHashes;

use crate::index::cached_wheel::CachedWheel;
use crate::source::{read_http_revision, REVISION};

/// A local index of distributions that originate from a registry, like `PyPI`.
#[derive(Debug)]
pub struct RegistryWheelIndex<'a> {
    cache: &'a Cache,
    tags: &'a Tags,
    index_locations: &'a IndexLocations,
    hashes: &'a RequiredHashes,
    index: FxHashMap<&'a PackageName, BTreeMap<Version, CachedRegistryDist>>,
}

impl<'a> RegistryWheelIndex<'a> {
    /// Initialize an index of registry distributions.
    pub fn new(
        cache: &'a Cache,
        tags: &'a Tags,
        index_locations: &'a IndexLocations,
        hashes: &'a RequiredHashes,
    ) -> Self {
        Self {
            cache,
            tags,
            index_locations,
            hashes,
            index: FxHashMap::default(),
        }
    }

    /// Return an iterator over available wheels for a given package.
    ///
    /// If the package is not yet indexed, this will index the package by reading from the cache.
    pub fn get(
        &mut self,
        name: &'a PackageName,
    ) -> impl Iterator<Item = (&Version, &CachedRegistryDist)> {
        self.get_impl(name).iter().rev()
    }

    /// Get the best wheel for the given package name and version.
    ///
    /// If the package is not yet indexed, this will index the package by reading from the cache.
    pub fn get_version(
        &mut self,
        name: &'a PackageName,
        version: &Version,
    ) -> Option<&CachedRegistryDist> {
        self.get_impl(name).get(version)
    }

    /// Get an entry in the index.
    fn get_impl(&mut self, name: &'a PackageName) -> &BTreeMap<Version, CachedRegistryDist> {
        let versions = match self.index.entry(name) {
            Entry::Occupied(entry) => entry.into_mut(),
            Entry::Vacant(entry) => entry.insert(Self::index(
                name,
                self.cache,
                self.tags,
                self.index_locations,
                self.hashes,
            )),
        };
        versions
    }

    /// Add a package to the index by reading from the cache.
    fn index(
        package: &PackageName,
        cache: &Cache,
        tags: &Tags,
        index_locations: &IndexLocations,
        hashes: &RequiredHashes,
    ) -> BTreeMap<Version, CachedRegistryDist> {
        let mut versions = BTreeMap::new();
        let hashes = hashes.get(package).unwrap_or_default();

        // Collect into owned `IndexUrl`
        let flat_index_urls: Vec<IndexUrl> = index_locations
            .flat_index()
            .filter_map(|flat_index| match flat_index {
                FlatIndexLocation::Path(path) => {
                    let path = fs_err::canonicalize(path).ok()?;
                    Some(IndexUrl::Url(VerbatimUrl::from_path(path)))
                }
                FlatIndexLocation::Url(url) => {
                    Some(IndexUrl::Url(VerbatimUrl::unknown(url.clone())))
                }
            })
            .collect();

        for index_url in index_locations.indexes().chain(flat_index_urls.iter()) {
            // Index all the wheels that were downloaded directly from the registry.
            let wheel_dir = cache.shard(
                CacheBucket::Wheels,
                WheelCache::Index(index_url).wheel_dir(package.to_string()),
            );

            // For registry wheels, the cache structure is: `<index>/<package-name>/<wheel>.http`
            // or `<index>/<package-name>/<version>/<wheel>.rev`.
            for file in files(&wheel_dir) {
                if file
                    .extension()
                    .is_some_and(|ext| ext.eq_ignore_ascii_case("http"))
                {
                    if let Some(wheel) = CachedWheel::from_http_pointer(&wheel_dir.join(&file)) {
                        // Enforce hash-checking based on the built distribution.
                        if wheel.satisfies(hashes) {
                            Self::add_wheel(wheel, tags, &mut versions);
                        }
                    }
                }

                if file
                    .extension()
                    .is_some_and(|ext| ext.eq_ignore_ascii_case("rev"))
                {
                    if let Some(wheel) = CachedWheel::from_revision_pointer(&wheel_dir.join(&file))
                    {
                        // Enforce hash-checking based on the built distribution.
                        if wheel.satisfies(hashes) {
                            Self::add_wheel(wheel, tags, &mut versions);
                        }
                    }
                }
            }

            // Index all the built wheels, created by downloading and building source distributions
            // from the registry.
            let cache_shard = cache.shard(
                CacheBucket::BuiltWheels,
                WheelCache::Index(index_url).wheel_dir(package.to_string()),
            );

            // For registry wheels, the cache structure is: `<index>/<package-name>/<version>/`.
            for shard in directories(&cache_shard) {
                // Read the existing metadata from the cache, if it exists.
                let cache_shard = cache_shard.shard(shard);
                let revision_entry = cache_shard.entry(REVISION);
                if let Ok(Some(revision)) = read_http_revision(&revision_entry) {
                    // Enforce hash-checking based on the source distribution.
                    if revision.satisfies(hashes) {
                        for wheel_dir in symlinks(cache_shard.join(revision.id())) {
                            if let Some(wheel) = CachedWheel::from_built_source(&wheel_dir) {
                                Self::add_wheel(wheel, tags, &mut versions);
                            }
                        }
                    }
                };
            }
        }

        versions
    }

    /// Add the [`CachedWheel`] to the index.
    fn add_wheel(
        wheel: CachedWheel,
        tags: &Tags,
        versions: &mut BTreeMap<Version, CachedRegistryDist>,
    ) {
        let dist_info = wheel.into_registry_dist();

        // Pick the wheel with the highest priority
        let compatibility = dist_info.filename.compatibility(tags);
        if let Some(existing) = versions.get_mut(&dist_info.filename.version) {
            // Override if we have better compatibility
            if compatibility > existing.filename.compatibility(tags) {
                *existing = dist_info;
            }
        } else if compatibility.is_compatible() {
            versions.insert(dist_info.filename.version.clone(), dist_info);
        }
    }
}