rustc_span: Optimize span parent get/set methods

petrochenkov · petrochenkov · commit 90b72b581bbc · 2024-06-16T14:01:55.000+03:00
diff --git a/compiler/rustc_span/src/lib.rs b/compiler/rustc_span/src/lib.rs
@@ -525,8 +525,9 @@ impl SpanData {
     fn with_ctxt(&self, ctxt: SyntaxContext) -> Span {
         Span::new(self.lo, self.hi, ctxt, self.parent)
     }
+    /// Avoid if possible, `Span::with_parent` should be preferred.
     #[inline]
-    pub fn with_parent(&self, parent: Option<LocalDefId>) -> Span {
+    fn with_parent(&self, parent: Option<LocalDefId>) -> Span {
         Span::new(self.lo, self.hi, self.ctxt, parent)
     }
     /// Returns `true` if this is a dummy span with any hygienic context.
@@ -580,14 +581,6 @@ impl Span {
     pub fn with_ctxt(self, ctxt: SyntaxContext) -> Span {
         self.map_ctxt(|_| ctxt)
     }
-    #[inline]
-    pub fn parent(self) -> Option<LocalDefId> {
-        self.data().parent
-    }
-    #[inline]
-    pub fn with_parent(self, ctxt: Option<LocalDefId>) -> Span {
-        self.data().with_parent(ctxt)
-    }
 
     #[inline]
     pub fn is_visible(self, sm: &SourceMap) -> bool {
diff --git a/compiler/rustc_span/src/span_encoding.rs b/compiler/rustc_span/src/span_encoding.rs
@@ -401,6 +401,63 @@ impl Span {
             }),
         }
     }
+
+    // For optimization we are interested in cases in which the parent is inline and the parent
+    // update doesn't change format. All non-inline or format changing scenarios require accessing
+    // interner and can fall back to `Span::new`.
+    #[inline]
+    pub fn with_parent(self, parent: Option<LocalDefId>) -> Span {
+        let data = match_span_kind! {
+            self,
+            InlineCtxt(span) => {
+                // Only if the new parent is `None` the format will be preserved.
+                if parent.is_none() {
+                    return self;
+                }
+                span.data()
+            },
+            InlineParent(span) => {
+                // Any `Some` small new parent will preserve the format.
+                // `None` may change the format to `InlineCtxt` if ctxt and len are small enough.
+                if let Some(parent) = parent
+                    && let parent32 = parent.local_def_index.as_u32()
+                    && parent32 <= MAX_CTXT {
+                    let old_parent = LocalDefId { local_def_index: DefIndex::from_u32(span.parent as u32) };
+                    // FIXME: Is this tracking necessary?
+                    (*SPAN_TRACK)(old_parent);
+                    return InlineParent::span(span.lo, span.len_with_tag, parent32 as u16);
+                }
+                span.data()
+            },
+            PartiallyInterned(span) => span.data(),
+            Interned(span) => span.data(),
+        };
+
+        // We could not keep the span in the same inline format, fall back to the complete logic.
+        if let Some(old_parent) = data.parent {
+            // FIXME: Is this tracking necessary?
+            (*SPAN_TRACK)(old_parent);
+        }
+        data.with_parent(parent)
+    }
+
+    #[inline]
+    pub fn parent(self) -> Option<LocalDefId> {
+        let interned_parent =
+            |index: u32| with_span_interner(|interner| interner.spans[index as usize].parent);
+        let parent = match_span_kind! {
+            self,
+            InlineCtxt(_span) => return None,
+            InlineParent(span) => Some(LocalDefId { local_def_index: DefIndex::from_u32(span.parent as u32) }),
+            PartiallyInterned(span) => interned_parent(span.index),
+            Interned(span) => interned_parent(span.index),
+        };
+        if let Some(parent) = parent {
+            // FIXME: Is this tracking necessary?
+            (*SPAN_TRACK)(parent);
+        }
+        parent
+    }
 }
 
 #[derive(Default)]
