uv_distribution/index/

built_wheel_index.rs

use std::borrow::Cow;

use uv_cache::{Cache, CacheBucket, CacheShard, WheelCache};
use uv_cache_info::CacheInfo;
use uv_distribution_types::{
    BuildInfo, BuildVariables, ConfigSettings, DirectUrlSourceDist, DirectorySourceDist,
    ExtraBuildRequirement, ExtraBuildRequires, ExtraBuildVariables, GitDirectorySourceDist,
    GitPathSourceDist, Hashed, PackageConfigSettings, PathSourceDist,
};
use uv_normalize::PackageName;
use uv_platform_tags::Tags;
use uv_pypi_types::HashDigests;
use uv_types::HashStrategy;

use crate::Error;
use crate::index::cached_wheel::{CachedWheel, ResolvedWheel};
use crate::source::{
    HASHES, HTTP_REVISION, HttpRevisionPointer, LOCAL_REVISION, LocalRevisionPointer,
    RevisionHashes,
};

/// A local index of built distributions for a specific source distribution.
#[derive(Debug)]
pub struct BuiltWheelIndex<'a> {
    cache: &'a Cache,
    tags: &'a Tags,
    hasher: &'a HashStrategy,
    config_settings: &'a ConfigSettings,
    config_settings_package: &'a PackageConfigSettings,
    extra_build_requires: &'a ExtraBuildRequires,
    extra_build_variables: &'a ExtraBuildVariables,
}

impl<'a> BuiltWheelIndex<'a> {
    /// Initialize an index of built distributions.
    pub fn new(
        cache: &'a Cache,
        tags: &'a Tags,
        hasher: &'a HashStrategy,
        config_settings: &'a ConfigSettings,
        config_settings_package: &'a PackageConfigSettings,
        extra_build_requires: &'a ExtraBuildRequires,
        extra_build_variables: &'a ExtraBuildVariables,
    ) -> Self {
        Self {
            cache,
            tags,
            hasher,
            config_settings,
            config_settings_package,
            extra_build_requires,
            extra_build_variables,
        }
    }

    /// Return the most compatible [`CachedWheel`] for a given source distribution at a direct URL.
    ///
    /// This method does not perform any freshness checks and assumes that the source distribution
    /// is already up-to-date.
    pub fn url(&self, source_dist: &DirectUrlSourceDist) -> Result<Option<CachedWheel>, Error> {
        // For direct URLs, cache directly under the hash of the URL itself.
        let cache_shard = self.cache.shard(
            CacheBucket::SourceDistributions,
            WheelCache::Url(source_dist.url.raw()).root(),
        );

        // Read the revision from the cache.
        let Some(pointer) = HttpRevisionPointer::read_from(cache_shard.entry(HTTP_REVISION))?
        else {
            return Ok(None);
        };

        // Enforce hash-checking by omitting any wheels that don't satisfy the required hashes.
        let revision = pointer.into_revision();
        if !revision.satisfies(self.hasher.get(source_dist)) {
            return Ok(None);
        }

        let cache_shard = cache_shard.shard(revision.id());

        // If there are build settings, we need to scope to a cache shard.
        let config_settings = self.config_settings_for(&source_dist.name);
        let extra_build_deps = self.extra_build_requires_for(&source_dist.name);
        let extra_build_vars = self.extra_build_variables_for(&source_dist.name);
        let build_info = BuildInfo::from_settings(
            config_settings.into_owned(),
            extra_build_deps.to_vec(),
            extra_build_vars.cloned(),
        );
        let cache_shard = build_info
            .cache_shard()
            .map(|digest| cache_shard.shard(digest))
            .unwrap_or(cache_shard);

        Ok(self.find(&cache_shard).map(|wheel| {
            CachedWheel::from_entry(
                wheel,
                revision.into_hashes(),
                CacheInfo::default(),
                build_info,
            )
        }))
    }

