Share a single Index across resolutions

crates/puffin-cli/src/commands/pip_compile.rs

@@ -2,6 +2,7 @@ use std::borrow::Cow;
 use std::env;
 use std::fmt::Write;
 use std::io::stdout;
+use std::ops::Deref;
 use std::path::Path;
 use std::str::FromStr;
 
@@ -24,7 +25,8 @@ use puffin_installer::Downloader;
 use puffin_interpreter::{Interpreter, PythonVersion};
 use puffin_normalize::ExtraName;
 use puffin_resolver::{
-    DisplayResolutionGraph, Manifest, PreReleaseMode, ResolutionMode, ResolutionOptions, Resolver,
+    DisplayResolutionGraph, InMemoryIndex, Manifest, PreReleaseMode, ResolutionMode,
+    ResolutionOptions, Resolver,
 };
 use puffin_traits::{InFlight, SetupPyStrategy};
 use requirements_txt::EditableRequirement;
@@ -128,6 +130,18 @@ pub(crate) async fn pip_compile(
         interpreter.sys_executable().display()
     );
 
+    // Create a shared in-memory index.
+    let source_index = InMemoryIndex::default();
+
+    // If we're resolving against a different Python version, use a separate index. Source
+    // distributions will be built against the installed version, and so the index may contain
+    // different package priorities than in the top-level resolution.
+    let top_level_index = if python_version.is_some() {
+        InMemoryIndexRef::Owned(InMemoryIndex::default())
+    } else {
+        InMemoryIndexRef::Borrowed(&source_index)
+    };
+
     // Determine the tags, markers, and interpreter to use for resolution.
     let tags = if let Some(python_version) = python_version.as_ref() {
         Cow::Owned(Tags::from_env(
@@ -164,6 +178,7 @@ pub(crate) async fn pip_compile(
         &interpreter,
         &index_locations,
         &flat_index,
+        &source_index,
         &in_flight,
         interpreter.sys_executable().to_path_buf(),
         setup_py,
@@ -240,6 +255,7 @@ pub(crate) async fn pip_compile(
         &tags,
         &client,
         &flat_index,
+        &top_level_index,
         &build_dispatch,
     )
     .with_reporter(ResolverReporter::from(printer));
@@ -341,3 +357,20 @@ pub(crate) fn extra_name_with_clap_error(arg: &str) -> Result<ExtraName> {
         )
     })
 }
+
+/// An owned or unowned [`InMemoryIndex`].
+enum InMemoryIndexRef<'a> {
+    Owned(InMemoryIndex),
+    Borrowed(&'a InMemoryIndex),
+}
+
+impl Deref for InMemoryIndexRef<'_> {
+    type Target = InMemoryIndex;
+
+    fn deref(&self) -> &Self::Target {
+        match self {
+            Self::Owned(index) => index,
+            Self::Borrowed(index) => index,
+        }
+    }
+}

crates/puffin-cli/src/commands/pip_install.rs

