Merge pull request dtolnay#438 from buchgr/invalid-state

Add a function to reset thread-local span data.

Author: dtolnay

src/fallback.rs

@@ -320,17 +320,30 @@ impl Debug for SourceFile {
     }
 }
 
+#[cfg(all(span_locations, not(fuzzing)))]
+fn dummy_file() -> FileInfo {
+    FileInfo {
+        source_text: String::new(),
+        span: Span { lo: 0, hi: 0 },
+        lines: vec![0],
+        char_index_to_byte_offset: BTreeMap::new(),
+    }
+}
+
 #[cfg(all(span_locations, not(fuzzing)))]
 thread_local! {
     static SOURCE_MAP: RefCell<SourceMap> = RefCell::new(SourceMap {
         // Start with a single dummy file which all call_site() and def_site()
         // spans reference.
-        files: vec![FileInfo {
-            source_text: String::new(),
-            span: Span { lo: 0, hi: 0 },
-            lines: vec![0],
-            char_index_to_byte_offset: BTreeMap::new(),
-        }],
+        files: vec![dummy_file()],
+    });
+}
+
+#[cfg(all(span_locations, not(fuzzing)))]
+pub fn invalidate_current_thread_spans() {
+    SOURCE_MAP.with(|sm| {
+        let mut sm = sm.borrow_mut();
+        sm.files = vec![dummy_file()];
     });
 }
 

src/lib.rs

@@ -170,6 +170,10 @@ use std::error::Error;
 #[cfg(procmacro2_semver_exempt)]
 use std::path::PathBuf;
 
+#[cfg(span_locations)]
+#[cfg_attr(doc_cfg, doc(cfg(feature = "span-locations")))]
+pub use crate::imp::invalidate_current_thread_spans;
+
 #[cfg(span_locations)]
 #[cfg_attr(doc_cfg, doc(cfg(feature = "span-locations")))]
 pub use crate::location::LineColumn;
@@ -1325,4 +1329,4 @@ pub mod token_stream {
             }
         }
     }
-}
+}

src/wrapper.rs

@@ -928,3 +928,30 @@ impl Debug for Literal {
         }
     }
 }
+
+/// Invalidates any `proc_macro2::Span` on the current thread
+///
+/// The implementation of the `proc_macro2::Span` type relies on thread-local
+/// memory and this function clears it. Calling any method on a
+/// `proc_macro2::Span` on the current thread and that was created before this
+/// function was called will either lead to incorrect results or abort your
+/// program.
+///
+/// This function is useful for programs that process more than 2^32 bytes of
+/// text on a single thread. The internal representation of `proc_macro2::Span`
+/// uses 32-bit integers to represent offsets and those will overflow when
+/// processing more than 2^32 bytes. This function resets all offsets
+/// and thereby also invalidates any previously created `proc_macro2::Span`.
+/// 
+/// This function requires the `span-locations` feature to be enabled. This
+/// function is not applicable to and will panic if called from a procedural macro.
+#[cfg(span_locations)]
+pub fn invalidate_current_thread_spans() {
+    if inside_proc_macro() {
+        panic!(
+            "proc_macro2::invalidate_current_thread_spans is not available in procedural macros"
+        );
+    } else {
+        crate::fallback::invalidate_current_thread_spans()
+    }
+}

tests/test.rs

@@ -757,3 +757,40 @@ fn byte_order_mark() {
     let string = "foo\u{feff}";
     string.parse::<TokenStream>().unwrap_err();
 }
+
+// Creates a new Span from a TokenStream
+#[cfg(all(test, span_locations))]
+fn create_span() -> proc_macro2::Span {
+    let tts: TokenStream = "1".parse().unwrap();
+    match tts.into_iter().next().unwrap() {
+        TokenTree::Literal(literal) => literal.span(),
+        _ => unreachable!(),
+    }
+}
+
+#[cfg(span_locations)]
+#[test]
+fn test_invalidate_current_thread_spans() {
+    let actual = format!("{:#?}", create_span());
+    assert_eq!(actual, "bytes(1..2)");
+    let actual = format!("{:#?}", create_span());
+    assert_eq!(actual, "bytes(3..4)");
+
+    proc_macro2::invalidate_current_thread_spans();
+
+    let actual = format!("{:#?}", create_span());
+    // Test that span offsets have been reset after the call
+    // to invalidate_current_thread_spans()
+    assert_eq!(actual, "bytes(1..2)");
+}
+
+#[cfg(span_locations)]
+#[test]
+#[should_panic]
+fn test_use_span_after_invalidation() {
+    let span = create_span();
+
+    proc_macro2::invalidate_current_thread_spans();
+
+    span.source_text();
+}