rustc_span: Optimize syntax context comparisons

Including comparisons with root context

9 files changed, 72 insertions, 48 deletions

diff --git a/compiler/rustc_expand/src/mbe/quoted.rs b/compiler/rustc_expand/src/mbe/quoted.rs
index 445be01bc97..cd4ba7a9a62 100644
--- a/compiler/rustc_expand/src/mbe/quoted.rs
+++ b/compiler/rustc_expand/src/mbe/quoted.rs
@@ -72,7 +72,7 @@ pub(super) fn parse(
                                             // `SyntaxContext::root()` from a foreign crate will
                                             // have the edition of that crate (which we manually
                                             // retrieve via the `edition` parameter).
-                                            if span.ctxt().is_root() {
+                                            if !span.from_expansion() {
                                                 edition
                                             } else {
                                                 span.edition()
diff --git a/compiler/rustc_lint/src/internal.rs b/compiler/rustc_lint/src/internal.rs
index 53d99c7f7f3..e3405aa2e55 100644
--- a/compiler/rustc_lint/src/internal.rs
+++ b/compiler/rustc_lint/src/internal.rs
@@ -549,7 +549,9 @@ declare_lint_pass!(SpanUseEqCtxt => [SPAN_USE_EQ_CTXT]);
 
 impl<'tcx> LateLintPass<'tcx> for SpanUseEqCtxt {
     fn check_expr(&mut self, cx: &LateContext<'tcx>, expr: &Expr<'_>) {
-        if let ExprKind::Binary(BinOp { node: BinOpKind::Eq, .. }, lhs, rhs) = expr.kind {
+        if let ExprKind::Binary(BinOp { node: BinOpKind::Eq | BinOpKind::Ne, .. }, lhs, rhs) =
+            expr.kind
+        {
             if is_span_ctxt_call(cx, lhs) && is_span_ctxt_call(cx, rhs) {
                 cx.emit_spanned_lint(SPAN_USE_EQ_CTXT, expr.span, SpanUseEqCtxtDiag);
             }
diff --git a/compiler/rustc_lint/src/non_fmt_panic.rs b/compiler/rustc_lint/src/non_fmt_panic.rs
index f0bbc03d747..479acd88d71 100644
--- a/compiler/rustc_lint/src/non_fmt_panic.rs
+++ b/compiler/rustc_lint/src/non_fmt_panic.rs
@@ -111,10 +111,11 @@ fn check_panic<'tcx>(cx: &LateContext<'tcx>, f: &'tcx hir::Expr<'tcx>, arg: &'tc
     let mut arg_span = arg.span;
     let mut arg_macro = None;
     while !span.contains(arg_span) {
-        let expn = arg_span.ctxt().outer_expn_data();
-        if expn.is_root() {
+        let ctxt = arg_span.ctxt();
+        if ctxt.is_root() {
             break;
         }
+        let expn = ctxt.outer_expn_data();
         arg_macro = expn.macro_def_id;
         arg_span = expn.call_site;
     }
diff --git a/compiler/rustc_passes/src/liveness.rs b/compiler/rustc_passes/src/liveness.rs
index cfe829f170f..75ef1cd38e8 100644
--- a/compiler/rustc_passes/src/liveness.rs
+++ b/compiler/rustc_passes/src/liveness.rs
@@ -1630,7 +1630,7 @@ impl<'tcx> Liveness<'_, 'tcx> {
                     let from_macro = non_shorthands
                         .iter()
                         .find(|(_, pat_span, ident_span)| {
-                            pat_span.ctxt() != ident_span.ctxt() && pat_span.from_expansion()
+                            !pat_span.eq_ctxt(*ident_span) && pat_span.from_expansion()
                         })
                         .map(|(_, pat_span, _)| *pat_span);
                     let non_shorthands = non_shorthands
diff --git a/compiler/rustc_span/src/hygiene.rs b/compiler/rustc_span/src/hygiene.rs
index b717229b68d..72642fa27bb 100644
--- a/compiler/rustc_span/src/hygiene.rs
+++ b/compiler/rustc_span/src/hygiene.rs
@@ -295,11 +295,13 @@ impl ExpnId {
     pub fn expansion_cause(mut self) -> Option<Span> {
         let mut last_macro = None;
         loop {
+            // Fast path to avoid locking.
+            if self == ExpnId::root() {
+                break;
+            }
             let expn_data = self.expn_data();
             // Stop going up the backtrace once include! is encountered
-            if expn_data.is_root()
-                || expn_data.kind == ExpnKind::Macro(MacroKind::Bang, sym::include)
-            {
+            if expn_data.kind == ExpnKind::Macro(MacroKind::Bang, sym::include) {
                 break;
             }
             self = expn_data.call_site.ctxt().outer_expn();
@@ -433,7 +435,7 @@ impl HygieneData {
 
     fn marks(&self, mut ctxt: SyntaxContext) -> Vec<(ExpnId, Transparency)> {
         let mut marks = Vec::new();
-        while ctxt != SyntaxContext::root() {
+        while !ctxt.is_root() {
             debug!("marks: getting parent of {:?}", ctxt);
             marks.push(self.outer_mark(ctxt));
             ctxt = self.parent_ctxt(ctxt);
diff --git a/compiler/rustc_span/src/lib.rs b/compiler/rustc_span/src/lib.rs
index 8f64eed9a87..3dc5434c7af 100644
--- a/compiler/rustc_span/src/lib.rs
+++ b/compiler/rustc_span/src/lib.rs
@@ -541,10 +541,6 @@ impl Span {
         self.data().with_hi(hi)
     }
     #[inline]
-    pub fn eq_ctxt(self, other: Span) -> bool {
-        self.data_untracked().ctxt == other.data_untracked().ctxt
-    }
-    #[inline]
     pub fn with_ctxt(self, ctxt: SyntaxContext) -> Span {
         self.data_untracked().with_ctxt(ctxt)
     }
@@ -565,7 +561,7 @@ impl Span {
     /// Returns `true` if this span comes from any kind of macro, desugaring or inlining.
     #[inline]
     pub fn from_expansion(self) -> bool {
-        self.ctxt() != SyntaxContext::root()
+        !self.ctxt().is_root()
     }
 
     /// Returns `true` if `span` originates in a macro's expansion where debuginfo should be
@@ -654,15 +650,15 @@ impl Span {
     /// Returns the source span -- this is either the supplied span, or the span for
     /// the macro callsite that expanded to it.
     pub fn source_callsite(self) -> Span {
-        let expn_data = self.ctxt().outer_expn_data();
-        if !expn_data.is_root() { expn_data.call_site.source_callsite() } else { self }
+        let ctxt = self.ctxt();
+        if !ctxt.is_root() { ctxt.outer_expn_data().call_site.source_callsite() } else { self }
     }
 
     /// The `Span` for the tokens in the previous macro expansion from which `self` was generated,
     /// if any.
     pub fn parent_callsite(self) -> Option<Span> {
-        let expn_data = self.ctxt().outer_expn_data();
-        if !expn_data.is_root() { Some(expn_data.call_site) } else { None }
+        let ctxt = self.ctxt();
+        (!ctxt.is_root()).then(|| ctxt.outer_expn_data().call_site)
     }
 
     /// Walk down the expansion ancestors to find a span that's contained within `outer`.
@@ -747,15 +743,14 @@ impl Span {
     /// else returns the `ExpnData` for the macro definition
     /// corresponding to the source callsite.
     pub fn source_callee(self) -> Option<ExpnData> {
-        let expn_data = self.ctxt().outer_expn_data();
-
-        // Create an iterator of call site expansions
-        iter::successors(Some(expn_data), |expn_data| {
-            Some(expn_data.call_site.ctxt().outer_expn_data())
-        })
-        // Find the last expansion which is not root
-        .take_while(|expn_data| !expn_data.is_root())
-        .last()
+        let mut ctxt = self.ctxt();
+        let mut opt_expn_data = None;
+        while !ctxt.is_root() {
+            let expn_data = ctxt.outer_expn_data();
+            ctxt = expn_data.call_site.ctxt();
+            opt_expn_data = Some(expn_data);
+        }
+        opt_expn_data
     }
 
     /// Checks if a span is "internal" to a macro in which `#[unstable]`
@@ -796,11 +791,12 @@ impl Span {
         let mut prev_span = DUMMY_SP;
         iter::from_fn(move || {
             loop {
-                let expn_data = self.ctxt().outer_expn_data();
-                if expn_data.is_root() {
+                let ctxt = self.ctxt();
+                if ctxt.is_root() {
                     return None;
                 }
 
+                let expn_data = ctxt.outer_expn_data();
                 let is_recursive = expn_data.call_site.source_equal(prev_span);
 
                 prev_span = self;
diff --git a/compiler/rustc_span/src/source_map.rs b/compiler/rustc_span/src/source_map.rs
index c61dbcaae95..8253ffefcaa 100644
--- a/compiler/rustc_span/src/source_map.rs
+++ b/compiler/rustc_span/src/source_map.rs
@@ -23,9 +23,15 @@ mod tests;
 /// otherwise return the call site span up to the `enclosing_sp` by
 /// following the `expn_data` chain.
 pub fn original_sp(sp: Span, enclosing_sp: Span) -> Span {
-    let expn_data1 = sp.ctxt().outer_expn_data();
-    let expn_data2 = enclosing_sp.ctxt().outer_expn_data();
-    if expn_data1.is_root() || !expn_data2.is_root() && expn_data1.call_site == expn_data2.call_site
+    let ctxt = sp.ctxt();
+    if ctxt.is_root() {
+        return sp;
+    }
+
+    let enclosing_ctxt = enclosing_sp.ctxt();
+    let expn_data1 = ctxt.outer_expn_data();
+    if !enclosing_ctxt.is_root()
+        && expn_data1.call_site == enclosing_ctxt.outer_expn_data().call_site
     {
         sp
     } else {
diff --git a/compiler/rustc_span/src/source_map/tests.rs b/compiler/rustc_span/src/source_map/tests.rs
index 130522a302d..5788d11ed43 100644
--- a/compiler/rustc_span/src/source_map/tests.rs
+++ b/compiler/rustc_span/src/source_map/tests.rs
@@ -18,7 +18,7 @@ impl SourceMap {
     ///    * the LHS span must start at or before the RHS span.
     fn merge_spans(&self, sp_lhs: Span, sp_rhs: Span) -> Option<Span> {
         // Ensure we're at the same expansion ID.
-        if sp_lhs.ctxt() != sp_rhs.ctxt() {
+        if !sp_lhs.eq_ctxt(sp_rhs) {
             return None;
         }
 
diff --git a/compiler/rustc_span/src/span_encoding.rs b/compiler/rustc_span/src/span_encoding.rs
index f7d17a267d6..e162695a13b 100644
--- a/compiler/rustc_span/src/span_encoding.rs
+++ b/compiler/rustc_span/src/span_encoding.rs
@@ -210,12 +210,10 @@ impl Span {
         }
     }
 
-    /// 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")]
-    #[inline]
-    pub fn ctxt(self) -> SyntaxContext {
-        if self.len_with_tag_or_marker != BASE_LEN_INTERNED_MARKER {
+    // 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 & PARENT_TAG == 0 {
                 // Inline-context format.
                 SyntaxContext::from_u32(self.ctxt_or_parent_or_marker as u32)
@@ -223,17 +221,36 @@ impl Span {
                 // Inline-parent format. We know that the SyntaxContext is root.
                 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)
         } 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)
-            } else {
-                // Fully-interned format.
-                let index = self.lo_or_index;
-                with_span_interner(|interner| interner.spans[index as usize].ctxt)
+            // Fully-interned format.
+            return Err(self.lo_or_index as usize);
+        })
+    }
+
+    /// 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")]
+    #[inline]
+    pub fn ctxt(self) -> SyntaxContext {
+        self.inline_ctxt()
+            .unwrap_or_else(|index| with_span_interner(|interner| interner.spans[index].ctxt))
+    }
+
+    #[inline]
+    pub fn eq_ctxt(self, other: Span) -> bool {
+        match (self.inline_ctxt(), other.inline_ctxt()) {
+            (Ok(ctxt1), Ok(ctxt2)) => ctxt1 == ctxt2,
+            (Ok(ctxt), Err(index)) | (Err(index), Ok(ctxt)) => {
+                with_span_interner(|interner| ctxt == interner.spans[index].ctxt)
             }
+            (Err(index1), Err(index2)) => with_span_interner(|interner| {
+                interner.spans[index1].ctxt == interner.spans[index2].ctxt
+            }),
         }
     }
 }