add windows one time initialization

5 files changed, 368 insertions, 32 deletions

diff --git a/src/tools/miri/src/concurrency/sync.rs b/src/tools/miri/src/concurrency/sync.rs
index 464f452ca76..18920951049 100644
--- a/src/tools/miri/src/concurrency/sync.rs
+++ b/src/tools/miri/src/concurrency/sync.rs
@@ -7,6 +7,7 @@ use log::trace;
 use rustc_data_structures::fx::FxHashMap;
 use rustc_index::vec::{Idx, IndexVec};
 
+use super::thread::MachineCallback;
 use super::vector_clock::VClock;
 use crate::*;
 
@@ -149,13 +150,68 @@ struct FutexWaiter {
     bitset: u32,
 }
 
+declare_id!(InitOnceId);
+
+struct InitOnceWaiter<'mir, 'tcx> {
+    thread: ThreadId,
+    callback: Box<dyn MachineCallback<'mir, 'tcx> + 'tcx>,
+}
+
+impl<'mir, 'tcx> std::fmt::Debug for InitOnceWaiter<'mir, 'tcx> {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("InitOnce")
+            .field("thread", &self.thread)
+            .field("callback", &"dyn MachineCallback")
+            .finish()
+    }
+}
+
+#[derive(Debug, Copy, Clone, PartialEq, Eq)]
+/// The current status of a one time initialization.
+pub enum InitOnceStatus {
+    Uninitialized,
+    Begun,
+    Complete,
+}
+
+impl Default for InitOnceStatus {
+    fn default() -> Self {
+        Self::Uninitialized
+    }
+}
+
+/// The one time initialization state.
+#[derive(Default, Debug)]
+struct InitOnce<'mir, 'tcx> {
+    status: InitOnceStatus,
+    waiters: VecDeque<InitOnceWaiter<'mir, 'tcx>>,
+    data_race: VClock,
+}
+
+impl<'mir, 'tcx> VisitTags for InitOnce<'mir, 'tcx> {
+    fn visit_tags(&self, visit: &mut dyn FnMut(SbTag)) {
+        for waiter in self.waiters.iter() {
+            waiter.callback.visit_tags(visit);
+        }
+    }
+}
+
 /// The state of all synchronization variables.
 #[derive(Default, Debug)]
-pub(crate) struct SynchronizationState {
+pub(crate) struct SynchronizationState<'mir, 'tcx> {
     mutexes: IndexVec<MutexId, Mutex>,
     rwlocks: IndexVec<RwLockId, RwLock>,
     condvars: IndexVec<CondvarId, Condvar>,
     futexes: FxHashMap<u64, Futex>,
+    init_onces: IndexVec<InitOnceId, InitOnce<'mir, 'tcx>>,
+}
+
+impl<'mir, 'tcx> VisitTags for SynchronizationState<'mir, 'tcx> {
+    fn visit_tags(&self, visit: &mut dyn FnMut(SbTag)) {
+        for init_once in self.init_onces.iter() {
+            init_once.visit_tags(visit);
+        }
+    }
 }
 
 // Private extension trait for local helper methods
@@ -581,4 +637,144 @@ pub trait EvalContextExt<'mir, 'tcx: 'mir>: crate::MiriInterpCxExt<'mir, 'tcx> {
             futex.waiters.retain(|waiter| waiter.thread != thread);
         }
     }