    /// Return the most compatible [`CachedWheel`] for a given source distribution at a local path.
    pub fn path(&self, source_dist: &PathSourceDist) -> Result<Option<CachedWheel>, Error> {
        let cache_shard = self.cache.shard(
            CacheBucket::SourceDistributions,
            WheelCache::Path(&source_dist.url).root(),
        );

        // Read the revision from the cache.
        let Some(pointer) = LocalRevisionPointer::read_from(cache_shard.entry(LOCAL_REVISION))?
        else {
            return Ok(None);
        };

        // If the distribution is stale, omit it from the index.
        let cache_info =
            CacheInfo::from_file(&source_dist.install_path).map_err(Error::CacheRead)?;
        if cache_info != *pointer.cache_info() {
            return Ok(None);
        }

        // Enforce hash-checking by omitting any wheels that don't satisfy the required hashes.
        let revision = pointer.into_revision();
        if !revision.satisfies(self.hasher.get(source_dist)) {
            return Ok(None);
        }

        let cache_shard = cache_shard.shard(revision.id());

        // If there are build settings, we need to scope to a cache shard.
        let config_settings = self.config_settings_for(&source_dist.name);
        let extra_build_deps = self.extra_build_requires_for(&source_dist.name);
        let extra_build_vars = self.extra_build_variables_for(&source_dist.name);
        let build_info = BuildInfo::from_settings(
            config_settings.into_owned(),
            extra_build_deps.to_vec(),
            extra_build_vars.cloned(),
        );
        let cache_shard = build_info
            .cache_shard()
            .map(|digest| cache_shard.shard(digest))
            .unwrap_or(cache_shard);

        Ok(self.find(&cache_shard).map(|wheel| {
            CachedWheel::from_entry(wheel, revision.into_hashes(), cache_info, build_info)
        }))
    }

    /// Return the most compatible [`CachedWheel`] for a given source distribution built from a
    /// local directory (source tree).
    pub fn directory(
        &self,
        source_dist: &DirectorySourceDist,
    ) -> Result<Option<CachedWheel>, Error> {
        let cache_shard = self.cache.shard(
            CacheBucket::SourceDistributions,
            if source_dist.editable.unwrap_or(false) {
                WheelCache::Editable(&source_dist.url).root()
            } else {
                WheelCache::Path(&source_dist.url).root()
            },
        );

        // Read the revision from the cache.
        let Some(pointer) = LocalRevisionPointer::read_from(cache_shard.entry(LOCAL_REVISION))?
        else {
            return Ok(None);
        };

        // If the distribution is stale, omit it from the index.
        let cache_info = CacheInfo::from_directory(&source_dist.install_path)?;
        if cache_info != *pointer.cache_info() {
            return Ok(None);
        }

        // Enforce hash-checking by omitting any wheels that don't satisfy the required hashes.
        let revision = pointer.into_revision();
        if !revision.satisfies(self.hasher.get(source_dist)) {
            return Ok(None);
        }

        let cache_shard = cache_shard.shard(revision.id());

        // If there are build settings, we need to scope to a cache shard.
        let config_settings = self.config_settings_for(&source_dist.name);
        let extra_build_deps = self.extra_build_requires_for(&source_dist.name);
        let extra_build_vars = self.extra_build_variables_for(&source_dist.name);
        let build_info = BuildInfo::from_settings(
            config_settings.into_owned(),
            extra_build_deps.to_vec(),
            extra_build_vars.cloned(),
        );
        let cache_shard = build_info
            .cache_shard()
            .map(|digest| cache_shard.shard(digest))
            .unwrap_or(cache_shard);

        Ok(self.find(&cache_shard).map(|wheel| {
            CachedWheel::from_entry(wheel, revision.into_hashes(), cache_info, build_info)
        }))
    }

    /// Return the most compatible [`CachedWheel`] for a given source distribution at a git URL.
    pub fn git_directory(&self, source_dist: &GitDirectorySourceDist) -> Option<CachedWheel> {
        // Enforce hash-checking, which isn't supported for Git distributions.
        if self.hasher.get(source_dist).requires_validation() {
            return None;
        }

        let git_sha = source_dist.git.precise()?;

        let cache_shard = self.cache.shard(
            CacheBucket::SourceDistributions,
            WheelCache::Git(&source_dist.url, git_sha.as_short_str()).root(),
        );

        // If there are build settings, we need to scope to a cache shard.
        let config_settings = self.config_settings_for(&source_dist.name);
        let extra_build_deps = self.extra_build_requires_for(&source_dist.name);
        let extra_build_vars = self.extra_build_variables_for(&source_dist.name);
        let build_info = BuildInfo::from_settings(
            config_settings.into_owned(),
            extra_build_deps.to_vec(),
            extra_build_vars.cloned(),
        );
        let cache_shard = build_info
            .cache_shard()
            .map(|digest| cache_shard.shard(digest))
            .unwrap_or(cache_shard);

        self.find(&cache_shard).map(|wheel| {
            CachedWheel::from_entry(
                wheel,
                HashDigests::empty(),
                CacheInfo::default(),
                build_info,
            )
        })
    }

