Merge pull request #412 from dtolnay/sourcetext

Cache and lazily build the mapping from char index to byte offset

Author: dtolnay

src/fallback.rs

@@ -4,6 +4,8 @@ use crate::parse::{self, Cursor};
 use crate::rcvec::{RcVec, RcVecBuilder, RcVecIntoIter, RcVecMut};
 use crate::{Delimiter, Spacing, TokenTree};
 #[cfg(all(span_locations, not(fuzzing)))]
+use alloc::collections::BTreeMap;
+#[cfg(all(span_locations, not(fuzzing)))]
 use core::cell::RefCell;
 #[cfg(span_locations)]
 use core::cmp;
@@ -327,6 +329,7 @@ thread_local! {
             source_text: String::new(),
             span: Span { lo: 0, hi: 0 },
             lines: vec![0],
+            char_index_to_byte_offset: BTreeMap::new(),
         }],
     });
 }
@@ -336,6 +339,7 @@ struct FileInfo {
     source_text: String,
     span: Span,
     lines: Vec<usize>,
+    char_index_to_byte_offset: BTreeMap<usize, usize>,
 }
 
 #[cfg(all(span_locations, not(fuzzing)))]
@@ -362,12 +366,34 @@ impl FileInfo {
         span.lo >= self.span.lo && span.hi <= self.span.hi
     }
 
-    fn source_text(&self, span: Span) -> String {
-        let lo = (span.lo - self.span.lo) as usize;
-        let trunc_lo = match self.source_text.char_indices().nth(lo) {
-            Some((offset, _ch)) => &self.source_text[offset..],
-            None => return String::new(),
+    fn source_text(&mut self, span: Span) -> String {
+        let lo_char = (span.lo - self.span.lo) as usize;
+
+        // Look up offset of the largest already-computed char index that is
+        // less than or equal to the current requested one. We resume counting
+        // chars from that point.
+        let (&last_char_index, &last_byte_offset) = self
+            .char_index_to_byte_offset
+            .range(..=lo_char)
+            .next_back()
+            .unwrap_or((&0, &0));
+
+        let lo_byte = if last_char_index == lo_char {
+            last_byte_offset
+        } else {
+            let total_byte_offset = match self.source_text[last_byte_offset..]
+                .char_indices()
+                .nth(lo_char - last_char_index)
+            {
+                Some((additional_offset, _ch)) => last_byte_offset + additional_offset,
+                None => self.source_text.len(),
+            };
+            self.char_index_to_byte_offset
+                .insert(lo_char, total_byte_offset);
+            total_byte_offset
         };
+
+        let trunc_lo = &self.source_text[lo_byte..];
         let char_len = (span.hi - span.lo) as usize;
         let source_text = match trunc_lo.char_indices().nth(char_len) {
             Some((offset, _ch)) => &trunc_lo[..offset],
@@ -421,6 +447,8 @@ impl SourceMap {
             source_text: src.to_owned(),
             span,
             lines,
+            // Populated lazily by source_text().
+            char_index_to_byte_offset: BTreeMap::new(),
         });
 
         span
@@ -448,6 +476,15 @@ impl SourceMap {
         }
         unreachable!("Invalid span with no related FileInfo!");
     }
+
+    fn fileinfo_mut(&mut self, span: Span) -> &mut FileInfo {
+        for file in &mut self.files {
+            if file.span_within(span) {
+                return file;
+            }
+        }
+        unreachable!("Invalid span with no related FileInfo!");
+    }
 }
 
 #[derive(Clone, Copy, PartialEq, Eq)]
@@ -572,7 +609,7 @@ impl Span {
             if self.is_call_site() {
                 None
             } else {
-                Some(SOURCE_MAP.with(|cm| cm.borrow().fileinfo(*self).source_text(*self)))
+                Some(SOURCE_MAP.with(|cm| cm.borrow_mut().fileinfo_mut(*self).source_text(*self)))
             }
         }
     }

tests/test.rs

@@ -328,17 +328,19 @@ fn literal_span() {
 #[cfg(span_locations)]
 #[test]
 fn source_text() {
-    let input = "    𓀕 c    ";
+    let input = "    𓀕 a z    ";
     let mut tokens = input
         .parse::<proc_macro2::TokenStream>()
         .unwrap()
         .into_iter();
 
-    let ident = tokens.next().unwrap();
-    assert_eq!("𓀕", ident.span().source_text().unwrap());
+    let first = tokens.next().unwrap();
+    assert_eq!("𓀕", first.span().source_text().unwrap());
 
-    let ident = tokens.next().unwrap();
-    assert_eq!("c", ident.span().source_text().unwrap());
+    let second = tokens.next().unwrap();
+    let third = tokens.next().unwrap();
+    assert_eq!("z", third.span().source_text().unwrap());
+    assert_eq!("a", second.span().source_text().unwrap());
 }
 
 #[test]