@@ -26,7 +26,8 @@ use puffin_installer::{
 use puffin_interpreter::{Interpreter, Virtualenv};
 use puffin_normalize::PackageName;
 use puffin_resolver::{
-    Manifest, PreReleaseMode, ResolutionGraph, ResolutionMode, ResolutionOptions, Resolver,
+    InMemoryIndex, Manifest, PreReleaseMode, ResolutionGraph, ResolutionMode, ResolutionOptions,
+    Resolver,
 };
 use puffin_traits::{InFlight, SetupPyStrategy};
 use requirements_txt::EditableRequirement;
@@ -144,6 +145,9 @@ pub(crate) async fn pip_install(
         FlatIndex::from_entries(entries, tags)
     };
 
+    // Create a shared in-memory index.
+    let index = InMemoryIndex::default();
+
     // Track in-flight downloads, builds, etc., across resolutions.
     let in_flight = InFlight::default();
 
@@ -155,6 +159,7 @@ pub(crate) async fn pip_install(
         &interpreter,
         &index_locations,
         &flat_index,
+        &index,
         &in_flight,
         venv.python_executable(),
         setup_py,
@@ -196,6 +201,7 @@ pub(crate) async fn pip_install(
         markers,
         &client,
         &flat_index,
+        &index,
         &resolve_dispatch,
         options,
         printer,
@@ -229,6 +235,7 @@ pub(crate) async fn pip_install(
             &interpreter,
             &index_locations,
             &flat_index,
+            &index,
             &in_flight,
             venv.python_executable(),
             setup_py,
@@ -369,6 +376,7 @@ async fn resolve(
     markers: &MarkerEnvironment,
     client: &RegistryClient,
     flat_index: &FlatIndex,
+    index: &InMemoryIndex,
     build_dispatch: &BuildDispatch<'_>,
     options: ResolutionOptions,
     mut printer: Printer,
@@ -415,6 +423,7 @@ async fn resolve(
         tags,
         client,
         flat_index,
+        index,
         build_dispatch,
     )
     .with_reporter(ResolverReporter::from(printer));

crates/puffin-cli/src/commands/pip_sync.rs

@@ -14,6 +14,7 @@ use puffin_client::{FlatIndex, FlatIndexClient, RegistryClient, RegistryClientBu
 use puffin_dispatch::BuildDispatch;
 use puffin_installer::{Downloader, InstallPlan, Reinstall, ResolvedEditable, SitePackages};
 use puffin_interpreter::Virtualenv;
+use puffin_resolver::InMemoryIndex;
 use puffin_traits::{InFlight, SetupPyStrategy};
 use pypi_types::Yanked;
 use requirements_txt::EditableRequirement;
@@ -70,6 +71,9 @@ pub(crate) async fn pip_sync(
         FlatIndex::from_entries(entries, tags)
     };
 
+    // Create a shared in-memory index.
+    let index = InMemoryIndex::default();
+
     // Track in-flight downloads, builds, etc., across resolutions.
     let in_flight = InFlight::default();
 
@@ -80,6 +84,7 @@ pub(crate) async fn pip_sync(
         venv.interpreter(),
         &index_locations,
         &flat_index,
+        &index,
         &in_flight,
         venv.python_executable(),
         setup_py,

crates/puffin-cli/src/commands/venv.rs

@@ -15,6 +15,7 @@ use puffin_cache::Cache;
 use puffin_client::{FlatIndex, FlatIndexClient, RegistryClientBuilder};
 use puffin_dispatch::BuildDispatch;
 use puffin_interpreter::Interpreter;
+use puffin_resolver::InMemoryIndex;
 use puffin_traits::{BuildContext, InFlight, SetupPyStrategy};
 
 use crate::commands::ExitStatus;
@@ -139,6 +140,9 @@ async fn venv_impl(
             FlatIndex::from_entries(entries, tags)
         };
 
+        // Create a shared in-memory index.
+        let index = InMemoryIndex::default();
+
         // Track in-flight downloads, builds, etc., across resolutions.
         let in_flight = InFlight::default();
 
@@ -149,6 +153,7 @@ async fn venv_impl(
             interpreter,
             index_locations,
             &flat_index,
+            &index,
             &in_flight,
             venv.python_executable(),
             SetupPyStrategy::default(),

crates/puffin-dev/src/build.rs

@@ -12,6 +12,7 @@ use puffin_cache::{Cache, CacheArgs};
 use puffin_client::{FlatIndex, RegistryClientBuilder};
 use puffin_dispatch::BuildDispatch;
 use puffin_interpreter::Virtualenv;
+use puffin_resolver::InMemoryIndex;
 use puffin_traits::{BuildContext, BuildKind, InFlight, SetupPyStrategy};
 
 #[derive(Parser)]
@@ -56,6 +57,7 @@ pub(crate) async fn build(args: BuildArgs) -> Result<PathBuf> {
     let client = RegistryClientBuilder::new(cache.clone()).build();
     let index_urls = IndexLocations::default();
     let flat_index = FlatIndex::default();
+    let index = InMemoryIndex::default();
     let setup_py = SetupPyStrategy::default();
     let in_flight = InFlight::default();
 
@@ -65,6 +67,7 @@ pub(crate) async fn build(args: BuildArgs) -> Result<PathBuf> {
         venv.interpreter(),
         &index_urls,
         &flat_index,
+        &index,
         &in_flight,
         venv.python_executable(),
         setup_py,

crates/puffin-dev/src/install_many.rs

@@ -24,7 +24,7 @@ use puffin_distribution::RegistryWheelIndex;
 use puffin_installer::Downloader;
 use puffin_interpreter::Virtualenv;
 use puffin_normalize::PackageName;
-use puffin_resolver::DistFinder;
+use puffin_resolver::{DistFinder, InMemoryIndex};
 use puffin_traits::{BuildContext, InFlight, SetupPyStrategy};
 
 #[derive(Parser)]
@@ -61,6 +61,7 @@ pub(crate) async fn install_many(args: InstallManyArgs) -> Result<()> {
     let client = RegistryClientBuilder::new(cache.clone()).build();
     let index_locations = IndexLocations::default();
     let flat_index = FlatIndex::default();
+    let index = InMemoryIndex::default();
     let setup_py = SetupPyStrategy::default();
     let in_flight = InFlight::default();
     let tags = venv.interpreter().tags()?;
@@ -71,6 +72,7 @@ pub(crate) async fn install_many(args: InstallManyArgs) -> Result<()> {
         venv.interpreter(),
         &index_locations,
         &flat_index,
+        &index,
         &in_flight,
         venv.python_executable(),
         setup_py,

crates/puffin-dev/src/resolve_cli.rs

@@ -16,7 +16,7 @@ use puffin_cache::{Cache, CacheArgs};
 use puffin_client::{FlatIndex, FlatIndexClient, RegistryClientBuilder};
 use puffin_dispatch::BuildDispatch;
 use puffin_interpreter::Virtualenv;
-use puffin_resolver::{Manifest, ResolutionOptions, Resolver};
+use puffin_resolver::{InMemoryIndex, Manifest, ResolutionOptions, Resolver};
 use puffin_traits::{InFlight, SetupPyStrategy};
 
 #[derive(ValueEnum, Default, Clone)]
@@ -65,6 +65,7 @@ pub(crate) async fn resolve_cli(args: ResolveCliArgs) -> Result<()> {
         let entries = client.fetch(index_locations.flat_indexes()).await?;
         FlatIndex::from_entries(entries, venv.interpreter().tags()?)
     };
+    let index = InMemoryIndex::default();
     let in_flight = InFlight::default();
 
     let build_dispatch = BuildDispatch::new(
@@ -73,6 +74,7 @@ pub(crate) async fn resolve_cli(args: ResolveCliArgs) -> Result<()> {
         venv.interpreter(),
         &index_locations,
         &flat_index,
+        &index,
         &in_flight,
         venv.python_executable(),
         SetupPyStrategy::default(),
@@ -89,6 +91,7 @@ pub(crate) async fn resolve_cli(args: ResolveCliArgs) -> Result<()> {
         tags,
         &client,
         &flat_index,
+        &index,
         &build_dispatch,
     );
     let resolution_graph = resolver.resolve().await.with_context(|| {

crates/puffin-dev/src/resolve_many.rs

@@ -20,6 +20,7 @@ use puffin_client::{FlatIndex, RegistryClient, RegistryClientBuilder};
 use puffin_dispatch::BuildDispatch;
 use puffin_interpreter::Virtualenv;
 use puffin_normalize::PackageName;
+use puffin_resolver::InMemoryIndex;
 use puffin_traits::{BuildContext, InFlight, SetupPyStrategy};
 
 #[derive(Parser)]
@@ -75,6 +76,7 @@ pub(crate) async fn resolve_many(args: ResolveManyArgs) -> Result<()> {
     let client = RegistryClientBuilder::new(cache.clone()).build();
     let index_locations = IndexLocations::default();
     let flat_index = FlatIndex::default();
+    let index = InMemoryIndex::default();
     let setup_py = SetupPyStrategy::default();
     let in_flight = InFlight::default();
 
@@ -84,6 +86,7 @@ pub(crate) async fn resolve_many(args: ResolveManyArgs) -> Result<()> {
         venv.interpreter(),
         &index_locations,
         &flat_index,
+        &index,
         &in_flight,
         venv.python_executable(),
         setup_py,

crates/puffin-dispatch/src/lib.rs

@@ -16,7 +16,7 @@ use puffin_cache::Cache;
 use puffin_client::{FlatIndex, RegistryClient};
 use puffin_installer::{Downloader, InstallPlan, Installer, Reinstall, SitePackages};
 use puffin_interpreter::{Interpreter, Virtualenv};
-use puffin_resolver::{Manifest, ResolutionOptions, Resolver};
+use puffin_resolver::{InMemoryIndex, Manifest, ResolutionOptions, Resolver};
 use puffin_traits::{BuildContext, BuildKind, InFlight, SetupPyStrategy};
 
 /// The main implementation of [`BuildContext`], used by the CLI, see [`BuildContext`]
@@ -27,6 +27,7 @@ pub struct BuildDispatch<'a> {
     interpreter: &'a Interpreter,
     index_locations: &'a IndexLocations,
     flat_index: &'a FlatIndex,
+    index: &'a InMemoryIndex,
     in_flight: &'a InFlight,
     base_python: PathBuf,
     setup_py: SetupPyStrategy,
@@ -43,6 +44,7 @@ impl<'a> BuildDispatch<'a> {
         interpreter: &'a Interpreter,
         index_locations: &'a IndexLocations,
         flat_index: &'a FlatIndex,
+        index: &'a InMemoryIndex,
         in_flight: &'a InFlight,
         base_python: PathBuf,
         setup_py: SetupPyStrategy,
@@ -54,6 +56,7 @@ impl<'a> BuildDispatch<'a> {
             interpreter,
             index_locations,
             flat_index,
+            index,
             in_flight,
             base_python,
             setup_py,
@@ -104,6 +107,7 @@ impl<'a> BuildContext for BuildDispatch<'a> {
             tags,
             self.client,
             self.flat_index,
+            self.index,
             self,
         );
         let graph = resolver.resolve().await.with_context(|| {

crates/puffin-resolver/src/lib.rs

@@ -5,7 +5,9 @@ pub use prerelease_mode::PreReleaseMode;
 pub use resolution::{Diagnostic, DisplayResolutionGraph, ResolutionGraph};
 pub use resolution_mode::ResolutionMode;
 pub use resolution_options::ResolutionOptions;
-pub use resolver::{BuildId, Reporter as ResolverReporter, Resolver, ResolverProvider};
+pub use resolver::{
+    BuildId, InMemoryIndex, Reporter as ResolverReporter, Resolver, ResolverProvider,
+};
 
 mod candidate_selector;
 mod error;

crates/puffin-resolver/src/resolver/index.rs

@@ -10,7 +10,7 @@ use crate::version_map::VersionMap;
 
 /// In-memory index of package metadata.
 #[derive(Default)]
-pub(crate) struct Index {
+pub struct InMemoryIndex {
     /// A map from package name to the metadata for that package and the index where the metadata
     /// came from.
     pub(crate) packages: OnceMap<PackageName, VersionMap>,
@@ -24,7 +24,7 @@ pub(crate) struct Index {
     pub(crate) redirects: DashMap<Url, Url>,
 }
 
-impl Index {
+impl InMemoryIndex {
     /// Cancel all waiting tasks.
     ///
     /// Warning: waiting on tasks that have been canceled will cause the index to hang.