Add a virtual wrapper for &Latch

Author: cuviper

rayon-core/src/broadcast/mod.rs

@@ -1,4 +1,5 @@
 use crate::job::{ArcJob, StackJob};
+use crate::latch::LatchRef;
 use crate::registry::{Registry, WorkerThread};
 use crate::scope::ScopeLatch;
 use std::fmt;
@@ -107,7 +108,9 @@ where
     let n_threads = registry.num_threads();
     let current_thread = WorkerThread::current().as_ref();
     let latch = ScopeLatch::with_count(n_threads, current_thread);
-    let jobs: Vec<_> = (0..n_threads).map(|_| StackJob::new(&f, &latch)).collect();
+    let jobs: Vec<_> = (0..n_threads)
+        .map(|_| StackJob::new(&f, LatchRef::new(&latch)))
+        .collect();
     let job_refs = jobs.iter().map(|job| job.as_job_ref());
 
     registry.inject_broadcast(job_refs);

rayon-core/src/latch.rs

@@ -1,3 +1,5 @@
+use std::marker::PhantomData;
+use std::ops::Deref;
 use std::sync::atomic::{AtomicUsize, Ordering};
 use std::sync::{Arc, Condvar, Mutex};
 use std::usize;
@@ -378,12 +380,35 @@ impl Latch for CountLockLatch {
     }
 }
 
-impl<'a, L> Latch for &'a L
-where
-    L: Latch,
-{
+/// `&L` without any implication of `dereferenceable` for `Latch::set`
+pub(super) struct LatchRef<'a, L> {
+    inner: *const L,
+    marker: PhantomData<&'a L>,
+}
+
+impl<L> LatchRef<'_, L> {
+    pub(super) fn new(inner: &L) -> LatchRef<'_, L> {
+        LatchRef {
+            inner,
+            marker: PhantomData,
+        }
+    }
+}
+
+unsafe impl<L: Sync> Sync for LatchRef<'_, L> {}
+
+impl<L> Deref for LatchRef<'_, L> {
+    type Target = L;
+
+    fn deref(&self) -> &L {
+        // SAFETY: if we have &self, the inner latch is still alive
+        unsafe { &*self.inner }
+    }
+}
+
+impl<L: Latch> Latch for LatchRef<'_, L> {
     #[inline]
     unsafe fn set(this: *const Self) {
-        L::set(&**this);
+        L::set((*this).inner);
     }
 }

rayon-core/src/registry.rs

@@ -1,5 +1,5 @@
 use crate::job::{JobFifo, JobRef, StackJob};
-use crate::latch::{AsCoreLatch, CoreLatch, CountLatch, Latch, LockLatch, SpinLatch};
+use crate::latch::{AsCoreLatch, CoreLatch, CountLatch, Latch, LatchRef, LockLatch, SpinLatch};
 use crate::log::Event::*;
 use crate::log::Logger;
 use crate::sleep::Sleep;
@@ -505,7 +505,7 @@ impl Registry {
                     assert!(injected && !worker_thread.is_null());
                     op(&*worker_thread, true)
                 },
-                l,
+                LatchRef::new(l),
             );
             self.inject(&[job.as_job_ref()]);
             job.latch.wait_and_reset(); // Make sure we can use the same latch again next time.