Rollup merge of #125829 - petrochenkov:upctxt2, r=michaelwoerister

rustc_span: Add conveniences for working with span formats

This is the refactoring part of https://github.com/rust-lang/rust/pull/125017.

compiler/rustc_expand/src/mbe/transcribe.rs

@@ -31,7 +31,7 @@ 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;
-        span.update_ctxt(|ctxt| {
+        *span = span.map_ctxt(|ctxt| {
             *cache
                 .entry(ctxt)
                 .or_insert_with(|| ctxt.apply_mark(expn_id.to_expn_id(), transparency))

compiler/rustc_span/src/lib.rs

@@ -520,6 +520,7 @@ impl SpanData {
     pub fn with_hi(&self, hi: BytePos) -> Span {
         Span::new(self.lo, hi, self.ctxt, self.parent)
     }
+    /// Avoid if possible, `Span::map_ctxt` should be preferred.
     #[inline]
     fn with_ctxt(&self, ctxt: SyntaxContext) -> Span {
         Span::new(self.lo, self.hi, ctxt, self.parent)
@@ -576,9 +577,8 @@ impl Span {
         self.data().with_hi(hi)
     }
     #[inline]
-    pub fn with_ctxt(mut self, ctxt: SyntaxContext) -> Span {
+    pub fn with_ctxt(self, ctxt: SyntaxContext) -> Span {
-        self.update_ctxt(|_| ctxt);
+        self.map_ctxt(|_| ctxt)
-        self
     }
     #[inline]
     pub fn parent(self) -> Option<LocalDefId> {
@@ -1059,9 +1059,8 @@ impl Span {
     }
 
     #[inline]
-    pub fn apply_mark(mut self, expn_id: ExpnId, transparency: Transparency) -> Span {
+    pub fn apply_mark(self, expn_id: ExpnId, transparency: Transparency) -> Span {
-        self.update_ctxt(|ctxt| ctxt.apply_mark(expn_id, transparency));
+        self.map_ctxt(|ctxt| ctxt.apply_mark(expn_id, transparency))
-        self
     }
 
     #[inline]
@@ -1109,15 +1108,13 @@ impl Span {
     }
 
     #[inline]
-    pub fn normalize_to_macros_2_0(mut self) -> Span {
+    pub fn normalize_to_macros_2_0(self) -> Span {
-        self.update_ctxt(|ctxt| ctxt.normalize_to_macros_2_0());
+        self.map_ctxt(|ctxt| ctxt.normalize_to_macros_2_0())
-        self
     }
 
     #[inline]
-    pub fn normalize_to_macro_rules(mut self) -> Span {
+    pub fn normalize_to_macro_rules(self) -> Span {
-        self.update_ctxt(|ctxt| ctxt.normalize_to_macro_rules());
+        self.map_ctxt(|ctxt| ctxt.normalize_to_macro_rules())
-        self
     }
 }
 

compiler/rustc_span/src/span_encoding.rs

@@ -87,43 +87,150 @@ pub struct Span {
     ctxt_or_parent_or_marker: u16,
 }
 
-impl Span {
+// Convenience structures for all span formats.
+#[derive(Clone, Copy)]
+struct InlineCtxt {
+    lo: u32,
+    len: u16,
+    ctxt: u16,
+}
+
+#[derive(Clone, Copy)]
+struct InlineParent {
+    lo: u32,
+    len_with_tag: u16,
+    parent: u16,
+}
+
+#[derive(Clone, Copy)]
+struct PartiallyInterned {
+    index: u32,
+    ctxt: u16,
+}
+
+#[derive(Clone, Copy)]
+struct Interned {
+    index: u32,
+}
+
+impl InlineCtxt {
     #[inline]
-    fn data_inline_ctxt(self) -> SpanData {
+    fn data(self) -> SpanData {
-        let len = self.len_with_tag_or_marker as u32;
+        let len = self.len as u32;
         debug_assert!(len <= MAX_LEN);
         SpanData {
-            lo: BytePos(self.lo_or_index),
+            lo: BytePos(self.lo),
-            hi: BytePos(self.lo_or_index.debug_strict_add(len)),
+            hi: BytePos(self.lo.debug_strict_add(len)),
-            ctxt: SyntaxContext::from_u32(self.ctxt_or_parent_or_marker as u32),
+            ctxt: SyntaxContext::from_u32(self.ctxt as u32),
             parent: None,
         }
     }
     #[inline]
-    fn data_inline_parent(self) -> SpanData {
+    fn span(lo: u32, len: u16, ctxt: u16) -> Span {
-        let len = (self.len_with_tag_or_marker & !PARENT_TAG) as u32;
+        Span { lo_or_index: lo, len_with_tag_or_marker: len, ctxt_or_parent_or_marker: ctxt }
+    }
+    #[inline]
+    fn from_span(span: Span) -> InlineCtxt {
+        let (lo, len, ctxt) =
+            (span.lo_or_index, span.len_with_tag_or_marker, span.ctxt_or_parent_or_marker);
+        InlineCtxt { lo, len, ctxt }
+    }
+}
+
+impl InlineParent {
+    #[inline]
+    fn data(self) -> SpanData {
+        let len = (self.len_with_tag & !PARENT_TAG) as u32;
         debug_assert!(len <= MAX_LEN);
-        let parent = LocalDefId {
-            local_def_index: DefIndex::from_u32(self.ctxt_or_parent_or_marker as u32),
-        };
         SpanData {
-            lo: BytePos(self.lo_or_index),
+            lo: BytePos(self.lo),
-            hi: BytePos(self.lo_or_index.debug_strict_add(len)),
+            hi: BytePos(self.lo.debug_strict_add(len)),
             ctxt: SyntaxContext::root(),
-            parent: Some(parent),
+            parent: Some(LocalDefId { local_def_index: DefIndex::from_u32(self.parent as u32) }),
         }
     }
     #[inline]
-    fn data_partially_interned(self) -> SpanData {
+    fn span(lo: u32, len: u16, parent: u16) -> Span {
+        let (lo_or_index, len_with_tag_or_marker, ctxt_or_parent_or_marker) =
+            (lo, PARENT_TAG | len, parent);
+        Span { lo_or_index, len_with_tag_or_marker, ctxt_or_parent_or_marker }
+    }
+    #[inline]
+    fn from_span(span: Span) -> InlineParent {
+        let (lo, len_with_tag, parent) =
+            (span.lo_or_index, span.len_with_tag_or_marker, span.ctxt_or_parent_or_marker);
+        InlineParent { lo, len_with_tag, parent }
+    }
+}
+
+impl PartiallyInterned {
+    #[inline]
+    fn data(self) -> SpanData {
         SpanData {
-            ctxt: SyntaxContext::from_u32(self.ctxt_or_parent_or_marker as u32),
+            ctxt: SyntaxContext::from_u32(self.ctxt as u32),
-            ..with_span_interner(|interner| interner.spans[self.lo_or_index as usize])
+            ..with_span_interner(|interner| interner.spans[self.index as usize])
         }
     }
     #[inline]
-    fn data_interned(self) -> SpanData {
+    fn span(index: u32, ctxt: u16) -> Span {
-        with_span_interner(|interner| interner.spans[self.lo_or_index as usize])
+        let (lo_or_index, len_with_tag_or_marker, ctxt_or_parent_or_marker) =
+            (index, BASE_LEN_INTERNED_MARKER, ctxt);
+        Span { lo_or_index, len_with_tag_or_marker, ctxt_or_parent_or_marker }
     }
+    #[inline]
+    fn from_span(span: Span) -> PartiallyInterned {
+        PartiallyInterned { index: span.lo_or_index, ctxt: span.ctxt_or_parent_or_marker }
+    }
+}
+
+impl Interned {
+    #[inline]
+    fn data(self) -> SpanData {
+        with_span_interner(|interner| interner.spans[self.index as usize])
+    }
+    #[inline]
+    fn span(index: u32) -> Span {
+        let (lo_or_index, len_with_tag_or_marker, ctxt_or_parent_or_marker) =
+            (index, BASE_LEN_INTERNED_MARKER, CTXT_INTERNED_MARKER);
+        Span { lo_or_index, len_with_tag_or_marker, ctxt_or_parent_or_marker }
+    }
+    #[inline]
+    fn from_span(span: Span) -> Interned {
+        Interned { index: span.lo_or_index }
+    }
+}
+
+// This code is very hot, and converting span to an enum and matching on it doesn't optimize away
+// properly. So we are using a macro emulating such a match, but expand it directly to an if-else
+// chain.
+macro_rules! match_span_kind {
+    (
+        $span:expr,
+        InlineCtxt($span1:ident) => $arm1:expr,
+        InlineParent($span2:ident) => $arm2:expr,
+        PartiallyInterned($span3:ident) => $arm3:expr,
+        Interned($span4:ident) => $arm4:expr,
+    ) => {
+        if $span.len_with_tag_or_marker != BASE_LEN_INTERNED_MARKER {
+            if $span.len_with_tag_or_marker & PARENT_TAG == 0 {
+                // Inline-context format.
+                let $span1 = InlineCtxt::from_span($span);
+                $arm1
+            } else {
+                // Inline-parent format.
+                let $span2 = InlineParent::from_span($span);
+                $arm2
+            }
+        } else if $span.ctxt_or_parent_or_marker != CTXT_INTERNED_MARKER {
+            // Partially-interned format.
+            let $span3 = PartiallyInterned::from_span($span);
+            $arm3
+        } else {
+            // Interned format.
+            let $span4 = Interned::from_span($span);
+            $arm4
+        }
+    };
 }
 
 // `MAX_LEN` is chosen so that `PARENT_TAG | MAX_LEN` is distinct from
@@ -154,23 +261,13 @@ impl Span {
         let (len, ctxt32) = (hi.0 - lo.0, ctxt.as_u32());
         if len <= MAX_LEN {
             if ctxt32 <= MAX_CTXT && parent.is_none() {
-                // Inline-context format.
+                return InlineCtxt::span(lo.0, len as u16, ctxt32 as u16);
-                return Span {
-                    lo_or_index: lo.0,
-                    len_with_tag_or_marker: len as u16,
-                    ctxt_or_parent_or_marker: ctxt32 as u16,
-                };
             } else if ctxt32 == 0
                 && let Some(parent) = parent
                 && let parent32 = parent.local_def_index.as_u32()
                 && parent32 <= MAX_CTXT
             {
-                // Inline-parent format.
+                return InlineParent::span(lo.0, len as u16, parent32 as u16);
-                return Span {
-                    lo_or_index: lo.0,
-                    len_with_tag_or_marker: PARENT_TAG | len as u16,
-                    ctxt_or_parent_or_marker: parent32 as u16,
-                };
             }
         }
 
@@ -179,20 +276,10 @@ impl Span {
             with_span_interner(|interner| interner.intern(&SpanData { lo, hi, ctxt, parent }))
         };
         if ctxt32 <= MAX_CTXT {
-            // Partially-interned format.
+            // Interned ctxt should never be read, so it can use any value.
-            Span {
+            PartiallyInterned::span(index(SyntaxContext::from_u32(u32::MAX)), ctxt32 as u16)
-                // Interned ctxt should never be read, so it can use any value.
-                lo_or_index: index(SyntaxContext::from_u32(u32::MAX)),
-                len_with_tag_or_marker: BASE_LEN_INTERNED_MARKER,
-                ctxt_or_parent_or_marker: ctxt32 as u16,
-            }
         } else {
-            // Interned format.
+            Interned::span(index(ctxt))
-            Span {
-                lo_or_index: index(ctxt),
-                len_with_tag_or_marker: BASE_LEN_INTERNED_MARKER,
-                ctxt_or_parent_or_marker: CTXT_INTERNED_MARKER,
-            }
         }
     }
 
@@ -209,20 +296,12 @@ impl Span {
     /// This function must not be used outside the incremental engine.
     #[inline]
     pub fn data_untracked(self) -> SpanData {
-        if self.len_with_tag_or_marker != BASE_LEN_INTERNED_MARKER {
+        match_span_kind! {
-            if self.len_with_tag_or_marker & PARENT_TAG == 0 {
+            self,
-                // Inline-context format.
+            InlineCtxt(span) => span.data(),
-                self.data_inline_ctxt()
+            InlineParent(span) => span.data(),
-            } else {
+            PartiallyInterned(span) => span.data(),
-                // Inline-parent format.
+            Interned(span) => span.data(),
-                self.data_inline_parent()
-            }
-        } else if self.ctxt_or_parent_or_marker != CTXT_INTERNED_MARKER {
-            // Partially-interned format.
-            self.data_partially_interned()
-        } else {
-            // Interned format.
-            self.data_interned()
         }
     }
 
@@ -247,68 +326,57 @@ impl Span {
     // 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 update_ctxt(&mut self, update: impl FnOnce(SyntaxContext) -> SyntaxContext) {
+    pub fn map_ctxt(self, update: impl FnOnce(SyntaxContext) -> SyntaxContext) -> Span {
         let (updated_ctxt32, data);
-        if self.len_with_tag_or_marker != BASE_LEN_INTERNED_MARKER {
+        match_span_kind! {
-            if self.len_with_tag_or_marker & PARENT_TAG == 0 {
+            self,
-                // Inline-context format.
+            InlineCtxt(span) => {
                 updated_ctxt32 =
-                    update(SyntaxContext::from_u32(self.ctxt_or_parent_or_marker as u32)).as_u32();
+                    update(SyntaxContext::from_u32(span.ctxt as u32)).as_u32();
                 // Any small new context including zero will preserve the format.
                 if updated_ctxt32 <= MAX_CTXT {
-                    self.ctxt_or_parent_or_marker = updated_ctxt32 as u16;
+                    return InlineCtxt::span(span.lo, span.len, updated_ctxt32 as u16);
-                    return;
                 }
-                data = self.data_inline_ctxt();
+                data = span.data();
-            } else {
+            },
-                // Inline-parent format.
+            InlineParent(span) => {
                 updated_ctxt32 = update(SyntaxContext::root()).as_u32();
                 // Only if the new context is zero the format will be preserved.
                 if updated_ctxt32 == 0 {
                     // Do nothing.
-                    return;
+                    return self;
                 }
-                data = self.data_inline_parent();
+                data = span.data();
-            }
+            },
-        } else if self.ctxt_or_parent_or_marker != CTXT_INTERNED_MARKER {
+            PartiallyInterned(span) => {
-            // Partially-interned format.
+                updated_ctxt32 = update(SyntaxContext::from_u32(span.ctxt as u32)).as_u32();
-            updated_ctxt32 =
+                // Any small new context excluding zero will preserve the format.
-                update(SyntaxContext::from_u32(self.ctxt_or_parent_or_marker as u32)).as_u32();
+                // Zero may change the format to `InlineParent` if parent and len are small enough.
-            // Any small new context excluding zero will preserve the format.
+                if updated_ctxt32 <= MAX_CTXT && updated_ctxt32 != 0 {
-            // Zero may change the format to `InlineParent` if parent and len are small enough.
+                    return PartiallyInterned::span(span.index, updated_ctxt32 as u16);
-            if updated_ctxt32 <= MAX_CTXT && updated_ctxt32 != 0 {
+                }
-                self.ctxt_or_parent_or_marker = updated_ctxt32 as u16;
+                data = span.data();
-                return;
+            },
-            }
+            Interned(span) => {
-            data = self.data_partially_interned();
+                data = span.data();
-        } else {
+                updated_ctxt32 = update(data.ctxt).as_u32();
-            // Interned format.
+            },
-            data = self.data_interned();
-            updated_ctxt32 = update(data.ctxt).as_u32();
         }
 
         // We could not keep the span in the same inline format, fall back to the complete logic.
-        *self = data.with_ctxt(SyntaxContext::from_u32(updated_ctxt32));
+        data.with_ctxt(SyntaxContext::from_u32(updated_ctxt32))
     }
 
     // Returns either syntactic context, if it can be retrieved without taking the interner lock,
     // or an index into the interner if it cannot.
     #[inline]
     fn inline_ctxt(self) -> Result<SyntaxContext, usize> {
-        if self.len_with_tag_or_marker != BASE_LEN_INTERNED_MARKER {
+        match_span_kind! {
-            if self.len_with_tag_or_marker & PARENT_TAG == 0 {
+            self,
-                // Inline-context format.
+            InlineCtxt(span) => Ok(SyntaxContext::from_u32(span.ctxt as u32)),
-                Ok(SyntaxContext::from_u32(self.ctxt_or_parent_or_marker as u32))
+            InlineParent(_span) => Ok(SyntaxContext::root()),
-            } else {
+            PartiallyInterned(span) => Ok(SyntaxContext::from_u32(span.ctxt as u32)),
-                // Inline-parent format.
+            Interned(span) => Err(span.index as usize),
-                Ok(SyntaxContext::root())
-            }
-        } else if self.ctxt_or_parent_or_marker != CTXT_INTERNED_MARKER {
-            // Partially-interned format.
-            Ok(SyntaxContext::from_u32(self.ctxt_or_parent_or_marker as u32))
-        } else {
-            // Interned format.
-            Err(self.lo_or_index as usize)
         }
     }