+
+    #[inline]
+    /// Create state for a new one time initialization.
+    fn init_once_create(&mut self) -> InitOnceId {
+        let this = self.eval_context_mut();
+        this.machine.threads.sync.init_onces.push(Default::default())
+    }
+
+    #[inline]
+    /// Provides the closure with the next InitOnceId. Creates that InitOnce if the closure returns None,
+    /// otherwise returns the value from the closure
+    fn init_once_get_or_create<F>(&mut self, existing: F) -> InterpResult<'tcx, InitOnceId>
+    where
+        F: FnOnce(
+            &mut MiriInterpCx<'mir, 'tcx>,
+            InitOnceId,
+        ) -> InterpResult<'tcx, Option<InitOnceId>>,
+    {
+        let this = self.eval_context_mut();
+        let next_index = this.machine.threads.sync.init_onces.next_index();
+        if let Some(old) = existing(this, next_index)? {
+            Ok(old)
+        } else {
+            let new_index = this.machine.threads.sync.init_onces.push(Default::default());
+            assert_eq!(next_index, new_index);
+            Ok(new_index)
+        }
+    }
+
+    #[inline]
+    fn init_once_status(&mut self, id: InitOnceId) -> InitOnceStatus {
+        let this = self.eval_context_ref();
+        this.machine.threads.sync.init_onces[id].status
+    }
+
+    #[inline]
+    /// Put the thread into the queue waiting for the initialization.
+    fn init_once_enqueue_and_block(
+        &mut self,
+        id: InitOnceId,
+        thread: ThreadId,
+        callback: Box<dyn MachineCallback<'mir, 'tcx> + 'tcx>,
+    ) {
+        let this = self.eval_context_mut();
+        let init_once = &mut this.machine.threads.sync.init_onces[id];
+        assert_ne!(init_once.status, InitOnceStatus::Complete, "queueing on complete init once");
+        init_once.waiters.push_back(InitOnceWaiter { thread, callback });
+        this.block_thread(thread);
+    }
+
+    #[inline]
+    fn init_once_begin(&mut self, id: InitOnceId) {
+        let this = self.eval_context_mut();
+        let init_once = &mut this.machine.threads.sync.init_onces[id];
+        assert_eq!(
+            init_once.status,
+            InitOnceStatus::Uninitialized,
+            "begining already begun or complete init once"
+        );
+        init_once.status = InitOnceStatus::Begun;
+    }
+
+    #[inline]
+    fn init_once_complete(&mut self, id: InitOnceId) -> InterpResult<'tcx> {
+        let this = self.eval_context_mut();
+        let current_thread = this.get_active_thread();
+        let init_once = &mut this.machine.threads.sync.init_onces[id];
+
+        assert_eq!(
+            init_once.status,
+            InitOnceStatus::Begun,
+            "completing already complete or uninit init once"
+        );
+
+        init_once.status = InitOnceStatus::Complete;
+
+        // Each complete happens-before the end of the wait
+        if let Some(data_race) = &this.machine.data_race {
+            data_race.validate_lock_release(&mut init_once.data_race, current_thread);
+        }
+
+        // need to take the queue to avoid having `this` be borrowed multiple times
+        for waiter in std::mem::take(&mut init_once.waiters) {
+            this.unblock_thread(waiter.thread);
+
+            this.set_active_thread(waiter.thread);
+            waiter.callback.call(this)?;
+            this.set_active_thread(current_thread);
+
+            if let Some(data_race) = &this.machine.data_race {
+                data_race.validate_lock_acquire(
+                    &this.machine.threads.sync.init_onces[id].data_race,
+                    waiter.thread,
+                );
+            }
+        }
+
+        Ok(())
+    }
+
+    #[inline]
+    fn init_once_fail(&mut self, id: InitOnceId) -> InterpResult<'tcx> {
+        let this = self.eval_context_mut();
+        let current_thread = this.get_active_thread();
+        let init_once = &mut this.machine.threads.sync.init_onces[id];
+        assert_eq!(
+            init_once.status,
+            InitOnceStatus::Begun,
+            "failing already completed or uninit init once"
+        );
+
+        // Each complete happens-before the end of the wait
+        if let Some(data_race) = &this.machine.data_race {
+            data_race.validate_lock_release(&mut init_once.data_race, current_thread);
+        }
+
+        // the behavior of failing the initialization is left vague by the docs
+        // it had to be determined experimentally
+        if let Some(waiter) = init_once.waiters.pop_front() {
+            // try initializing again on a different thread
+            init_once.status = InitOnceStatus::Begun;
+
+            this.unblock_thread(waiter.thread);
+
+            this.set_active_thread(waiter.thread);
+            waiter.callback.call(this)?;
+            this.set_active_thread(current_thread);
+
+            if let Some(data_race) = &this.machine.data_race {
+                data_race.validate_lock_acquire(
+                    &this.machine.threads.sync.init_onces[id].data_race,
+                    waiter.thread,
+                );
+            }
+        } else {
+            init_once.status = InitOnceStatus::Uninitialized;
+        }
+
+        Ok(())
+    }
 }
