Reduce memory usage for Path values with a single Ident, or where we can reconstruct the Span

Cargo.lock

@@ -3599,6 +3599,7 @@ name = "rustc_ast"
 version = "0.0.0"
 dependencies = [
  "bitflags",
+ "either",
  "memchr",
  "rustc-literal-escaper",
  "rustc_ast_ir",
@@ -3758,6 +3759,7 @@ dependencies = [
 name = "rustc_builtin_macros"
 version = "0.0.0"
 dependencies = [
+ "itertools",
  "rustc_ast",
  "rustc_ast_pretty",
  "rustc_attr_ir",
@@ -4589,6 +4591,7 @@ name = "rustc_parse"
 version = "0.0.0"
 dependencies = [
  "bitflags",
+ "itertools",
  "rustc-literal-escaper",
  "rustc_ast",
  "rustc_ast_pretty",

compiler/rustc_ast/Cargo.toml

@@ -6,6 +6,7 @@ edition = "2024"
 [dependencies]
 # tidy-alphabetical-start
 bitflags = "2.4.1"
+either = "1.15.0"
 memchr = "2.7.6"
 rustc-literal-escaper = "0.0.8"
 rustc_ast_ir = { path = "../rustc_ast_ir" }

compiler/rustc_ast/src/ast.rs

@@ -18,11 +18,14 @@
 //! - [`Attribute`]: Metadata associated with item.
 //! - [`UnOp`], [`BinOp`], and [`BinOpKind`]: Unary and binary operators.
 
+// ignore-tidy-file-filelength
+
 use std::borrow::{Borrow, Cow};
-use std::{cmp, fmt};
+use std::{cmp, fmt, iter};
 
 pub use GenericArgs::*;
 pub use UnsafeSource::*;
+use either::Either;
 pub use rustc_ast_ir::{FloatTy, IntTy, Movability, Mutability, Pinnedness, UintTy};
 use rustc_data_structures::packed::Pu128;
 use rustc_data_structures::stable_hash::{StableHash, StableHashCtxt, StableHasher};
@@ -85,72 +88,222 @@ impl fmt::Display for Lifetime {
 
 /// A "Path" is essentially Rust's notion of a name.
 ///
-/// It's represented as a sequence of identifiers,
-/// along with a bunch of supporting information.
+/// We separate the common case of a single identifier (e.g. `x`) from the general case of a
+/// sequence of identifiers that might also have generics attached (e.g. `std::cmp::PartialEq`,
+/// `Vec::<T>::new`).
 ///
-/// E.g., `std::cmp::PartialEq`.
+/// Canonical form is desirable but not mandatory. Single idents may use General if their span
+/// differs from the ident's span, or occasionally simply because they didn't get canonicalized.
 #[derive(Clone, Encodable, Decodable, Debug, Walkable)]
-pub struct Path {
-    pub span: Span,
-    /// The segments in the path: the things separated by `::`.
-    /// Global paths begin with `kw::PathRoot`.
-    pub segments: ThinVec<PathSegment>,
+pub enum Path {
+    /// The common case of a single identifier (e.g. `x`)
+    Ident { ident: Ident, id: NodeId },
+    /// The common case of a general path, which must have at least one segment, and the span can be
+    /// trivially reconstructed from the spans of the first and last segments.
+    NoSpan {
+        /// The segments in the path: the things separated by `::`. Global paths begin with
+        /// `kw::PathRoot`.
+        segments: ThinVec<PathSegment>,
+    },
+    /// A fully general path where the path's span differs (e.g. from recovery). Boxed to avoid
+    /// making `Path` larger.
+    General(Box<(ThinVec<PathSegment>, Span)>),
 }
 
 // Succeeds if the path has a single segment that is arg-free and matches the given symbol.
 impl PartialEq<Symbol> for Path {
     #[inline]
     fn eq(&self, name: &Symbol) -> bool {
-        if let [segment] = self.segments.as_ref()
-            && segment == name
-        {
-            true
-        } else {
-            false
-        }
+        self.as_single_argless_name() == Some(*name)
     }
 }
 
 // Succeeds if the path has segments that are arg-free and match the given symbols.
 impl PartialEq<&[Symbol]> for Path {
     #[inline]
     fn eq(&self, names: &&[Symbol]) -> bool {
-        self.segments.iter().eq(*names)
+        match self {
+            Path::Ident { ident, .. } => {
+                let [name] = names else {
+                    return false;
+                };
+                ident.name == *name
+            }
+            Path::NoSpan { segments } | Path::General((segments, _)) => segments.iter().eq(*names),
+        }
     }
 }
 
 impl StableHash for Path {
     fn stable_hash<Hcx: StableHashCtxt>(&self, hcx: &mut Hcx, hasher: &mut StableHasher) {
-        self.segments.len().stable_hash(hcx, hasher);
-        for segment in &self.segments {
-            segment.ident.stable_hash(hcx, hasher);
-        }
+        self.num_segments().stable_hash(hcx, hasher);
+        self.iter_idents().for_each(|ident| ident.stable_hash(hcx, hasher));
     }
 }
 
 impl Path {
     /// Convert a span and an identifier to the corresponding
     /// one-segment path.
     pub fn from_ident(ident: Ident) -> Path {
-        Path { segments: thin_vec![PathSegment::from_ident(ident)], span: ident.span }
+        Path::Ident { ident, id: DUMMY_NODE_ID }
+    }
+
+    /// Convert a set of segments and a span to the corresponding `Path`.
+    #[inline]
+    pub fn from_segments(segments: ThinVec<PathSegment>, span: Span) -> Path {
+        if let [segment] = segments.as_slice()
+            && segment.ident.span == span
+            && segment.args.is_none()
+        {
+            Path::Ident { ident: segment.ident, id: segment.id }
+        } else if !segments.is_empty() && segments_span(&segments) == span {
+            Path::NoSpan { segments }
+        } else {
+            Path::General(Box::new((segments, span)))
+        }
+    }
+
+    #[inline]
+    pub fn span(&self) -> Span {
+        match self {
+            Path::Ident { ident, .. } => ident.span,
+            Path::NoSpan { segments } => segments_span(segments),
+            Path::General((_, span)) => *span,
+        }
+    }
+
+    pub fn path_root(span: Span) -> Self {
+        Path::from_ident(Ident::new(kw::PathRoot, span))
     }
 
     pub fn is_global(&self) -> bool {
-        self.segments.first().is_some_and(|segment| segment.ident.name == kw::PathRoot)
+        self.iter_idents().next().is_some_and(|i| i.name == kw::PathRoot)
     }
 
     /// Checks if this path is just a simple one-word `PATH` - i.e. the inverse of
     /// [`Path::from_ident`]
+    #[inline]
     pub fn is_single_argless_ident(&self) -> bool {
-        self.segments.len() == 1 && self.segments[0].args.is_none()
+        self.as_single_argless_ident().is_some()
     }
 
     /// The inverse of [`Path::from_ident`] - if this path is just a simple one-word `PATH`
+    #[inline]
     pub fn as_single_argless_ident(&self) -> Option<Ident> {
-        self.is_single_argless_ident().then(|| self.segments[0].ident)
+        // This can't *exclusively* handle the `Path::Ident` case, because a single ident can use
+        // `Path::General` if the path span differs from the ident span, or `Path::NoSpan` if it
+        // isn't in canonical form.
+        match self {
+            Path::Ident { ident, .. } => Some(*ident),
+            Path::NoSpan { segments } | Path::General((segments, _)) => {
+                let [segment] = segments.as_ref() else {
+                    return None;
+                };
+                if segment.args.is_some() {
+                    return None;
+                }
+                Some(segment.ident)
+            }
+        }
+    }
+
+    #[inline]
+    pub fn as_single_argless_name(&self) -> Option<Symbol> {
+        self.as_single_argless_ident().map(|i| i.name)
+    }
+
+    #[inline]
+    pub fn is_empty(&self) -> bool {
+        match self {
+            Path::General((segments, _)) => segments.is_empty(),
+            // Path::NoSpan always has non-empty segments. Path::Ident is non-empty by definition.
+            _ => false,
+        }
+    }
+
+    #[inline]
+    pub fn num_segments(&self) -> usize {
+        match self {
+            Path::Ident { .. } => 1,
+            Path::NoSpan { segments } | Path::General((segments, _)) => segments.len(),
+        }
+    }
+
+    #[inline]
+    pub fn iter_idents(&self) -> impl DoubleEndedIterator<Item = &Ident> + ExactSizeIterator {
+        match self {
+            Path::Ident { ident, .. } => Either::Left(iter::once(ident)),
+            Path::NoSpan { segments } | Path::General((segments, _)) => {
+                Either::Right(segments.iter().map(|s| &s.ident))
+            }
+        }
+    }
+
+    #[inline]
+    pub fn iter_segments(
+        &self,
+    ) -> impl DoubleEndedIterator<Item = PathSegmentRef<'_>> + ExactSizeIterator {
+        match self {
+            &Path::Ident { ref ident, id } => {
+                Either::Left(iter::once(PathSegmentRef { ident, id, args: None }))
+            }
+            Path::NoSpan { segments } | Path::General((segments, _)) => {
+                Either::Right(segments.iter().map(PathSegment::as_ref))
+            }
+        }
+    }
+
+    #[inline]
+    pub fn last_segment(&self) -> Option<PathSegmentRef<'_>> {
+        match self {
+            &Path::Ident { ref ident, id } => Some(PathSegmentRef { ident, id, args: None }),
+            Path::NoSpan { segments } | Path::General((segments, _)) => {
+                segments.last().map(PathSegment::as_ref)
+            }
+        }
+    }
+
+    #[inline]
+    pub fn last_ident(&self) -> Option<Ident> {
+        match self {
+            &Path::Ident { ident, .. } => Some(ident),
+            Path::NoSpan { segments } | Path::General((segments, _)) => {
+                segments.last().map(|s| s.ident)
+            }
+        }
+    }
+
+    pub fn force_general_mut(&mut self) -> (&mut ThinVec<PathSegment>, &mut Span) {
+        match self {
+            Path::General((segments, span)) => return (segments, span),
+            Path::NoSpan { segments } => {
+                let span = segments_span(segments);
+                *self = Path::General(Box::new((std::mem::take(segments), span)))
+            }
+            &mut Path::Ident { ident, id } => {
+                *self = Path::General(Box::new((
+                    thin_vec![PathSegment { ident, id, args: None }],
+                    ident.span,
+                )));
+            }
+        }
+        match self {
+            Path::General((segments, span)) => (segments, span),
+            _ => unreachable!(),
+        }
     }
 }
 
+#[inline]
+fn segments_span(segments: &[PathSegment]) -> Span {
+    let last = segments.last().unwrap();
+    let last_span = match last.args.as_deref() {
+        Some(last_args) => last_args.span(),
+        None => last.ident.span,
+    };
+    segments[0].ident.span.to(last_span)
+}
+
 /// Joins multiple symbols with "::" into a path, e.g. "a::b::c". If the first
 /// segment is `kw::PathRoot` it will be printed as empty, e.g. "::b::c".
 ///
@@ -240,12 +393,34 @@ impl PathSegment {
         PathSegment::from_ident(Ident::new(kw::PathRoot, span))
     }
 
+    #[inline]
     pub fn span(&self) -> Span {
         match &self.args {
             Some(args) => self.ident.span.to(args.span()),
             None => self.ident.span,
         }
     }
+
+    pub fn as_ref(&self) -> PathSegmentRef<'_> {
+        PathSegmentRef { ident: &self.ident, id: self.id, args: self.args.as_deref() }
+    }
+}
+
+#[derive(Copy, Clone, Debug)]
+pub struct PathSegmentRef<'a> {
+    pub ident: &'a Ident,
+    pub id: NodeId,
+    pub args: Option<&'a GenericArgs>,
+}
+
+impl PathSegmentRef<'_> {
+    #[inline]
+    pub fn span(&self) -> Span {
+        match self.args {
+            Some(args) => self.ident.span.to(args.span()),
+            None => self.ident.span,
+        }
+    }
 }
 
 /// The generic arguments and associated item constraints of a path segment.
