From 2a073bfe00e9527aca48260abfb9ebc03f74b497 Mon Sep 17 00:00:00 2001
From: Jon Gjengset <jon@thesquareplanet.com>
Date: Thu, 9 Apr 2020 12:53:52 -0400
Subject: [PATCH] Track location for atomic new/drop

---
 src/rt/atomic.rs | 4 ++++
 src/rt/mod.rs    | 8 ++++----
 tests/atomic.rs  | 8 ++++++++
 3 files changed, 16 insertions(+), 4 deletions(-)

diff --git a/src/rt/atomic.rs b/src/rt/atomic.rs
index 0c60f798..5b5baa87 100644
--- a/src/rt/atomic.rs
+++ b/src/rt/atomic.rs
@@ -313,6 +313,9 @@ impl<T: Numeric> Atomic<T> {
                     state.stores[index].value = T::into_u64(self.0);
 
                     if !std::thread::panicking() {
+                        state
+                            .unsync_mut_locations
+                            .track(location!(), &execution.threads);
                         state.track_unsync_mut(&execution.threads);
                     }
                 });
@@ -358,6 +361,7 @@ impl State {
         };
 
         // All subsequent accesses must happen-after.
+        state.unsync_mut_locations.track(location, &threads);
         state.track_unsync_mut(threads);
 
         // Store the initial thread
diff --git a/src/rt/mod.rs b/src/rt/mod.rs
index 52fd55bf..a24e1bb0 100644
--- a/src/rt/mod.rs
+++ b/src/rt/mod.rs
@@ -1,3 +1,7 @@
+#[macro_use]
+mod location;
+pub(crate) use self::location::Location;
+
 mod access;
 use self::access::Access;
 
@@ -10,10 +14,6 @@ pub(crate) use self::arc::Arc;
 mod atomic;
 pub(crate) use self::atomic::{fence, Atomic};
 
-#[macro_use]
-mod location;
-pub(crate) use self::location::Location;
-
 mod cell;
 pub(crate) use self::cell::Cell;
 
diff --git a/tests/atomic.rs b/tests/atomic.rs
index 9150e0b0..f561b21e 100644
--- a/tests/atomic.rs
+++ b/tests/atomic.rs
@@ -6,6 +6,14 @@ use loom::thread;
 use std::sync::atomic::Ordering::{AcqRel, Acquire, Relaxed, Release};
 use std::sync::Arc;
 
+#[test]
+fn load_after_new() {
+    loom::model(|| {
+        let a = AtomicUsize::new(0);
+        a.load(Relaxed);
+    });
+}
+
 #[test]
 #[should_panic]
 fn invalid_unsync_load_relaxed() {