    /// Return the most compatible [`CachedWheel`] for a given source distribution at a git URL.
    pub fn git_path(&self, source_dist: &GitPathSourceDist) -> Result<Option<CachedWheel>, Error> {
        let Some(git_sha) = source_dist.git.precise() else {
            return Ok(None);
        };

        let cache_shard = self.cache.shard(
            CacheBucket::SourceDistributions,
            WheelCache::Git(&source_dist.url, git_sha.as_short_str()).root(),
        );

        // Read the revision from the cache.
        let Some(revision) = RevisionHashes::read_from(cache_shard.entry(HASHES))? else {
            return Ok(None);
        };

        // Enforce hash-checking by omitting any wheels that don't satisfy the required hashes.
        if !revision.satisfies(self.hasher.get(source_dist)) {
            return Ok(None);
        }

        // If there are build settings, we need to scope to a cache shard.
        let config_settings = self.config_settings_for(&source_dist.name);
        let extra_build_deps = self.extra_build_requires_for(&source_dist.name);
        let extra_build_vars = self.extra_build_variables_for(&source_dist.name);
        let build_info = BuildInfo::from_settings(
            config_settings.into_owned(),
            extra_build_deps.to_vec(),
            extra_build_vars.cloned(),
        );
        let cache_shard = build_info
            .cache_shard()
            .map(|digest| cache_shard.shard(digest))
            .unwrap_or(cache_shard);

        Ok(self.find(&cache_shard).map(|wheel| {
            CachedWheel::from_entry(
                wheel,
                revision.into_hashes(),
                CacheInfo::default(),
                build_info,
            )
        }))
    }

    /// Find the "best" distribution in the index for a given source distribution.
    ///
    /// This lookup prefers newer versions over older versions, and aims to maximize compatibility
    /// with the target platform.
    ///
    /// The `shard` should point to a directory containing the built distributions for a specific
    /// source distribution. For example, given the built wheel cache structure:
    /// ```text
    /// built-wheels-v0/
    /// └── pypi
    ///     └── django-allauth-0.51.0.tar.gz
    ///         ├── django_allauth-0.51.0-py3-none-any.whl
    ///         └── metadata.json
    /// ```
    ///
    /// The `shard` should be `built-wheels-v0/pypi/django-allauth-0.51.0.tar.gz`.
    fn find(&self, shard: &CacheShard) -> Option<ResolvedWheel> {
        let mut candidate: Option<ResolvedWheel> = None;

        // Unzipped wheels are stored as symlinks into the archive directory.
        for wheel_dir in uv_fs::entries(shard).ok().into_iter().flatten() {
            // Ignore any `.lock` files.
            if wheel_dir
                .extension()
                .is_some_and(|ext| ext.eq_ignore_ascii_case("lock"))
            {
                continue;
            }

            match ResolvedWheel::from_built_source(&wheel_dir, self.cache) {
                None => {}
                Some(dist_info) => {
                    // Pick the wheel with the highest priority
                    let compatibility = dist_info.filename.compatibility(self.tags);

                    // Only consider wheels that are compatible with our tags.
                    if !compatibility.is_compatible() {
                        continue;
                    }

                    if let Some(existing) = candidate.as_ref() {
                        // Override if the wheel is newer, or "more" compatible.
                        if dist_info.filename.version > existing.filename.version
                            || compatibility > existing.filename.compatibility(self.tags)
                        {
                            candidate = Some(dist_info);
                        }
                    } else {
                        candidate = Some(dist_info);
                    }
                }
            }
        }

        candidate
    }

    /// Determine the [`ConfigSettings`] for the given package name.
    fn config_settings_for(&self, name: &PackageName) -> Cow<'_, ConfigSettings> {
        if let Some(package_settings) = self.config_settings_package.get(name) {
            Cow::Owned(package_settings.clone().merge(self.config_settings.clone()))
        } else {
            Cow::Borrowed(self.config_settings)
        }
    }

    /// Determine the extra build requirements for the given package name.
    fn extra_build_requires_for(&self, name: &PackageName) -> &[ExtraBuildRequirement] {
        self.extra_build_requires
            .get(name)
            .map(Vec::as_slice)
            .unwrap_or(&[])
    }

    /// Determine the extra build variables for the given package name.
    fn extra_build_variables_for(&self, name: &PackageName) -> Option<&BuildVariables> {
        self.extra_build_variables.get(name)
    }
}