From b2a17e9bd10b6a777485d9e37fa0289bf5ecf62f Mon Sep 17 00:00:00 2001
From: Camille GILLOT <gillot.camille@gmail.com>
Date: Sun, 3 Jul 2022 12:20:40 +0200
Subject: [PATCH] Keep ctxt in encoded span representation.

---
 compiler/rustc_span/src/lib.rs           |  3 ---
 compiler/rustc_span/src/span_encoding.rs | 32 +++++++++++++++++-------
 2 files changed, 23 insertions(+), 12 deletions(-)

diff --git a/compiler/rustc_span/src/lib.rs b/compiler/rustc_span/src/lib.rs
index cf306928151..c8dc1c1adcf 100644
--- a/compiler/rustc_span/src/lib.rs
+++ b/compiler/rustc_span/src/lib.rs
@@ -533,9 +533,6 @@ impl Span {
         self.data().with_hi(hi)
     }
     #[inline]
-    pub fn ctxt(self) -> SyntaxContext {
-        self.data_untracked().ctxt
-    }
     pub fn eq_ctxt(self, other: Span) -> bool {
         self.data_untracked().ctxt == other.data_untracked().ctxt
     }
diff --git a/compiler/rustc_span/src/span_encoding.rs b/compiler/rustc_span/src/span_encoding.rs
index 3ee329e9736..708ca43f4be 100644
--- a/compiler/rustc_span/src/span_encoding.rs
+++ b/compiler/rustc_span/src/span_encoding.rs
@@ -28,12 +28,12 @@ use rustc_data_structures::fx::FxIndexSet;
 /// Inline (compressed) format:
 /// - `span.base_or_index == span_data.lo`
 /// - `span.len_or_tag == len == span_data.hi - span_data.lo` (must be `<= MAX_LEN`)
-/// - `span.ctxt == span_data.ctxt` (must be `<= MAX_CTXT`)
+/// - `span.ctxt == span_data.ctxt` (must be `< MAX_CTXT`)
 ///
 /// Interned format:
 /// - `span.base_or_index == index` (indexes into the interner table)
 /// - `span.len_or_tag == LEN_TAG` (high bit set, all other bits are zero)
-/// - `span.ctxt == 0`
+/// - `span.ctxt == span_data.ctxt` (must be < `MAX_CTXT`) or `MAX_CTXT` otherwise
 ///
 /// The inline form uses 0 for the tag value (rather than 1) so that we don't
 /// need to mask out the tag bit when getting the length, and so that the
@@ -65,7 +65,7 @@ use rustc_data_structures::fx::FxIndexSet;
 pub struct Span {
     base_or_index: u32,
     len_or_tag: u16,
-    ctxt_or_zero: u16,
+    ctxt_or_max: u16,
 }
 
 const LEN_TAG: u16 = 0b1000_0000_0000_0000;
@@ -73,7 +73,7 @@ const MAX_LEN: u32 = 0b0111_1111_1111_1111;
 const MAX_CTXT: u32 = 0b1111_1111_1111_1111;
 
 /// Dummy span, both position and length are zero, syntax context is zero as well.
-pub const DUMMY_SP: Span = Span { base_or_index: 0, len_or_tag: 0, ctxt_or_zero: 0 };
+pub const DUMMY_SP: Span = Span { base_or_index: 0, len_or_tag: 0, ctxt_or_max: 0 };
 
 impl Span {
     #[inline]
@@ -89,14 +89,15 @@ impl Span {
 
         let (base, len, ctxt2) = (lo.0, hi.0 - lo.0, ctxt.as_u32());
 
-        if len <= MAX_LEN && ctxt2 <= MAX_CTXT && parent.is_none() {
+        if len <= MAX_LEN && ctxt2 < MAX_CTXT && parent.is_none() {
             // Inline format.
-            Span { base_or_index: base, len_or_tag: len as u16, ctxt_or_zero: ctxt2 as u16 }
+            Span { base_or_index: base, len_or_tag: len as u16, ctxt_or_max: ctxt2 as u16 }
         } else {
             // Interned format.
             let index =
                 with_span_interner(|interner| interner.intern(&SpanData { lo, hi, ctxt, parent }));
-            Span { base_or_index: index, len_or_tag: LEN_TAG, ctxt_or_zero: 0 }
+            let ctxt_or_max = if ctxt2 < MAX_CTXT { ctxt2 } else { MAX_CTXT } as u16;
+            Span { base_or_index: index, len_or_tag: LEN_TAG, ctxt_or_max }
         }
     }
 
@@ -119,16 +120,29 @@ impl Span {
             SpanData {
                 lo: BytePos(self.base_or_index),
                 hi: BytePos(self.base_or_index + self.len_or_tag as u32),
-                ctxt: SyntaxContext::from_u32(self.ctxt_or_zero as u32),
+                ctxt: SyntaxContext::from_u32(self.ctxt_or_max as u32),
                 parent: None,
             }
         } else {
             // Interned format.
-            debug_assert!(self.ctxt_or_zero == 0);
             let index = self.base_or_index;
             with_span_interner(|interner| interner.spans[index as usize])
         }
     }
+
+    /// This function is used as a fast path when decoding the full `SpanData` is not necessary.
+    #[inline]
+    pub fn ctxt(self) -> SyntaxContext {
+        let ctxt_or_max = self.ctxt_or_max as u32;
+        if ctxt_or_max < MAX_CTXT {
+            // Inline format or interned format with inline ctxt.
+            SyntaxContext::from_u32(ctxt_or_max)
+        } else {
+            // Interned format.
+            let index = self.base_or_index;
+            with_span_interner(|interner| interner.spans[index as usize].ctxt)
+        }
+    }
 }
 
 #[derive(Default)]