Fix/Make proc_macro span_close() and span_open() more accurate after set_span() calls

compiler/rustc_ast/src/tokenstream.rs

@@ -978,7 +978,54 @@ pub struct DelimSpan {
 
 impl DelimSpan {
     pub fn from_single(sp: Span) -> Self {
-        DelimSpan { open: sp, close: sp }
+        let default_open = sp.shrink_to_lo();
+        let default_close = sp.shrink_to_hi();
+
+        let Some(sm) = rustc_span::source_map::get_source_map() else {
+            // No source map available.
+            return Self { open: default_open, close: default_close };
+        };
+
+        let (open, close) = sm
+            .span_to_source(sp, |src, start, end| {
+                let Some(src) = src.get(start..end) else {
+                    return Ok((default_close, default_close));
+                };
+
+                // Only check the first and last characters.
+                // If there is white space or other characters
+                // other than `( ... )`, `[ ... ]`, and `{ ... }`.
+                // I assume that is intentionally included in this
+                // span so we don't want to shrink the span by
+                // searching for the delimiters, and setting
+                // the open and close spans to some more interior
+                // position.
+                let mut chars = src.chars();
+                // just using '_' as a place holder.
+                let first = chars.next().unwrap_or('_');
+                let last = chars.last().unwrap_or('_');
+
+                // Thought maybe scan through if first is '(', '[',  or '{'
+                // and see if the last matches up e.g. make sure it's not some
+                // extra mismatched delimiter.
+
+                let open = sp.subspan(0..first.len_utf8()).unwrap_or(default_open);
+
+                let len = (sp.hi() - sp.lo()).0 as usize;
+                let pos = len.checked_sub(last.len_utf8()).unwrap_or(0);
+                let close = sp.subspan(pos..).unwrap_or(default_close);
+
+                Ok(match (first, last) {
+                    ('(', ')') | ('{', '}') | ('[', ']') => (open, close),
+                    ('(', _) | ('{', _) | ('[', _) => (open, default_close),
+                    (_, ')') | (_, '}') | (_, ']') => (default_open, close),
+                    (_, _) => (default_close, default_open),
+                })
+            })
+            .ok()
+            .unwrap_or((default_open, default_close));
+
+        DelimSpan { open, close }
     }
 
     pub fn from_pair(open: Span, close: Span) -> Self {

compiler/rustc_span/src/lib.rs

@@ -77,7 +77,7 @@ use std::cmp::{self, Ordering};
 use std::fmt::Display;
 use std::hash::Hash;
 use std::io::{self, Read};
-use std::ops::{Add, Range, Sub};
+use std::ops::{Add, Bound, Range, RangeBounds, Sub};
 use std::path::{Path, PathBuf};
 use std::str::FromStr;
 use std::sync::Arc;
@@ -1176,6 +1176,37 @@ impl Span {
     pub fn normalize_to_macro_rules(self) -> Span {
         self.map_ctxt(|ctxt| ctxt.normalize_to_macro_rules())
     }
+
+    /// This function is similar to `Span::from_inner`, but it
+    /// will return `None` if the relative Range span exceeds
+    /// the bounds of span.
+    pub fn subspan<I: Copy, R: RangeBounds<I>>(self, subspan: R) -> Option<Span>
+    where
+        u32: TryFrom<I>,
+    {
+        let lo = self.lo().0;
+        let hi = self.hi().0;
+
+        let start = match subspan.start_bound() {
+            Bound::Included(s) => u32::try_from(*s).ok()?,
+            Bound::Excluded(s) => u32::try_from(*s).ok()?.checked_add(1)?,
+            Bound::Unbounded => 0,
+        };
+
+        let end = match subspan.end_bound() {
+            Bound::Included(e) => u32::try_from(*e).ok()?.checked_add(1)?,
+            Bound::Excluded(e) => u32::try_from(*e).ok()?,
+            Bound::Unbounded => hi - lo,
+        };
+
+        let new_lo = lo.checked_add(start)?;
+        let new_hi = lo.checked_add(end)?;
+        if new_lo > hi || new_hi > hi {
+            return None;
+        }
+
+        Some(self.with_lo(BytePos(new_lo)).with_hi(BytePos(new_hi)))
+    }
 }
 
 impl Default for Span {

library/proc_macro/src/lib.rs

@@ -46,7 +46,7 @@ mod to_tokens;
 
 use core::ops::BitOr;
 use std::ffi::CStr;
-use std::ops::{Range, RangeBounds};
+use std::ops::{Bound, Range, RangeBounds};
 use std::path::PathBuf;
 use std::str::FromStr;
 use std::{error, fmt};
@@ -880,7 +880,20 @@ impl Group {
     /// tokens at the level of the `Group`.
     #[stable(feature = "proc_macro_lib2", since = "1.29.0")]
     pub fn set_span(&mut self, span: Span) {
-        self.0.span = bridge::DelimSpan::from_single(span.0);
+        // This is an invisible delimiter so they have no width, so just use start and end.
+        if matches!(self.0.delimiter, Delimiter::None) {
+            self.0.span =
+                bridge::DelimSpan { open: span.0.start(), close: span.0.end(), entire: span.0 };
+            return;
+        }
+
+        let rng = span.0.byte_range();
+        let beg = (rng.end - rng.start).saturating_sub(1);
+        self.0.span = bridge::DelimSpan {
+            open: span.0.subspan(Bound::Included(0), Bound::Excluded(1)).unwrap_or(span.0.start()),
+            close: span.0.subspan(Bound::Included(beg), Bound::Unbounded).unwrap_or(span.0.end()),
+            entire: span.0,
+        };
     }
 }