diff --git a/src/tools/miri/src/concurrency/thread.rs b/src/tools/miri/src/concurrency/thread.rs
index ec1da4138d4..3432f10f7a9 100644
--- a/src/tools/miri/src/concurrency/thread.rs
+++ b/src/tools/miri/src/concurrency/thread.rs
@@ -30,8 +30,7 @@ pub enum SchedulingAction {
     Stop,
 }
 
-/// Timeout callbacks can be created by synchronization primitives to tell the
-/// scheduler that they should be called once some period of time passes.
+/// Trait for callbacks that can be executed when some event happens, such as after a timeout.
 pub trait MachineCallback<'mir, 'tcx>: VisitTags {
     fn call(&self, ecx: &mut InterpCx<'mir, 'tcx, MiriMachine<'mir, 'tcx>>) -> InterpResult<'tcx>;
 }
@@ -269,7 +268,7 @@ pub struct ThreadManager<'mir, 'tcx> {
     threads: IndexVec<ThreadId, Thread<'mir, 'tcx>>,
     /// This field is pub(crate) because the synchronization primitives
     /// (`crate::sync`) need a way to access it.
-    pub(crate) sync: SynchronizationState,
+    pub(crate) sync: SynchronizationState<'mir, 'tcx>,
     /// A mapping from a thread-local static to an allocation id of a thread
     /// specific allocation.
     thread_local_alloc_ids: RefCell<FxHashMap<(DefId, ThreadId), Pointer<Provenance>>>,
@@ -303,7 +302,7 @@ impl VisitTags for ThreadManager<'_, '_> {
             timeout_callbacks,
             active_thread: _,
             yield_active_thread: _,
-            sync: _,
+            sync,
         } = self;
 
         for thread in threads {
@@ -315,6 +314,7 @@ impl VisitTags for ThreadManager<'_, '_> {
         for callback in timeout_callbacks.values() {
             callback.callback.visit_tags(visit);
         }
+        sync.visit_tags(visit);
     }
 }
 
diff --git a/src/tools/miri/src/lib.rs b/src/tools/miri/src/lib.rs
index 8de0b0413a4..f7c22b76f4f 100644
--- a/src/tools/miri/src/lib.rs
+++ b/src/tools/miri/src/lib.rs
@@ -87,7 +87,7 @@ pub use crate::concurrency::{
         AtomicFenceOrd, AtomicReadOrd, AtomicRwOrd, AtomicWriteOrd,
         EvalContextExt as DataRaceEvalContextExt,
     },
-    sync::{CondvarId, EvalContextExt as SyncEvalContextExt, MutexId, RwLockId},
+    sync::{CondvarId, EvalContextExt as SyncEvalContextExt, InitOnceId, MutexId, RwLockId},
     thread::{
         EvalContextExt as ThreadsEvalContextExt, SchedulingAction, ThreadId, ThreadManager,
         ThreadState, Time,
diff --git a/src/tools/miri/src/shims/windows/foreign_items.rs b/src/tools/miri/src/shims/windows/foreign_items.rs
index b0670358f9c..d998bdf420f 100644
--- a/src/tools/miri/src/shims/windows/foreign_items.rs
+++ b/src/tools/miri/src/shims/windows/foreign_items.rs
@@ -261,6 +261,18 @@ pub trait EvalContextExt<'mir, 'tcx: 'mir>: crate::MiriInterpCxExt<'mir, 'tcx> {
                 let ret = this.TryAcquireSRWLockShared(ptr)?;
                 this.write_scalar(ret, dest)?;
             }
+            "InitOnceBeginInitialize" => {
+                let [ptr, flags, pending, context] =
+                    this.check_shim(abi, Abi::System { unwind: false }, link_name, args)?;
+                let result = this.InitOnceBeginInitialize(ptr, flags, pending, context)?;
+                this.write_scalar(result, dest)?;
+            }
+            "InitOnceComplete" => {
+                let [ptr, flags, context] =
+                    this.check_shim(abi, Abi::System { unwind: false }, link_name, args)?;
+                let result = this.InitOnceComplete(ptr, flags, context)?;
+                this.write_scalar(result, dest)?;
+            }
 
             // Dynamic symbol loading
             "GetProcAddress" => {
diff --git a/src/tools/miri/src/shims/windows/sync.rs b/src/tools/miri/src/shims/windows/sync.rs
index 88b117c54be..feed90fad16 100644
--- a/src/tools/miri/src/shims/windows/sync.rs
+++ b/src/tools/miri/src/shims/windows/sync.rs
@@ -1,20 +1,24 @@
+use crate::concurrency::sync::InitOnceStatus;
+use crate::concurrency::thread::MachineCallback;
 use crate::*;
 
-// Locks are pointer-sized pieces of data, initialized to 0.
-// We use the first 4 bytes to store the RwLockId.
-
-fn srwlock_get_or_create_id<'mir, 'tcx: 'mir>(
-    ecx: &mut MiriInterpCx<'mir, 'tcx>,
-    lock_op: &OpTy<'tcx, Provenance>,
-) -> InterpResult<'tcx, RwLockId> {
-    let value_place = ecx.deref_operand_and_offset(lock_op, 0, ecx.machine.layouts.u32)?;
+impl<'mir, 'tcx> EvalContextExtPriv<'mir, 'tcx> for crate::MiriInterpCx<'mir, 'tcx> {}
+trait EvalContextExtPriv<'mir, 'tcx: 'mir>: crate::MiriInterpCxExt<'mir, 'tcx> {
+    // These synchronization structures are pointer-sized pieces of data, initialized to 0.
+    // We use the first 4 bytes to store the id.
+    fn get_or_create_id(
+        &mut self,
+        next_id: Scalar<Provenance>,
+        lock_op: &OpTy<'tcx, Provenance>,
+    ) -> InterpResult<'tcx, Option<u32>> {
+        let this = self.eval_context_mut();
+        let value_place = this.deref_operand_and_offset(lock_op, 0, this.machine.layouts.u32)?;
 
-    ecx.rwlock_get_or_create(|ecx, next_id| {
-        let (old, success) = ecx
+        let (old, success) = this
             .atomic_compare_exchange_scalar(
                 &value_place,
-                &ImmTy::from_uint(0u32, ecx.machine.layouts.u32),
-                next_id.to_u32_scalar(),
+                &ImmTy::from_uint(0u32, this.machine.layouts.u32),
+                next_id,
                 AtomicRwOrd::Relaxed,
                 AtomicReadOrd::Relaxed,
                 false,
@@ -25,17 +29,38 @@ fn srwlock_get_or_create_id<'mir, 'tcx: 'mir>(
             // Caller of the closure needs to allocate next_id
             None
         } else {
-            Some(RwLockId::from_u32(old.to_u32().expect("layout is u32")))
+            Some(old.to_u32().expect("layout is u32"))
+        })
+    }
+
+    fn srwlock_get_or_create_id(
+        &mut self,
+        lock_op: &OpTy<'tcx, Provenance>,
+    ) -> InterpResult<'tcx, RwLockId> {
+        let this = self.eval_context_mut();
+        this.rwlock_get_or_create(|ecx, next_id| {
+            Ok(ecx.get_or_create_id(next_id.to_u32_scalar(), lock_op)?.map(RwLockId::from_u32))
+        })
+    }
+
+    fn init_once_get_or_create_id(
+        &mut self,
+        lock_op: &OpTy<'tcx, Provenance>,
+    ) -> InterpResult<'tcx, InitOnceId> {
+        let this = self.eval_context_mut();
+        this.init_once_get_or_create(|ecx, next_id| {
+            Ok(ecx.get_or_create_id(next_id.to_u32_scalar(), lock_op)?.map(InitOnceId::from_u32))
         })
