rustc_span: Optimize syntax context updates in spans

petrochenkov · petrochenkov · commit d364faaceecc · 2024-06-03T18:33:09.000+03:00
diff --git a/compiler/rustc_expand/src/mbe/transcribe.rs b/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))
+        });
     }
 }
 
diff --git a/compiler/rustc_span/src/lib.rs b/compiler/rustc_span/src/lib.rs
@@ -521,7 +521,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]
@@ -576,8 +576,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> {
@@ -1051,9 +1052,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]
@@ -1101,15 +1102,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
     }
 }
 
diff --git a/compiler/rustc_span/src/span_encoding.rs b/compiler/rustc_span/src/span_encoding.rs
@@ -87,6 +87,45 @@ pub struct Span {
     ctxt_or_parent_or_marker: u16,
 }
 
+impl Span {
+    #[inline]
+    fn data_inline_ctxt(self) -> SpanData {
+        let len = self.len_with_tag_or_marker 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),
+            parent: None,
+        }
+    }
+    #[inline]
+    fn data_inline_parent(self) -> SpanData {
+        let len = (self.len_with_tag_or_marker & !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)),
+            ctxt: SyntaxContext::root(),
+            parent: Some(parent),
+        }
+    }
+    #[inline]
+    fn data_partially_interned(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])
+        }
+    }
+    #[inline]
+    fn data_interned(self) -> SpanData {
+        with_span_interner(|interner| interner.spans[self.lo_or_index as usize])
+    }
+}
+
 // `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;
@@ -111,42 +150,49 @@ impl Span {
             std::mem::swap(&mut lo, &mut hi);
         }
 
-        let (lo2, len, ctxt2) = (lo.0, hi.0 - lo.0, ctxt.as_u32());
-
+        // Small len may enable one of fully inline formats (or may not).
+        let (len, ctxt32) = (hi.0 - lo.0, ctxt.as_u32());
         if len <= MAX_LEN {
-            if ctxt2 <= MAX_CTXT && parent.is_none() {
+            if ctxt32 <= MAX_CTXT && parent.is_none() {
                 // Inline-context format.
                 return Span {
-                    lo_or_index: lo2,
+                    lo_or_index: lo.0,
                     len_with_tag_or_marker: len as u16,
-                    ctxt_or_parent_or_marker: ctxt2 as u16,
+                    ctxt_or_parent_or_marker: ctxt32 as u16,
                 };
-            } else if ctxt2 == SyntaxContext::root().as_u32()
+            } else if ctxt32 == 0
                 && let Some(parent) = parent
-                && let parent2 = parent.local_def_index.as_u32()
-                && parent2 <= MAX_CTXT
+                && let parent32 = parent.local_def_index.as_u32()
+                && parent32 <= MAX_CTXT
             {
                 // Inline-parent format.
                 return Span {
-                    lo_or_index: lo2,
+                    lo_or_index: lo.0,
                     len_with_tag_or_marker: PARENT_TAG | len as u16,
-                    ctxt_or_parent_or_marker: parent2 as u16,
+                    ctxt_or_parent_or_marker: parent32 as u16,
                 };
             }
         }
 
-        // 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
-        } else {
-            CTXT_INTERNED_MARKER // fully-interned
+        // Otherwise small ctxt may enable the partially inline format.
+        let index = |ctxt| {
+            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,
-            ctxt_or_parent_or_marker,
+        if ctxt32 <= MAX_CTXT {
+            // Partially-interned format.
+            Span {
+                // Any value, should never be read.
+                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.
+            Span {
+                lo_or_index: index(ctxt),
+                len_with_tag_or_marker: BASE_LEN_INTERNED_MARKER,
+                ctxt_or_parent_or_marker: CTXT_INTERNED_MARKER,
+            }
         }
     }
 
@@ -166,34 +212,17 @@ impl Span {
         if self.len_with_tag_or_marker != BASE_LEN_INTERNED_MARKER {
             if self.len_with_tag_or_marker & PARENT_TAG == 0 {
                 // Inline-context format.
-                let len = self.len_with_tag_or_marker 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),
-                    parent: None,
-                }
+                self.data_inline_ctxt()
             } else {
                 // Inline-parent format.
-                let len = (self.len_with_tag_or_marker & !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)),
-                    ctxt: SyntaxContext::root(),
-                    parent: Some(parent),
-                }
+                self.data_inline_parent()
             }
+        } else if self.ctxt_or_parent_or_marker != CTXT_INTERNED_MARKER {
+            // Partially-interned format.
+            self.data_partially_interned()
         } else {
-            // Fully-interned or partially-interned format. In either case,
-            // 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])
+            // Interned format.
+            self.data_interned()
         }
     }
 
@@ -214,26 +243,71 @@ impl Span {
         }
     }
 
