rustc_span: Add conveniences for working with span formats

Author: petrochenkov

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::update_ctxt` should be preferred.
     #[inline]
     fn with_ctxt(&self, ctxt: SyntaxContext) -> Span {
         Span::new(self.lo, self.hi, ctxt, self.parent)

compiler/rustc_span/src/span_encoding.rs

@@ -4,10 +4,10 @@ use crate::SPAN_TRACK;
 use crate::{BytePos, SpanData};
 
 use rustc_data_structures::fx::FxIndexSet;
-
 // This code is very hot and uses lots of arithmetic, avoid overflow checks for performance.
 // See https://github.com/rust-lang/rust/pull/119440#issuecomment-1874255727
 use rustc_serialize::int_overflow::DebugStrictAdd;
+use std::mem::transmute;
 
 /// A compressed span.
 ///
@@ -87,45 +87,130 @@ 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,
+    _marker1: u16,
+    ctxt: u16,
+}
+
+#[derive(Clone, Copy)]
+struct Interned {
+    index: u32,
+    _marker1: u16,
+    _marker2: u16,
+}
+
+impl InlineCtxt {
     #[inline]
-    fn data_inline_ctxt(self) -> SpanData {
-        let len = self.len_with_tag_or_marker as u32;
+    fn data(self) -> SpanData {
+        let len = self.len as u32;
         debug_assert!(len <= MAX_LEN);
         SpanData {
-            lo: BytePos(self.lo_or_index),
-            hi: BytePos(self.lo_or_index.debug_strict_add(len)),
-            ctxt: SyntaxContext::from_u32(self.ctxt_or_parent_or_marker as u32),
+            lo: BytePos(self.lo),
+            hi: BytePos(self.lo.debug_strict_add(len)),
+            ctxt: SyntaxContext::from_u32(self.ctxt as u32),
             parent: None,
         }
     }
     #[inline]
-    fn data_inline_parent(self) -> SpanData {
-        let len = (self.len_with_tag_or_marker & !PARENT_TAG) as u32;
+    fn span(lo: u32, len: u16, ctxt: u16) -> Span {
+        unsafe { transmute(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),
-            hi: BytePos(self.lo_or_index.debug_strict_add(len)),
+            lo: BytePos(self.lo),
+            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_with_tag: u16, parent: u16) -> Span {
+        unsafe { transmute(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),
-            ..with_span_interner(|interner| interner.spans[self.lo_or_index as usize])
+            ctxt: SyntaxContext::from_u32(self.ctxt as u32),
+            ..with_span_interner(|interner| interner.spans[self.index as usize])
         }
     }
     #[inline]
-    fn data_interned(self) -> SpanData {
-        with_span_interner(|interner| interner.spans[self.lo_or_index as usize])
+    fn span(index: u32, ctxt: u16) -> Span {
+        unsafe { transmute(PartiallyInterned { index, _marker1: BASE_LEN_INTERNED_MARKER, ctxt }) }
     }
 }
 
+impl Interned {
+    #[inline]
+    fn data(self) -> SpanData {
+        with_span_interner(|interner| interner.spans[self.index as usize])
+    }
+    #[inline]
+    fn span(index: u32) -> Span {
+        let _marker1 = BASE_LEN_INTERNED_MARKER;
+        unsafe { transmute(Interned { index, _marker1, _marker2: CTXT_INTERNED_MARKER }) }
+    }
+}
+
+// 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: &mut InlineCtxt = unsafe { transmute(&mut *$span) };
+                $arm1
+            } else {
+                // Inline-parent format.
+                let $span2: &mut InlineParent = unsafe { transmute(&mut *$span) };
+                $arm2
+            }
+        } else if $span.ctxt_or_parent_or_marker != CTXT_INTERNED_MARKER {
+            // Partially-interned format.
+            let $span3: &mut PartiallyInterned = unsafe { transmute(&mut *$span) };
+            $arm3
+        } else {
+            // Interned format.
+            let $span4: &mut Interned = unsafe { transmute(&mut *$span) };
+            $arm4
+        }
+    };
+}
+
 // `MAX_LEN` is chosen so that `PARENT_TAG | MAX_LEN` is distinct from
 // `BASE_LEN_INTERNED_MARKER`. (If `MAX_LEN` was 1 higher, this wouldn't be true.)
 const MAX_LEN: u32 = 0b0111_1111_1111_1110;
@@ -154,23 +239,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 Span {
-                    lo_or_index: lo.0,
-                    len_with_tag_or_marker: len as u16,
-                    ctxt_or_parent_or_marker: ctxt32 as u16,
-                };
+                return InlineCtxt::span(lo.0, len as u16, 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 Span {
-                    lo_or_index: lo.0,
-                    len_with_tag_or_marker: PARENT_TAG | len as u16,
-                    ctxt_or_parent_or_marker: parent32 as u16,
-                };
+                return InlineParent::span(lo.0, PARENT_TAG | len as u16, parent32 as u16);
             }
         }
 
@@ -179,20 +254,10 @@ impl Span {
             with_span_interner(|interner| interner.intern(&SpanData { lo, hi, ctxt, parent }))
         };
         if ctxt32 <= MAX_CTXT {
-            // Partially-interned format.
-            Span {
-                // 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,
-            }
+            // Interned ctxt should never be read, so it can use any value.
+            PartiallyInterned::span(index(SyntaxContext::from_u32(u32::MAX)), ctxt32 as u16)
         } else {
-            // Interned format.
-            Span {
-                lo_or_index: index(ctxt),
-                len_with_tag_or_marker: BASE_LEN_INTERNED_MARKER,
-                ctxt_or_parent_or_marker: CTXT_INTERNED_MARKER,
-            }
+            Interned::span(index(ctxt))
         }
     }
 
@@ -208,21 +273,13 @@ impl Span {
     /// Internal function to translate between an encoded span and the expanded representation.
     /// 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 {
-            if self.len_with_tag_or_marker & PARENT_TAG == 0 {
-                // Inline-context format.
-                self.data_inline_ctxt()
-            } else {
-                // Inline-parent format.
-                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()
+    pub fn data_untracked(mut self) -> SpanData {
+        match_span_kind! {
+            &mut self,
+            InlineCtxt(span) => span.data(),
+            InlineParent(span) => span.data(),
+            PartiallyInterned(span) => span.data(),
+            Interned(span) => span.data(),
         }
     }
 
@@ -249,42 +306,41 @@ impl Span {
     #[inline]
     pub fn update_ctxt(&mut self, update: impl FnOnce(SyntaxContext) -> SyntaxContext) {
         let (updated_ctxt32, 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.
+        match_span_kind! {
+            self,
+            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;
+                    span.ctxt = updated_ctxt32 as u16;
                     return;
                 }
-                data = self.data_inline_ctxt();
-            } else {
-                // Inline-parent format.
+                data = span.data();
+            },
+            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;
                 }
-                data = self.data_inline_parent();
-            }
-        } else if self.ctxt_or_parent_or_marker != CTXT_INTERNED_MARKER {
-            // Partially-interned format.
-            updated_ctxt32 =
-                update(SyntaxContext::from_u32(self.ctxt_or_parent_or_marker as u32)).as_u32();
-            // Any small new context excluding zero will preserve the format.
-            // Zero may change the format to `InlineParent` if parent and len are small enough.
-            if updated_ctxt32 <= MAX_CTXT && updated_ctxt32 != 0 {
-                self.ctxt_or_parent_or_marker = updated_ctxt32 as u16;
-                return;
-            }
-            data = self.data_partially_interned();
-        } else {
-            // Interned format.
-            data = self.data_interned();
-            updated_ctxt32 = update(data.ctxt).as_u32();
+                data = span.data();
+            },
+            PartiallyInterned(span) => {
+                updated_ctxt32 = update(SyntaxContext::from_u32(span.ctxt as u32)).as_u32();
+                // Any small new context excluding zero will preserve the format.
+                // Zero may change the format to `InlineParent` if parent and len are small enough.
+                if updated_ctxt32 <= MAX_CTXT && updated_ctxt32 != 0 {
+                    span.ctxt = updated_ctxt32 as u16;
+                    return;
+                }
+                data = span.data();
+            },
+            Interned(span) => {
+                data = span.data();
+                updated_ctxt32 = update(data.ctxt).as_u32();
+            },
         }
 
         // We could not keep the span in the same inline format, fall back to the complete logic.
@@ -294,21 +350,13 @@ impl Span {
     // 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 {
-            if self.len_with_tag_or_marker & PARENT_TAG == 0 {
-                // Inline-context format.
-                Ok(SyntaxContext::from_u32(self.ctxt_or_parent_or_marker as u32))
-            } else {
-                // Inline-parent format.
-                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)
+    fn inline_ctxt(mut self) -> Result<SyntaxContext, usize> {
+        match_span_kind! {
+            &mut self,
+            InlineCtxt(span) => Ok(SyntaxContext::from_u32(span.ctxt as u32)),
+            InlineParent(_span) => Ok(SyntaxContext::root()),
+            PartiallyInterned(span) => Ok(SyntaxContext::from_u32(span.ctxt as u32)),
+            Interned(span) => Err(span.index as usize),
         }
     }