Store unsupported tags in wheel filename (#10665)

## Summary

We can retain the small-size advantage of our new tags by moving the
"unknown tag" case into `WheelTagLarge`. This ensures that we can still
represent unknown tags, but avoid paying the cost for them.

crates/uv-distribution-filename/src/lib.rs

@@ -7,7 +7,7 @@ pub use build_tag::{BuildTag, BuildTagError};
 pub use egg::{EggInfoFilename, EggInfoFilenameError};
 pub use extension::{DistExtension, ExtensionError, SourceDistExtension};
 pub use source_dist::{SourceDistFilename, SourceDistFilenameError};
-pub use wheel::{TagSet, WheelFilename, WheelFilenameError};
+pub use wheel::{WheelFilename, WheelFilenameError};
 
 mod build_tag;
 mod egg;
@@ -15,6 +15,7 @@ mod extension;
 mod source_dist;
 mod splitter;
 mod wheel;
+mod wheel_tag;
 
 #[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
 pub enum DistFilename {
@@ -97,13 +98,9 @@ impl Display for DistFilename {
 #[cfg(test)]
 mod tests {
     use crate::WheelFilename;
-    use uv_platform_tags::{AbiTag, LanguageTag, PlatformTag};
 
     #[test]
     fn wheel_filename_size() {
-        assert_eq!(size_of::<WheelFilename>(), 72);
-        assert_eq!(size_of::<LanguageTag>(), 16);
-        assert_eq!(size_of::<AbiTag>(), 16);
-        assert_eq!(size_of::<PlatformTag>(), 16);
+        assert_eq!(size_of::<WheelFilename>(), 48);
     }
 }

crates/uv-distribution-filename/src/snapshots/uv_distribution_filename__wheel__tests__ok_build_tag.snap

@@ -28,6 +28,7 @@ Ok(
                 platform_tag: [
                     Any,
                 ],
+                repr: "202206090410-py3-none-any",
             },
         },
     },

crates/uv-distribution-filename/src/snapshots/uv_distribution_filename__wheel__tests__ok_multiple_tags.snap

@@ -38,6 +38,7 @@ Ok(
                         arch: X86_64,
                     },
                 ],
+                repr: "cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64",
             },
         },
     },

crates/uv-distribution-filename/src/wheel.rs

@@ -14,6 +14,7 @@ use uv_platform_tags::{
 };
 
 use crate::splitter::MemchrSplitter;
+use crate::wheel_tag::{WheelTag, WheelTagLarge, WheelTagSmall};
 use crate::{BuildTag, BuildTagError};
 
 #[derive(
@@ -177,7 +178,7 @@ impl WheelFilename {
             ));
         };
 
