distribution-filename: speed up is_compatible

This commit tweaks the representation of `Tags` in order to offer a
faster implementation of `WheelFilename::is_compatible`. We now use a
nested map of tags that lets us avoid looping over every supported
platform tag. As the code comments suggest, that is the essential gain.
We still do not mind looping over the tags in each wheel name since they
tend to be quite small. And pushing our thumb on that side of things can
make things worse overall since it would likely slow down WheelFilename
construction itself.

For micro-benchmarks, we improve considerably for compatibility
checking:

    $ critcmp base test3
    group                                                   base                                    test3
    -----                                                   ----                                    -----
    build_platform_tags/burntsushi-archlinux                1.00     46.2±0.28µs        ? ?/sec     2.48    114.8±0.45µs        ? ?/sec
    wheelname_parsing/flyte-long-compatible                 1.00    624.8±3.31ns   174.0 MB/sec     1.01    629.4±4.30ns   172.7 MB/sec
    wheelname_parsing/flyte-long-incompatible               1.00    743.6±4.23ns   165.4 MB/sec     1.00    746.9±4.62ns   164.7 MB/sec
    wheelname_parsing/flyte-short-compatible                1.00    526.7±4.76ns    54.3 MB/sec     1.01    530.2±5.81ns    54.0 MB/sec
    wheelname_parsing/flyte-short-incompatible              1.00    540.4±4.93ns    60.0 MB/sec     1.01    545.7±5.31ns    59.4 MB/sec
    wheelname_parsing_failure/flyte-long-extension          1.00     13.6±0.13ns     3.2 GB/sec     1.01     13.7±0.14ns     3.2 GB/sec
    wheelname_parsing_failure/flyte-short-extension         1.00     14.0±0.20ns  1160.4 MB/sec     1.01     14.1±0.14ns  1146.5 MB/sec
    wheelname_tag_compatibility/flyte-long-compatible       11.33   159.8±2.79ns   680.5 MB/sec     1.00     14.1±0.23ns     7.5 GB/sec
    wheelname_tag_compatibility/flyte-long-incompatible     237.60 1671.8±37.99ns    73.6 MB/sec    1.00      7.0±0.08ns    17.1 GB/sec
    wheelname_tag_compatibility/flyte-short-compatible      16.07   223.5±8.60ns   128.0 MB/sec     1.00     13.9±0.30ns     2.0 GB/sec
    wheelname_tag_compatibility/flyte-short-incompatible    149.83   628.3±2.13ns    51.6 MB/sec    1.00      4.2±0.10ns     7.6 GB/sec

We do regress slightly on the time it takes for `Tags::new` to run, but
this is somewhat expected. And in absolute terms, 114us is perfectly
acceptable given that it's only executed ~once for each `puffin`
invocation.

Ad hoc benchmarks indicate an overall 25% perf improvement in `puffin
pip-compile` times. This roughly corresponds with how much time
`is_compatible` was taking. Indeed, a profile confirm that it has
virtually disappeared from the profile.

Fixes #157

crates/distribution-filename/src/wheel.rs

@@ -125,15 +125,7 @@ impl Display for WheelFilename {
 impl WheelFilename {
     /// Returns `true` if the wheel is compatible with the given tags.
     pub fn is_compatible(&self, compatible_tags: &Tags) -> bool {
-        for tag in compatible_tags.iter() {
-            if self.python_tag.contains(&tag.0)
-                && self.abi_tag.contains(&tag.1)
-                && self.platform_tag.contains(&tag.2)
-            {
-                return true;
-            }
-        }
-        false
+        compatible_tags.is_compatible(&self.python_tag, &self.abi_tag, &self.platform_tag)
     }
 
     /// Get the tag for this wheel.

crates/platform-tags/Cargo.toml

@@ -10,4 +10,5 @@ authors = { workspace = true }
 license = { workspace = true }
 
 [dependencies]
+fxhash = { workspace = true }
 platform-host = { path = "../platform-host" }

crates/platform-tags/src/lib.rs

@@ -1,14 +1,29 @@
+use fxhash::{FxHashMap, FxHashSet};
+
 use platform_host::{Arch, Os, Platform, PlatformError};
 
-/// A set of compatible tags for a given Python version and platform, in
-/// (`python_tag`, `abi_tag`, `platform_tag`) format.
+/// A set of compatible tags for a given Python version and platform.
+///
+/// Its principle function is to determine whether the tags for a particular
+/// wheel are compatible with the current environment.
 #[derive(Debug)]
-pub struct Tags(Vec<(String, String, String)>);
+pub struct Tags {
+    /// python_tag |--> abi_tag |--> {platform_tag}
+    map: FxHashMap<String, FxHashMap<String, FxHashSet<String>>>,
+}
 
 impl Tags {
     /// Create a new set of tags.
     pub fn new(tags: Vec<(String, String, String)>) -> Self {
-        Self(tags)
+        let mut map = FxHashMap::default();
+        for (py, abi, platform) in tags {
+            map.entry(py.to_string())
+                .or_insert_with(|| FxHashMap::default())
+                .entry(abi.to_string())
+                .or_insert_with(|| FxHashSet::default())
+                .insert(platform.to_string());
+        }
+        Self { map }
     }
 
     /// Returns the compatible tags for the given Python version and platform.
@@ -79,11 +94,41 @@ impl Tags {
             "any".to_string(),
         ));
         tags.sort();
-        Ok(Self(tags))
+        Ok(Self::new(tags))
     }
 
-    pub fn iter(&self) -> impl Iterator<Item = &(String, String, String)> {
-        self.0.iter()
+    /// Returns true when there exists at least one tag for this platform
+    /// whose individal components all appear in each of the slices given.
+    pub fn is_compatible(
+        &self,
+        wheel_python_tags: &[String],
+        wheel_abi_tags: &[String],
+        wheel_platform_tags: &[String],
+    ) -> bool {
+        // NOTE: A typical work-load is a context in which the platform tags
+        // are quite large, but the tags of a wheel are quite small. It is
+        // common, for example, for the lengths of the slices given to all be
+        // 1. So while the looping here might look slow, the key thing we want
+        // to avoid is looping over all of the platform tags. We avoid that
+        // with hashmap lookups.
+
+        let pythons = &self.map;
+        for wheel_py in wheel_python_tags.iter() {
+            let Some(abis) = pythons.get(wheel_py) else {
+                continue;
+            };
+            for wheel_abi in wheel_abi_tags.iter() {
+                let Some(platforms) = abis.get(wheel_abi) else {
+                    continue;
+                };
+                for wheel_platform in wheel_platform_tags.iter() {
+                    if platforms.contains(wheel_platform) {
+                        return true;
+                    }
+                }
+            }
+        }
+        false
     }
 }