add tests for track_caller in closures and generators

1 files changed, 148 insertions, 5 deletions

diff --git a/src/tools/miri/tests/pass/track-caller-attribute.rs b/src/tools/miri/tests/pass/track-caller-attribute.rs
index d0473f746cd..1b0226e61b5 100644
--- a/src/tools/miri/tests/pass/track-caller-attribute.rs
+++ b/src/tools/miri/tests/pass/track-caller-attribute.rs
@@ -1,17 +1,25 @@
 #![feature(core_intrinsics)]
+#![feature(stmt_expr_attributes)]
+#![feature(closure_track_caller)]
+#![feature(generator_trait)]
+#![feature(generators)]
 
+use std::ops::{Generator, GeneratorState};
 use std::panic::Location;
+use std::pin::Pin;
+
+type Loc = &'static Location<'static>;
 
 #[track_caller]
-fn tracked() -> &'static Location<'static> {
+fn tracked() -> Loc {
     Location::caller() // most importantly, we never get line 7
 }
 
-fn nested_intrinsic() -> &'static Location<'static> {
+fn nested_intrinsic() -> Loc {
     Location::caller()
 }
 
-fn nested_tracked() -> &'static Location<'static> {
+fn nested_tracked() -> Loc {
     tracked()
 }
 
@@ -36,12 +44,12 @@ fn test_basic() {
 
     let nested = nested_intrinsic();
     assert_eq!(nested.file(), file!());
-    assert_eq!(nested.line(), 11);
+    assert_eq!(nested.line(), 19);
     assert_eq!(nested.column(), 5);
 
     let contained = nested_tracked();
     assert_eq!(contained.file(), file!());
-    assert_eq!(contained.line(), 15);
+    assert_eq!(contained.line(), 23);
     assert_eq!(contained.column(), 5);
 
     // `Location::caller()` in a macro should behave similarly to `file!` and `line!`,
@@ -125,9 +133,144 @@ fn test_trait_obj2() {
     assert_eq!(loc.line(), expected_line);
 }
 
+fn test_closure() {
+    #[track_caller]
+    fn mono_invoke_fn<F: Fn(&'static str, bool) -> (&'static str, bool, Loc)>(
+        val: &F,
+    ) -> (&'static str, bool, Loc) {
+        val("from_mono", false)
+    }
+
+    #[track_caller]
+    fn mono_invoke_fn_once<F: FnOnce(&'static str, bool) -> (&'static str, bool, Loc)>(
+        val: F,
+    ) -> (&'static str, bool, Loc) {
+        val("from_mono", false)
+    }
+
+    #[track_caller]
+    fn dyn_invoke_fn_mut(
+        val: &mut dyn FnMut(&'static str, bool) -> (&'static str, bool, Loc),
+    ) -> (&'static str, bool, Loc) {
+        val("from_dyn", false)
+    }
+
+    #[track_caller]
+    fn dyn_invoke_fn_once(
+        val: Box<dyn FnOnce(&'static str, bool) -> (&'static str, bool, Loc)>,
+    ) -> (&'static str, bool, Loc) {
+        val("from_dyn", false)
+    }
+
+    let mut track_closure = #[track_caller]
+    |first: &'static str, second: bool| (first, second, Location::caller());
+    let (first_arg, first_bool, first_loc) = track_closure("first_arg", true);
+    let first_line = line!() - 1;
+    assert_eq!(first_arg, "first_arg");
+    assert_eq!(first_bool, true);
+    assert_eq!(first_loc.file(), file!());
+    assert_eq!(first_loc.line(), first_line);
+    assert_eq!(first_loc.column(), 46);
+
+    let (dyn_arg, dyn_bool, dyn_loc) = dyn_invoke_fn_mut(&mut track_closure);
+    assert_eq!(dyn_arg, "from_dyn");
+    assert_eq!(dyn_bool, false);
+    // `FnMut::call_mut` does not have `#[track_caller]`,
+    // so this will not match
+    assert_ne!(dyn_loc.file(), file!());
+
+    let (dyn_arg, dyn_bool, dyn_loc) = dyn_invoke_fn_once(Box::new(track_closure));
+    assert_eq!(dyn_arg, "from_dyn");
+    assert_eq!(dyn_bool, false);
+    // `FnOnce::call_once` does not have `#[track_caller]`
+    // so this will not match
+    assert_ne!(dyn_loc.file(), file!());
+
+    let (mono_arg, mono_bool, mono_loc) = mono_invoke_fn(&track_closure);
+    let mono_line = line!() - 1;
+    assert_eq!(mono_arg, "from_mono");
+    assert_eq!(mono_bool, false);
+    assert_eq!(mono_loc.file(), file!());
+    assert_eq!(mono_loc.line(), mono_line);
+    assert_eq!(mono_loc.column(), 43);
+
+    let (mono_arg, mono_bool, mono_loc) = mono_invoke_fn_once(track_closure);
+    let mono_line = line!() - 1;
+    assert_eq!(mono_arg, "from_mono");
+    assert_eq!(mono_bool, false);
+    assert_eq!(mono_loc.file(), file!());
+    assert_eq!(mono_loc.line(), mono_line);
+    assert_eq!(mono_loc.column(), 43);
+
+    let non_tracked_caller = || Location::caller();
+    let non_tracked_line = line!() - 1; // This is the line of the closure, not its caller
+    let non_tracked_loc = non_tracked_caller();
+    assert_eq!(non_tracked_loc.file(), file!());
+    assert_eq!(non_tracked_loc.line(), non_tracked_line);
+    assert_eq!(non_tracked_loc.column(), 33);
+}
+
+fn test_generator() {
+    #[track_caller]
+    fn mono_generator<F: Generator<String, Yield = (&'static str, String, Loc), Return = ()>>(
+        val: Pin<&mut F>,
+    ) -> (&'static str, String, Loc) {
+        match val.resume("Mono".to_string()) {
+            GeneratorState::Yielded(val) => val,
+            _ => unreachable!(),
+        }
+    }
+
+    #[track_caller]
+    fn dyn_generator(
+        val: Pin<&mut dyn Generator<String, Yield = (&'static str, String, Loc), Return = ()>>,
+    ) -> (&'static str, String, Loc) {
+        match val.resume("Dyn".to_string()) {
+            GeneratorState::Yielded(val) => val,
+            _ => unreachable!(),
+        }
+    }
+
+    #[rustfmt::skip]
+    let generator = #[track_caller] |arg: String| {
+        yield ("first", arg.clone(), Location::caller());
+        yield ("second", arg.clone(), Location::caller());
+    };
+
+    let mut pinned = Box::pin(generator);
+    let (dyn_ret, dyn_arg, dyn_loc) = dyn_generator(pinned.as_mut());
+    assert_eq!(dyn_ret, "first");
+    assert_eq!(dyn_arg, "Dyn".to_string());
+    // The `Generator` trait does not have `#[track_caller]` on `resume`, so
+    // this will not match.
+    assert_ne!(dyn_loc.file(), file!());
+
+    let (mono_ret, mono_arg, mono_loc) = mono_generator(pinned.as_mut());
+    let mono_line = line!() - 1;
+    assert_eq!(mono_ret, "second");
+    // The generator ignores the argument to the second `resume` call
+    assert_eq!(mono_arg, "Dyn".to_string());
+    assert_eq!(mono_loc.file(), file!());
+    assert_eq!(mono_loc.line(), mono_line);
+    assert_eq!(mono_loc.column(), 42);
+
+    #[rustfmt::skip]
+    let non_tracked_generator = || { yield Location::caller(); };
+    let non_tracked_line = line!() - 1; // This is the line of the generator, not its caller
+    let non_tracked_loc = match Box::pin(non_tracked_generator).as_mut().resume(()) {
+        GeneratorState::Yielded(val) => val,
+        _ => unreachable!(),
+    };
+    assert_eq!(non_tracked_loc.file(), file!());
+    assert_eq!(non_tracked_loc.line(), non_tracked_line);
+    assert_eq!(non_tracked_loc.column(), 44);
+}
+
 fn main() {
     test_basic();
     test_fn_ptr();
     test_trait_obj();
     test_trait_obj2();
+    test_closure();
+    test_generator();
 }