@@ -2060,7 +2235,11 @@ pub struct MacCall {
 
 impl MacCall {
     pub fn span(&self) -> Span {
-        self.path.span.to(self.args.dspan.entire())
+        let first_span = match &self.path {
+            Path::NoSpan { segments } => segments[0].ident.span,
+            _ => self.path.span(),
+        };
+        first_span.to(self.args.dspan.entire())
     }
 }
 
@@ -2588,14 +2767,7 @@ impl TyKind {
     }
 
     pub fn is_simple_path(&self) -> Option<Symbol> {
-        if let TyKind::Path(None, Path { segments, .. }) = &self
-            && let [segment] = &segments[..]
-            && segment.args.is_none()
-        {
-            Some(segment.ident.name)
-        } else {
-            None
-        }
+        if let TyKind::Path(None, path) = &self { path.as_single_argless_name() } else { None }
     }
 
     /// Returns `true` if this type is considered a scalar primitive (e.g.,
@@ -3340,7 +3512,7 @@ impl UseTree {
         match self.kind {
             UseTreeKind::Simple(Some(rename)) => rename,
             UseTreeKind::Simple(None) => {
-                self.prefix.segments.last().expect("empty prefix in a simple import").ident
+                self.prefix.last_ident().expect("empty prefix in a simple import")
             }
             _ => panic!("`UseTree::ident` can only be used on a simple import"),
         }
@@ -3350,7 +3522,10 @@ impl UseTree {
     /// closing `}` or nested spans, `*` of glob spans or the end of the
     /// identifier of simple spans.
     pub fn span(&self) -> Span {
-        self.prefix.span.to(self.hi_span())
+        match self.kind {
+            UseTreeKind::Simple(None) => self.prefix.span(),
+            _ => self.prefix.span().to(self.hi_span()),
+        }
     }
 
     /// Returns the trailing element's span. So for a nested
@@ -3360,7 +3535,7 @@ impl UseTree {
     /// path if no rename is specified.
     pub fn hi_span(&self) -> Span {
         match self.kind {
-            UseTreeKind::Simple(None) => self.prefix.span,
+            UseTreeKind::Simple(None) => self.prefix.span(),
             UseTreeKind::Simple(Some(name)) => name.span,
             UseTreeKind::Nested { span, .. } => span,
             UseTreeKind::Glob(span) => span,
@@ -3994,7 +4169,7 @@ pub struct Delegation {
 
 impl Delegation {
     pub fn last_segment_span(&self) -> Span {
-        self.path.segments.last().unwrap().ident.span
+        self.path.last_ident().unwrap().span
     }
 }