-        let (name, version, build_tag, python_tag, abi_tag, platform_tag, is_large) =
+        let (name, version, build_tag, python_tag, abi_tag, platform_tag, is_small) =
             if let Some(platform_tag) = splitter.next() {
                 if splitter.next().is_some() {
                     return Err(WheelFilenameError::InvalidWheelFileName(
@@ -193,7 +194,7 @@ impl WheelFilename {
                     &stem[abi_tag_or_platform_tag + 1..platform_tag],
                     &stem[platform_tag + 1..],
                     // Always take the slow path if a build tag is present.
-                    true,
+                    false,
                 )
             } else {
                 (
@@ -206,7 +207,7 @@ impl WheelFilename {
                     // Determine whether any of the tag types contain a period, which would indicate
                     // that at least one of the tag types includes multiple tags (which in turn
                     // necessitates taking the slow path).
-                    memchr(b'.', stem[build_tag_or_python_tag..].as_bytes()).is_some(),
+                    memchr(b'.', stem[build_tag_or_python_tag..].as_bytes()).is_none(),
                 )
             };
 
@@ -221,44 +222,38 @@ impl WheelFilename {
             })
             .transpose()?;
 
-        let tags = if is_large {
+        let tags = if let Some(small) = is_small
+            .then(|| {
+                Some(WheelTagSmall {
+                    python_tag: LanguageTag::from_str(python_tag).ok()?,
+                    abi_tag: AbiTag::from_str(abi_tag).ok()?,
+                    platform_tag: PlatformTag::from_str(platform_tag).ok()?,
+                })
+            })
+            .flatten()
+        {
+            WheelTag::Small { small }
+        } else {
+            // Store the plaintext representation of the tags.
+            let repr = &stem[build_tag_or_python_tag + 1..];
             WheelTag::Large {
                 large: Box::new(WheelTagLarge {
                     build_tag,
                     python_tag: MemchrSplitter::split(python_tag, b'.')
                         .map(LanguageTag::from_str)
-                        .collect::<Result<_, _>>()
-                        .map_err(|err| {
-                            WheelFilenameError::InvalidLanguageTag(filename.to_string(), err)
-                        })?,
+                        .filter_map(Result::ok)
+                        .collect(),
                     abi_tag: MemchrSplitter::split(abi_tag, b'.')
                         .map(AbiTag::from_str)
-                        .collect::<Result<_, _>>()
-                        .map_err(|err| {
-                            WheelFilenameError::InvalidAbiTag(filename.to_string(), err)
-                        })?,
+                        .filter_map(Result::ok)
+                        .collect(),
                     platform_tag: MemchrSplitter::split(platform_tag, b'.')
                         .map(PlatformTag::from_str)
-                        .collect::<Result<_, _>>()
-                        .map_err(|err| {
-                            WheelFilenameError::InvalidPlatformTag(filename.to_string(), err)
-                        })?,
+                        .filter_map(Result::ok)
+                        .collect(),
+                    repr: repr.into(),
                 }),
             }
-        } else {
-            WheelTag::Small {
-                small: WheelTagSmall {
-                    python_tag: LanguageTag::from_str(python_tag).map_err(|err| {
-                        WheelFilenameError::InvalidLanguageTag(filename.to_string(), err)
-                    })?,
-                    abi_tag: AbiTag::from_str(abi_tag).map_err(|err| {
-                        WheelFilenameError::InvalidAbiTag(filename.to_string(), err)
-                    })?,
-                    platform_tag: PlatformTag::from_str(platform_tag).map_err(|err| {
-                        WheelFilenameError::InvalidPlatformTag(filename.to_string(), err)
-                    })?,
-                },
-            }
         };
 
         Ok(Self {
@@ -311,124 +306,6 @@ impl Serialize for WheelFilename {
     }
 }
 
-/// A [`SmallVec`] type for storing tags.
-///
-/// Wheels tend to include a single language, ABI, and platform tag, so we use a [`SmallVec`] with a
-/// capacity of 1 to optimize for this common case.
-pub type TagSet<T> = smallvec::SmallVec<[T; 3]>;
-
-/// The portion of the wheel filename following the name and version: the optional build tag, along
-/// with the Python tag(s), ABI tag(s), and platform tag(s).
-///
-/// Most wheels consist of a single Python, ABI, and platform tag (and no build tag). We represent
-/// such wheels with [`WheelTagSmall`], a variant with a smaller memory footprint and (generally)
-/// zero allocations. The [`WheelTagLarge`] variant is used for wheels with multiple tags and/or a
-/// build tag.
-#[derive(
-    Debug,
-    Clone,
-    Eq,
-    PartialEq,
-    Ord,
-    PartialOrd,
-    Hash,
-    rkyv::Archive,
-    rkyv::Deserialize,
-    rkyv::Serialize,
-)]
-#[rkyv(derive(Debug))]
-enum WheelTag {
-    Small { small: WheelTagSmall },
-    Large { large: Box<WheelTagLarge> },
-}
-
-impl Display for WheelTag {
-    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
-        match self {
-            Self::Small { small } => write!(f, "{small}"),
-            Self::Large { large } => write!(f, "{large}"),
-        }
-    }
-}
-
-#[derive(
-    Debug,
-    Clone,
-    Eq,
-    PartialEq,
-    Ord,
-    PartialOrd,
-    Hash,
-    rkyv::Archive,
-    rkyv::Deserialize,
-    rkyv::Serialize,
-)]
-#[rkyv(derive(Debug))]
-#[allow(clippy::struct_field_names)]
-struct WheelTagSmall {
-    python_tag: LanguageTag,
-    abi_tag: AbiTag,
-    platform_tag: PlatformTag,
-}
-
-impl Display for WheelTagSmall {
-    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
-        write!(
-            f,
-            "{}-{}-{}",
-            self.python_tag, self.abi_tag, self.platform_tag
-        )
-    }
-}
-
-#[derive(
-    Debug,
-    Clone,
-    Eq,
-    PartialEq,
-    Ord,
-    PartialOrd,
-    Hash,
-    rkyv::Archive,
-    rkyv::Deserialize,
-    rkyv::Serialize,
-)]
-#[rkyv(derive(Debug))]
-#[allow(clippy::struct_field_names)]
-pub struct WheelTagLarge {
-    build_tag: Option<BuildTag>,
-    python_tag: TagSet<LanguageTag>,
-    abi_tag: TagSet<AbiTag>,
-    platform_tag: TagSet<PlatformTag>,
-}
-
-impl Display for WheelTagLarge {
-    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
-        if let Some(build_tag) = &self.build_tag {
-            write!(f, "{build_tag}-")?;
-        }
-        write!(
-            f,
-            "{}-{}-{}",
-            self.python_tag
-                .iter()
-                .map(ToString::to_string)
-                .collect::<Vec<_>>()
-                .join("."),
-            self.abi_tag
-                .iter()
-                .map(ToString::to_string)
-                .collect::<Vec<_>>()
-                .join("."),
-            self.platform_tag
-                .iter()
-                .map(ToString::to_string)
-                .collect::<Vec<_>>()
-                .join("."),
-        )
-    }
-}
-
 #[derive(Error, Debug)]
 pub enum WheelFilenameError {
     #[error("The wheel filename \"{0}\" is invalid: {1}")]

crates/uv-distribution-filename/src/wheel_tag.rs

@@ -0,0 +1,115 @@
+use std::fmt::{Display, Formatter};
+
+use crate::BuildTag;
+use uv_platform_tags::{AbiTag, LanguageTag, PlatformTag};
+use uv_small_str::SmallString;
+
+/// A [`SmallVec`] type for storing tags.
+///
+/// Wheels tend to include a single language, ABI, and platform tag, so we use a [`SmallVec`] with a
+/// capacity of 1 to optimize for this common case.
+pub(crate) type TagSet<T> = smallvec::SmallVec<[T; 3]>;
+
+/// The portion of the wheel filename following the name and version: the optional build tag, along
+/// with the Python tag(s), ABI tag(s), and platform tag(s).
+///
+/// Most wheels consist of a single Python, ABI, and platform tag (and no build tag). We represent
+/// such wheels with [`WheelTagSmall`], a variant with a smaller memory footprint and (generally)
+/// zero allocations. The [`WheelTagLarge`] variant is used for wheels with multiple tags, a build
+/// tag, or an unsupported tag (i.e., a tag that can't be represented by [`LanguageTag`],
+/// [`AbiTag`], or [`PlatformTag`]). (Unsupported tags are filtered out, but retained in the display
+/// representation of [`WheelTagLarge`].)
+#[derive(
+    Debug,
+    Clone,
+    Eq,
+    PartialEq,
+    Ord,
+    PartialOrd,
+    Hash,
+    rkyv::Archive,
+    rkyv::Deserialize,
+    rkyv::Serialize,
+)]
+#[rkyv(derive(Debug))]
+pub(crate) enum WheelTag {
+    Small { small: WheelTagSmall },
+    Large { large: Box<WheelTagLarge> },
+}
+
+impl Display for WheelTag {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        match self {
+            Self::Small { small } => write!(f, "{small}"),
+            Self::Large { large } => write!(f, "{large}"),
+        }
+    }
+}
+
+#[derive(
+    Debug,
+    Clone,
+    Eq,
+    PartialEq,
+    Ord,
+    PartialOrd,
+    Hash,
+    rkyv::Archive,
+    rkyv::Deserialize,
+    rkyv::Serialize,
+)]
+#[rkyv(derive(Debug))]
+#[allow(clippy::struct_field_names)]
+pub(crate) struct WheelTagSmall {
+    /// The Python tag, e.g., `py3` in `1.2.3-py3-none-any`.
+    pub(crate) python_tag: LanguageTag,
+    /// The ABI tag, e.g., `none` in `1.2.3-py3-none-any`.
+    pub(crate) abi_tag: AbiTag,
+    /// The platform tag, e.g., `none` in `1.2.3-py3-none-any`.
+    pub(crate) platform_tag: PlatformTag,
+}
+
+impl Display for WheelTagSmall {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        write!(
+            f,
+            "{}-{}-{}",
+            self.python_tag, self.abi_tag, self.platform_tag
+        )
+    }
+}
+
+#[derive(
+    Debug,
+    Clone,
+    Eq,
+    PartialEq,
+    Ord,
+    PartialOrd,
+    Hash,
+    rkyv::Archive,
+    rkyv::Deserialize,
+    rkyv::Serialize,
+)]
+#[rkyv(derive(Debug))]
+#[allow(clippy::struct_field_names)]
+pub(crate) struct WheelTagLarge {
+    /// The optional build tag, e.g., `73` in `1.2.3-73-py3-none-any`.
+    pub(crate) build_tag: Option<BuildTag>,
+    /// The Python tag(s), e.g., `py3` in `1.2.3-73-py3-none-any`.
+    pub(crate) python_tag: TagSet<LanguageTag>,
+    /// The ABI tag(s), e.g., `none` in `1.2.3-73-py3-none-any`.
+    pub(crate) abi_tag: TagSet<AbiTag>,
+    /// The platform tag(s), e.g., `none` in `1.2.3-73-py3-none-any`.
+    pub(crate) platform_tag: TagSet<PlatformTag>,
+    /// The string representation of the tag.
+    ///
+    /// Preserves any unsupported tags that were filtered out when parsing the wheel filename.
+    pub(crate) repr: SmallString,
+}
+
+impl Display for WheelTagLarge {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        write!(f, "{}", self.repr)
+    }
+}