-    })
+    }
 }
 
 impl<'mir, 'tcx> EvalContextExt<'mir, 'tcx> for crate::MiriInterpCx<'mir, 'tcx> {}
+
+#[allow(non_snake_case)]
 pub trait EvalContextExt<'mir, 'tcx: 'mir>: crate::MiriInterpCxExt<'mir, 'tcx> {
-    #[allow(non_snake_case)]
     fn AcquireSRWLockExclusive(&mut self, lock_op: &OpTy<'tcx, Provenance>) -> InterpResult<'tcx> {
         let this = self.eval_context_mut();
-        let id = srwlock_get_or_create_id(this, lock_op)?;
+        let id = this.srwlock_get_or_create_id(lock_op)?;
         let active_thread = this.get_active_thread();
 
         if this.rwlock_is_locked(id) {
@@ -54,13 +79,12 @@ pub trait EvalContextExt<'mir, 'tcx: 'mir>: crate::MiriInterpCxExt<'mir, 'tcx> {
         Ok(())
     }
 
-    #[allow(non_snake_case)]
     fn TryAcquireSRWLockExclusive(
         &mut self,
         lock_op: &OpTy<'tcx, Provenance>,
     ) -> InterpResult<'tcx, Scalar<Provenance>> {
         let this = self.eval_context_mut();
-        let id = srwlock_get_or_create_id(this, lock_op)?;
+        let id = this.srwlock_get_or_create_id(lock_op)?;
         let active_thread = this.get_active_thread();
 
         if this.rwlock_is_locked(id) {
@@ -72,10 +96,9 @@ pub trait EvalContextExt<'mir, 'tcx: 'mir>: crate::MiriInterpCxExt<'mir, 'tcx> {
         }
     }
 
-    #[allow(non_snake_case)]
     fn ReleaseSRWLockExclusive(&mut self, lock_op: &OpTy<'tcx, Provenance>) -> InterpResult<'tcx> {
         let this = self.eval_context_mut();
-        let id = srwlock_get_or_create_id(this, lock_op)?;
+        let id = this.srwlock_get_or_create_id(lock_op)?;
         let active_thread = this.get_active_thread();
 
         if !this.rwlock_writer_unlock(id, active_thread) {
@@ -88,10 +111,9 @@ pub trait EvalContextExt<'mir, 'tcx: 'mir>: crate::MiriInterpCxExt<'mir, 'tcx> {
         Ok(())
     }
 
-    #[allow(non_snake_case)]
     fn AcquireSRWLockShared(&mut self, lock_op: &OpTy<'tcx, Provenance>) -> InterpResult<'tcx> {
         let this = self.eval_context_mut();
-        let id = srwlock_get_or_create_id(this, lock_op)?;
+        let id = this.srwlock_get_or_create_id(lock_op)?;
         let active_thread = this.get_active_thread();
 
         if this.rwlock_is_write_locked(id) {
@@ -103,13 +125,12 @@ pub trait EvalContextExt<'mir, 'tcx: 'mir>: crate::MiriInterpCxExt<'mir, 'tcx> {
         Ok(())
     }
 
-    #[allow(non_snake_case)]
     fn TryAcquireSRWLockShared(
         &mut self,
         lock_op: &OpTy<'tcx, Provenance>,
     ) -> InterpResult<'tcx, Scalar<Provenance>> {
         let this = self.eval_context_mut();
-        let id = srwlock_get_or_create_id(this, lock_op)?;
+        let id = this.srwlock_get_or_create_id(lock_op)?;
         let active_thread = this.get_active_thread();
 
         if this.rwlock_is_write_locked(id) {
@@ -120,10 +141,9 @@ pub trait EvalContextExt<'mir, 'tcx: 'mir>: crate::MiriInterpCxExt<'mir, 'tcx> {
         }
     }
 
-    #[allow(non_snake_case)]
     fn ReleaseSRWLockShared(&mut self, lock_op: &OpTy<'tcx, Provenance>) -> InterpResult<'tcx> {
         let this = self.eval_context_mut();
-        let id = srwlock_get_or_create_id(this, lock_op)?;
+        let id = this.srwlock_get_or_create_id(lock_op)?;
         let active_thread = this.get_active_thread();
 
         if !this.rwlock_reader_unlock(id, active_thread) {
@@ -135,4 +155,112 @@ pub trait EvalContextExt<'mir, 'tcx: 'mir>: crate::MiriInterpCxExt<'mir, 'tcx> {
 
         Ok(())
     }
+
+    fn InitOnceBeginInitialize(
+        &mut self,
+        init_once_op: &OpTy<'tcx, Provenance>,
+        flags_op: &OpTy<'tcx, Provenance>,
+        pending_op: &OpTy<'tcx, Provenance>,
+        context_op: &OpTy<'tcx, Provenance>,
+    ) -> InterpResult<'tcx, Scalar<Provenance>> {
+        let this = self.eval_context_mut();
+        let active_thread = this.get_active_thread();
+
+        let id = this.init_once_get_or_create_id(init_once_op)?;
+        let flags = this.read_scalar(flags_op)?.to_u32()?;
+        let pending_place = this.deref_operand(pending_op)?.into();
+        let context = this.read_pointer(context_op)?;
+
+        if flags != 0 {
+            throw_unsup_format!("unsupported `dwFlags` {flags} in `InitOnceBeginInitialize`");
+        }
+
+        if !this.ptr_is_null(context)? {
+            throw_unsup_format!("non-null `lpContext` in `InitOnceBeginInitialize`");
+        }
+
+        struct Callback<'tcx> {
+            init_once_id: InitOnceId,
+            pending_place: PlaceTy<'tcx, Provenance>,
+        }
+
+        impl<'tcx> VisitTags for Callback<'tcx> {
+            fn visit_tags(&self, visit: &mut dyn FnMut(SbTag)) {
+                let Callback { init_once_id: _, pending_place } = self;
+                pending_place.visit_tags(visit);
+            }
+        }
+
+        impl<'mir, 'tcx> MachineCallback<'mir, 'tcx> for Callback<'tcx> {
+            fn call(&self, this: &mut MiriInterpCx<'mir, 'tcx>) -> InterpResult<'tcx> {
+                let pending = match this.init_once_status(self.init_once_id) {
+                    InitOnceStatus::Uninitialized =>
+                        unreachable!("status should have either been set to begun or complete"),
+                    InitOnceStatus::Begun => this.eval_windows("c", "TRUE")?,
+                    InitOnceStatus::Complete => this.eval_windows("c", "FALSE")?,
+                };
+
+                this.write_scalar(pending, &self.pending_place)?;
+
+                Ok(())
+            }
+        }
+
+        match this.init_once_status(id) {
+            InitOnceStatus::Uninitialized => {
+                this.init_once_begin(id);
+                this.write_scalar(this.eval_windows("c", "TRUE")?, &pending_place)?;
+            }
+            InitOnceStatus::Begun =>
+                this.init_once_enqueue_and_block(
+                    id,
+                    active_thread,
+                    Box::new(Callback { init_once_id: id, pending_place }),
+                ),
+            InitOnceStatus::Complete =>
+                this.write_scalar(this.eval_windows("c", "FALSE")?, &pending_place)?,
+        }
+
+        Ok(Scalar::from_i32(1))
+    }
+
+    fn InitOnceComplete(
+        &mut self,
+        init_once_op: &OpTy<'tcx, Provenance>,
+        flags_op: &OpTy<'tcx, Provenance>,
+        context_op: &OpTy<'tcx, Provenance>,
+    ) -> InterpResult<'tcx, Scalar<Provenance>> {
+        let this = self.eval_context_mut();
+
+        let id = this.init_once_get_or_create_id(init_once_op)?;
+        let flags = this.read_scalar(flags_op)?.to_u32()?;
+        let context = this.read_pointer(context_op)?;
+
+        let success = if flags == 0 {
+            true
+        } else if flags == this.eval_windows("c", "INIT_ONCE_INIT_FAILED")?.to_u32()? {
+            false
+        } else {
+            throw_unsup_format!("unsupported `dwFlags` {flags} in `InitOnceBeginInitialize`");
+        };
+
+        if !this.ptr_is_null(context)? {
+            throw_unsup_format!("non-null `lpContext` in `InitOnceBeginInitialize`");
+        }
+
+        if this.init_once_status(id) != InitOnceStatus::Begun {
+            // The docs do not say anything about this case, but it seems better to not allow it.
+            throw_ub_format!(
+                "calling InitOnceComplete on a one time initialization that has not begun or is already completed"
+            );
+        }
+
+        if success {
+            this.init_once_complete(id)?;
+        } else {
+            this.init_once_fail(id)?;
+        }
+
+        Ok(Scalar::from_i32(1))
+    }
 }