Auto merge of #125017 - petrochenkov:upctxt, r=<try>

rustc_span: Optimize syntax context updates in spans

By updating inline contexts in place, without touching span interner, when possible.

TODO: Run benchmarks and cleanup code if the results are good.

Author: bors

compiler/rustc_expand/src/mbe/transcribe.rs

@@ -30,11 +30,11 @@ impl MutVisitor for Marker {
         // it's some advanced case with macro-generated macros. So if we cache the marked version
         // of that context once, we'll typically have a 100% cache hit rate after that.
         let Marker(expn_id, transparency, ref mut cache) = *self;
-        let data = span.data();
-        let marked_ctxt = *cache
-            .entry(data.ctxt)
-            .or_insert_with(|| data.ctxt.apply_mark(expn_id.to_expn_id(), transparency));
-        *span = data.with_ctxt(marked_ctxt);
+        span.update_ctxt(|ctxt| {
+            *cache
+                .entry(ctxt)
+                .or_insert_with(|| ctxt.apply_mark(expn_id.to_expn_id(), transparency))
+        });
     }
 }
 

compiler/rustc_span/src/lib.rs

@@ -520,7 +520,7 @@ impl SpanData {
         Span::new(self.lo, hi, self.ctxt, self.parent)
     }
     #[inline]
-    pub fn with_ctxt(&self, ctxt: SyntaxContext) -> Span {
+    fn with_ctxt(&self, ctxt: SyntaxContext) -> Span {
         Span::new(self.lo, self.hi, ctxt, self.parent)
     }
     #[inline]
@@ -575,8 +575,9 @@ impl Span {
         self.data().with_hi(hi)
     }
     #[inline]
-    pub fn with_ctxt(self, ctxt: SyntaxContext) -> Span {
-        self.data_untracked().with_ctxt(ctxt)
+    pub fn with_ctxt(mut self, ctxt: SyntaxContext) -> Span {
+        self.update_ctxt(|_| ctxt);
+        self
     }
     #[inline]
     pub fn parent(self) -> Option<LocalDefId> {
@@ -1050,9 +1051,9 @@ impl Span {
     }
 
     #[inline]
-    pub fn apply_mark(self, expn_id: ExpnId, transparency: Transparency) -> Span {
-        let span = self.data();
-        span.with_ctxt(span.ctxt.apply_mark(expn_id, transparency))
+    pub fn apply_mark(mut self, expn_id: ExpnId, transparency: Transparency) -> Span {
+        self.update_ctxt(|ctxt| ctxt.apply_mark(expn_id, transparency));
+        self
     }
 
     #[inline]
@@ -1100,15 +1101,15 @@ impl Span {
     }
 
     #[inline]
-    pub fn normalize_to_macros_2_0(self) -> Span {
-        let span = self.data();
-        span.with_ctxt(span.ctxt.normalize_to_macros_2_0())
+    pub fn normalize_to_macros_2_0(mut self) -> Span {
+        self.update_ctxt(|ctxt| ctxt.normalize_to_macros_2_0());
+        self
     }
 
     #[inline]
-    pub fn normalize_to_macro_rules(self) -> Span {
-        let span = self.data();
-        span.with_ctxt(span.ctxt.normalize_to_macro_rules())
+    pub fn normalize_to_macro_rules(mut self) -> Span {
+        self.update_ctxt(|ctxt| ctxt.normalize_to_macro_rules());
+        self
     }
 }
 

compiler/rustc_span/src/span_encoding.rs

@@ -136,13 +136,13 @@ impl Span {
         }
 
         // Partially-interned or fully-interned format.
-        let index =
-            with_span_interner(|interner| interner.intern(&SpanData { lo, hi, ctxt, parent }));
-        let ctxt_or_parent_or_marker = if ctxt2 <= MAX_CTXT {
-            ctxt2 as u16 // partially-interned
+        let (ctxt_or_parent_or_marker, ctxt) = if ctxt2 <= MAX_CTXT {
+            (ctxt2 as u16, SyntaxContext::from_u32(u32::MAX)) // partially-interned
         } else {
-            CTXT_INTERNED_MARKER // fully-interned
+            (CTXT_INTERNED_MARKER, ctxt) // fully-interned
         };
+        let index =
+            with_span_interner(|interner| interner.intern(&SpanData { lo, hi, ctxt, parent }));
         Span {
             lo_or_index: index,
             len_with_tag_or_marker: BASE_LEN_INTERNED_MARKER,
@@ -193,7 +193,11 @@ impl Span {
             // the interned value contains all the data, so we don't need to
             // distinguish them.
             let index = self.lo_or_index;
-            with_span_interner(|interner| interner.spans[index as usize])
+            let mut data = with_span_interner(|interner| interner.spans[index as usize]);
+            if data.ctxt.as_u32() == u32::MAX {
+                data.ctxt = SyntaxContext::from_u32(u32::from(self.ctxt_or_parent_or_marker));
+            }
+            data
         }
     }
 
@@ -236,6 +240,46 @@ impl Span {
         })
     }
 
+    pub fn update_ctxt(&mut self, update: impl FnOnce(SyntaxContext) -> SyntaxContext) {
+        let updated_ctxt32;
+        let data;
+        if self.len_with_tag_or_marker != BASE_LEN_INTERNED_MARKER {
+            if self.len_with_tag_or_marker & PARENT_TAG == 0 {
+                // Inline-context format.
+                updated_ctxt32 =
+                    update(SyntaxContext::from_u32(self.ctxt_or_parent_or_marker as u32)).as_u32();
+                if updated_ctxt32 <= MAX_CTXT {
+                    self.ctxt_or_parent_or_marker = updated_ctxt32 as u16;
+                    return;
+                }
+            } else {
+                // Inline-parent format. We know that the SyntaxContext is root.
+                updated_ctxt32 = update(SyntaxContext::root()).as_u32();
+                if updated_ctxt32 == 0 {
+                    // do nothing
+                    return;
+                }
+            }
+            data = self.data_untracked();
+        } else if self.ctxt_or_parent_or_marker != CTXT_INTERNED_MARKER {
+            // Partially-interned format. This path avoids looking up the
+            // interned value, and is the whole point of the
+            // partially-interned format.
+            updated_ctxt32 =
+                update(SyntaxContext::from_u32(self.ctxt_or_parent_or_marker as u32)).as_u32();
+            if updated_ctxt32 <= MAX_CTXT {
+                self.ctxt_or_parent_or_marker = updated_ctxt32 as u16;
+                return;
+            }
+            data = self.data_untracked();
+        } else {
+            data = self.data_untracked();
+            updated_ctxt32 = update(data.ctxt).as_u32();
+        };
+
+        *self = data.with_ctxt(SyntaxContext::from_u32(updated_ctxt32));
+    }
+
     /// This function is used as a fast path when decoding the full `SpanData` is not necessary.
     /// It's a cut-down version of `data_untracked`.
     #[cfg_attr(not(test), rustc_diagnostic_item = "SpanCtxt")]