+    // For optimization we are interested in cases in which the context is inline and the context
+    // 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) {
+        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.
+                updated_ctxt32 =
+                    update(SyntaxContext::from_u32(self.ctxt_or_parent_or_marker 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;
+                }
+                data = self.data_inline_ctxt();
+            } else {
+                // Inline-parent format.
+                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();
+        }
+
+        *self = 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.
     fn inline_ctxt(self) -> Result<SyntaxContext, usize> {
-        Ok(if self.len_with_tag_or_marker != BASE_LEN_INTERNED_MARKER {
+        if self.len_with_tag_or_marker != BASE_LEN_INTERNED_MARKER {
             if self.len_with_tag_or_marker & PARENT_TAG == 0 {
                 // Inline-context format.
-                SyntaxContext::from_u32(self.ctxt_or_parent_or_marker as u32)
+                Ok(SyntaxContext::from_u32(self.ctxt_or_parent_or_marker as u32))
             } else {
-                // Inline-parent format. We know that the SyntaxContext is root.
-                SyntaxContext::root()
+                // Inline-parent format.
+                Ok(SyntaxContext::root())
             }
         } 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.
-            SyntaxContext::from_u32(self.ctxt_or_parent_or_marker as u32)
+            // Partially-interned format.
+            Ok(SyntaxContext::from_u32(self.ctxt_or_parent_or_marker as u32))
         } else {
-            // Fully-interned format.
-            return Err(self.lo_or_index as usize);
-        })
+            // Interned format.
+            Err(self.lo_or_index as usize)
+        }
     }
 
     /// This function is used as a fast path when decoding the full `SpanData` is not necessary.

Original file line number	Diff line number	Diff line change
`@@ -521,7 +521,7 @@ impl SpanData {`
`521`	`521`	`Span::new(self.lo, hi, self.ctxt, self.parent)`
`522`	`522`	`}`
`523`	`523`	`#[inline]`
`524`		`- pub fn with_ctxt(&self, ctxt: SyntaxContext) -> Span {`
	`524`	`+ fn with_ctxt(&self, ctxt: SyntaxContext) -> Span {`
`525`	`525`	`Span::new(self.lo, self.hi, ctxt, self.parent)`
`526`	`526`	`}`
`527`	`527`	`#[inline]`
`@@ -576,8 +576,9 @@ impl Span {`
`576`	`576`	`self.data().with_hi(hi)`
`577`	`577`	`}`
`578`	`578`	`#[inline]`
`579`		`- pub fn with_ctxt(self, ctxt: SyntaxContext) -> Span {`
`580`		`- self.data_untracked().with_ctxt(ctxt)`
	`579`	`+ pub fn with_ctxt(mut self, ctxt: SyntaxContext) -> Span {`
	`580`	`+ self.update_ctxt(\|_\| ctxt);`
	`581`	`+ self`
`581`	`582`	`}`
`582`	`583`	`#[inline]`
`583`	`584`	`pub fn parent(self) -> Option<LocalDefId> {`
`@@ -1051,9 +1052,9 @@ impl Span {`
`1051`	`1052`	`}`
`1052`	`1053`
`1053`	`1054`	`#[inline]`
`1054`		`- pub fn apply_mark(self, expn_id: ExpnId, transparency: Transparency) -> Span {`
`1055`		`- let span = self.data();`
`1056`		`- span.with_ctxt(span.ctxt.apply_mark(expn_id, transparency))`
	`1055`	`+ pub fn apply_mark(mut self, expn_id: ExpnId, transparency: Transparency) -> Span {`
	`1056`	`+ self.update_ctxt(\|ctxt\| ctxt.apply_mark(expn_id, transparency));`
	`1057`	`+ self`
`1057`	`1058`	`}`
`1058`	`1059`
`1059`	`1060`	`#[inline]`
`@@ -1101,15 +1102,15 @@ impl Span {`
`1101`	`1102`	`}`
`1102`	`1103`
`1103`	`1104`	`#[inline]`
`1104`		`- pub fn normalize_to_macros_2_0(self) -> Span {`
`1105`		`- let span = self.data();`
`1106`		`- span.with_ctxt(span.ctxt.normalize_to_macros_2_0())`
	`1105`	`+ pub fn normalize_to_macros_2_0(mut self) -> Span {`
	`1106`	`+ self.update_ctxt(\|ctxt\| ctxt.normalize_to_macros_2_0());`
	`1107`	`+ self`
`1107`	`1108`	`}`
`1108`	`1109`
`1109`	`1110`	`#[inline]`
`1110`		`- pub fn normalize_to_macro_rules(self) -> Span {`
`1111`		`- let span = self.data();`
`1112`		`- span.with_ctxt(span.ctxt.normalize_to_macro_rules())`
	`1111`	`+ pub fn normalize_to_macro_rules(mut self) -> Span {`
	`1112`	`+ self.update_ctxt(\|ctxt\| ctxt.normalize_to_macro_rules());`
	`1113`	`+ self`
`1113`	`1114`	`}`
`1114`	`1115`	`}`
`1115`	`1116`