Skip to main content

ksync/
kmutex.rs

1// Copyright 2026 The Fuchsia Authors. All rights reserved.
2// Use of this source code is governed by a BSD-style license that can be
3// found in the LICENSE file.
4
5use crate::{ExclusiveLockPolicy, LockPolicy, LockToken, RawLock, RawMutex};
6use core::marker::PhantomData;
7use core::pin::Pin;
8use core::ptr::addr_of_mut;
9use lockdep::LockClass;
10use pin_init::{PinInit, pin_data, pin_init, pin_init_from_closure, pinned_drop};
11
12/// A safe, Zircon-compatible mutual exclusion lock supporting compile-time order validation.
13///
14/// `KMutex` wraps a platform-specific `RawLock` abstraction. It is pinned in memory to support FFI
15/// loop-detector active list registrations safely under the lock class `Class`.
16#[repr(transparent)] // Ensure KMutex has the same layout as the underlying RawLock M.
17#[pin_data]
18pub struct KMutex<Class: LockClass, M = RawMutex> {
19    #[pin]
20    mutex: M,
21    _marker: PhantomData<Class>,
22}
23
24impl<Class: LockClass, M> KMutex<Class, M> {
25    /// Create a new KMutex with a pre-initialized raw lock.
26    pub const fn new(mutex: M) -> Self {
27        Self { mutex, _marker: PhantomData }
28    }
29}
30
31impl<Class: LockClass, M: RawLock> KMutex<Class, M> {
32    /// Safe dynamic initialization of the validation lock inside pin context.
33    pub fn init() -> impl PinInit<Self, core::convert::Infallible> {
34        pin_init!(Self {
35            // SAFETY: `Self::class_id()` is either null or a valid static `LockClassId`.
36            mutex <- unsafe { M::init(Self::class_id()) },
37            _marker: PhantomData,
38        })
39    }
40
41    /// Acquires the lock, using the default policy, and registers the active loop node.
42    #[inline]
43    pub fn lock(&self) -> impl PinInit<KMutexGuard<'_, Class, M>, core::convert::Infallible>
44    where
45        <M as RawLock>::DefaultPolicy: ExclusiveLockPolicy<M, AcquireArgs = ()>,
46    {
47        KMutexGuard::new(self)
48    }
49
50    /// Acquires the lock with arguments, using the default policy, and registers the active loop
51    /// node.
52    #[inline]
53    pub fn lock_with(
54        &self,
55        args: <<M as RawLock>::DefaultPolicy as LockPolicy<M>>::AcquireArgs,
56    ) -> impl PinInit<KMutexGuard<'_, Class, M>, core::convert::Infallible> {
57        KMutexGuard::new_with_args(self, args)
58    }
59
60    /// Acquires the lock, using the specified policy, and registers the active loop node.
61    #[inline]
62    pub fn lock_policy<P: ExclusiveLockPolicy<M, AcquireArgs = ()>>(
63        &self,
64    ) -> impl PinInit<KMutexGuard<'_, Class, M, P>, core::convert::Infallible> {
65        KMutexGuard::new(self)
66    }
67
68    /// Acquires the lock with arguments, using the specified policy, and registers the active loop
69    /// node.
70    #[inline]
71    pub fn lock_policy_with<P: ExclusiveLockPolicy<M>>(
72        &self,
73        args: P::AcquireArgs,
74    ) -> impl PinInit<KMutexGuard<'_, Class, M, P>, core::convert::Infallible> {
75        KMutexGuard::new_with_args(self, args)
76    }
77
78    /// Acquires this lock and aliases it with `alias`, returning a guard that proves ownership
79    /// of both `Class` and `AliasClass`.
80    #[inline]
81    pub fn aliased_lock<'a, AliasClass: LockClass, M2: RawLock>(
82        &'a self,
83        alias: &'a KMutex<AliasClass, M2>,
84    ) -> impl PinInit<KMutexAliasedGuard<'a, Class, AliasClass, M, M2>, core::convert::Infallible>
85    where
86        <M as RawLock>::DefaultPolicy: ExclusiveLockPolicy<M, AcquireArgs = ()>,
87    {
88        KMutexAliasedGuard::new(self, alias)
89    }
90
91    /// Acquires this lock with arguments and aliases it with `alias`, returning a guard that proves
92    /// ownership of both `Class` and `AliasClass`.
93    #[inline]
94    pub fn aliased_lock_with<'a, AliasClass: LockClass, M2: RawLock>(
95        &'a self,
96        alias: &'a KMutex<AliasClass, M2>,
97        args: <<M as RawLock>::DefaultPolicy as LockPolicy<M>>::AcquireArgs,
98    ) -> impl PinInit<KMutexAliasedGuard<'a, Class, AliasClass, M, M2>, core::convert::Infallible>
99    {
100        KMutexAliasedGuard::new_with_args(self, alias, args)
101    }
102
103    /// Acquires this lock with policy `P` and aliases it with `alias`.
104    #[inline]
105    pub fn aliased_lock_policy<
106        'a,
107        AliasClass: LockClass,
108        M2: RawLock,
109        P: ExclusiveLockPolicy<M, AcquireArgs = ()>,
110    >(
111        &'a self,
112        alias: &'a KMutex<AliasClass, M2>,
113    ) -> impl PinInit<KMutexAliasedGuard<'a, Class, AliasClass, M, M2, P>, core::convert::Infallible>
114    {
115        KMutexAliasedGuard::new(self, alias)
116    }
117
118    /// Acquires this lock with policy `P` and arguments and aliases it with `alias`.
119    #[inline]
120    pub fn aliased_lock_policy_with<
121        'a,
122        AliasClass: LockClass,
123        M2: RawLock,
124        P: ExclusiveLockPolicy<M>,
125    >(
126        &'a self,
127        alias: &'a KMutex<AliasClass, M2>,
128        args: P::AcquireArgs,
129    ) -> impl PinInit<KMutexAliasedGuard<'a, Class, AliasClass, M, M2, P>, core::convert::Infallible>
130    {
131        KMutexAliasedGuard::new_with_args(self, alias, args)
132    }
133
134    /// Acquires multiple locks of the same class and type in ascending memory address order,
135    /// using the default policy.
136    #[inline]
137    pub fn lock_multiple<'a, const N: usize>(
138        locks: [&'a Self; N],
139    ) -> impl PinInit<KMutexGuardMultiple<'a, N, Class, M>, core::convert::Infallible>
140    where
141        <M as RawLock>::DefaultPolicy: ExclusiveLockPolicy<M, AcquireArgs = ()>,
142    {
143        KMutexGuardMultiple::new(locks)
144    }
145
146    /// Acquires multiple locks of the same class and type in ascending memory address order with
147    /// arguments, using the default policy.
148    #[inline]
149    pub fn lock_multiple_with<'a, const N: usize>(
150        locks: [&'a Self; N],
151        args: <<M as RawLock>::DefaultPolicy as LockPolicy<M>>::AcquireArgs,
152    ) -> impl PinInit<KMutexGuardMultiple<'a, N, Class, M>, core::convert::Infallible> {
153        KMutexGuardMultiple::new_with_args(locks, args)
154    }
155
156    /// Acquires multiple locks of the same class and type in ascending memory address order,
157    /// using policy `P`.
158    #[inline]
159    pub fn lock_multiple_policy<'a, const N: usize, P: ExclusiveLockPolicy<M, AcquireArgs = ()>>(
160        locks: [&'a Self; N],
161    ) -> impl PinInit<KMutexGuardMultiple<'a, N, Class, M, P>, core::convert::Infallible> {
162        KMutexGuardMultiple::new(locks)
163    }
164
165    /// Acquires multiple locks of the same class and type in ascending memory address order with
166    /// policy `P` and arguments.
167    #[inline]
168    pub fn lock_multiple_policy_with<'a, const N: usize, P: ExclusiveLockPolicy<M>>(
169        locks: [&'a Self; N],
170        args: P::AcquireArgs,
171    ) -> impl PinInit<KMutexGuardMultiple<'a, N, Class, M, P>, core::convert::Infallible> {
172        KMutexGuardMultiple::new_with_args(locks, args)
173    }
174
175    /// Returns a reference to the underlying raw lock.
176    #[inline]
177    pub fn raw_mutex(&self) -> &M {
178        &self.mutex
179    }
180
181    const fn class_id() -> *const core::ffi::c_void {
182        if cfg!(any(feature = "lock_dep", feature = "lock_metadata_only")) {
183            Class::ID
184        } else {
185            core::ptr::null()
186        }
187    }
188}
189
190impl<Class: LockClass, M> core::fmt::Debug for KMutex<Class, M> {
191    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
192        f.debug_struct("KMutex").field("class", &core::any::type_name::<Class>()).finish()
193    }
194}
195
196/// A validation guard representing exclusive lock ownership and active list participation.
197///
198/// The guard is pinned in memory to ensure that its `state` remains safe and valid
199/// inside the C++ loop detector active thread list.
200#[repr(C)]
201pub struct KMutexGuard<
202    'a,
203    Class: LockClass,
204    M: RawLock = RawMutex,
205    P: ExclusiveLockPolicy<M> = <M as RawLock>::DefaultPolicy,
206> {
207    mutex: Option<&'a KMutex<Class, M>>,
208
209    state: P::GuardState,
210
211    token: LockToken<'a, Class>,
212}
213
214impl<'a, Class: LockClass, M: RawLock, P: ExclusiveLockPolicy<M>> KMutexGuard<'a, Class, M, P> {
215    /// Creates a new stack-pinned validation guard initialization block.
216    pub fn new(mutex: &'a KMutex<Class, M>) -> impl PinInit<Self, core::convert::Infallible>
217    where
218        P: ExclusiveLockPolicy<M, AcquireArgs = ()>,
219    {
220        Self::new_with_args(mutex, ())
221    }
222
223    /// Creates a new stack-pinned validation guard initialization block with policy acquire
224    /// arguments.
225    pub fn new_with_args(
226        mutex: &'a KMutex<Class, M>,
227        args: P::AcquireArgs,
228    ) -> impl PinInit<Self, core::convert::Infallible> {
229        // SAFETY: The closure correctly initializes all fields of the allocated `KMutexGuard`
230        // and satisfies all safety requirements of `pin_init_from_closure`.
231        unsafe {
232            pin_init_from_closure(move |this: *mut Self| -> Result<(), core::convert::Infallible> {
233                // SAFETY: `this` is a valid pointer to uninitialized memory allocated for
234                // `KMutexGuard`.
235
236                let mutex_addr = addr_of_mut!((*this).mutex);
237                core::ptr::write(mutex_addr, Some(mutex));
238
239                let state_addr = addr_of_mut!((*this).state);
240                core::ptr::write(state_addr, P::GuardState::default());
241
242                P::acquire(&mutex.mutex, &mut *state_addr, args);
243
244                let token_addr = addr_of_mut!((*this).token);
245                core::ptr::write(token_addr, LockToken::new());
246
247                Ok(())
248            })
249        }
250    }
251
252    /// Returns a shared reference to the lock proof `LockToken`.
253    #[inline]
254    pub fn token(&self) -> &LockToken<'a, Class> {
255        assert!(self.mutex.is_some(), "cannot access token of a released KMutexGuard");
256        &self.token
257    }
258
259    /// Returns a mutable reference to the lock proof `LockToken` inside this pinned projection.
260    #[inline]
261    pub fn token_mut(self: Pin<&mut Self>) -> &mut LockToken<'a, Class> {
262        // SAFETY: Modifying the non-pinned raw `token` field does not violate pinning invariants
263        // since the token has no drop logic or pointer-location sensitivity.
264        let me = unsafe { self.get_unchecked_mut() };
265        assert!(me.mutex.is_some(), "cannot access token of a released KMutexGuard");
266        &mut me.token
267    }
268
269    /// # Safety
270    ///
271    /// The guard has not been released.
272    #[inline]
273    pub unsafe fn token_unchecked(&self) -> &LockToken<'a, Class> {
274        debug_assert!(self.mutex.is_some(), "cannot access token of a released KMutexGuard");
275        &self.token
276    }
277
278    /// # Safety
279    ///
280    /// The guard has not been released.
281    #[inline]
282    pub unsafe fn token_mut_unchecked(self: Pin<&mut Self>) -> &mut LockToken<'a, Class> {
283        // SAFETY: Modifying the non-pinned raw `token` field does not violate pinning invariants
284        // since the token has no drop logic or pointer-location sensitivity.
285        let me = unsafe { self.get_unchecked_mut() };
286        debug_assert!(me.mutex.is_some(), "cannot access token of a released KMutexGuard");
287        &mut me.token
288    }
289
290    /// Returns true if this guard currently holds the given `target` mutex instance.
291    #[inline]
292    pub fn is_locked_by(&self, target: &KMutex<Class, M>) -> bool {
293        core::mem::size_of::<M>() > 0 && self.mutex.is_some_and(|m| core::ptr::eq(m, target))
294    }
295
296    /// Explicitly releases the lock before the guard is dropped.
297    ///
298    /// If the lock has already been released, this is a no-op.
299    #[inline]
300    pub fn release(self: Pin<&mut Self>) {
301        // SAFETY: `state` is pinned on the stack and valid.
302        unsafe {
303            let me = self.get_unchecked_mut();
304            if let Some(mutex) = me.mutex.take() {
305                P::release(&mutex.mutex, &mut me.state);
306            }
307        }
308    }
309
310    /// Temporarily releases the lock before executing the given callable `f` and then
311    /// re-acquires the lock.
312    #[inline]
313    pub fn call_unlocked<R, F: FnOnce() -> R>(self: Pin<&mut Self>, f: F) -> R {
314        // SAFETY: `state` is pinned on the stack and valid.
315        unsafe {
316            let me = self.get_unchecked_mut();
317            let mutex = me.mutex.expect("cannot call_unlocked on a released KMutexGuard");
318            P::release(&mutex.mutex, &mut me.state);
319            let result = f();
320            P::reacquire(&mutex.mutex, &mut me.state);
321            result
322        }
323    }
324
325    /// Calls a closure while temporarily disabling lockdep tracking for the lock held by this
326    /// guard.
327    #[inline]
328    pub fn call_untracked<R, F: FnOnce(&mut LockToken<'a, Class>) -> R>(
329        self: Pin<&mut Self>,
330        f: F,
331    ) -> R {
332        // SAFETY: `state` is pinned on the stack and valid.
333        unsafe {
334            let me = self.get_unchecked_mut();
335            let mutex = me.mutex.expect("cannot call_untracked on a released KMutexGuard");
336            P::untrack(&mutex.mutex, &mut me.state);
337            let result = f(&mut me.token);
338            P::retrack(&mutex.mutex, &mut me.state);
339            result
340        }
341    }
342}
343
344// Not `#[pin_data(PinnedDrop)]`, whose generated `Drop` is not `#[inline]`.
345impl<'a, Class: LockClass, M: RawLock, P: ExclusiveLockPolicy<M>> Drop
346    for KMutexGuard<'a, Class, M, P>
347{
348    #[inline]
349    fn drop(&mut self) {
350        // SAFETY: A guard is only created pinned, by `new_with_args`, and `self` is not moved
351        // after it is dropped.
352        unsafe { Pin::new_unchecked(self) }.release();
353    }
354}
355
356/// Guards simultaneous acquisitions of multiple locks of the same class and type.
357///
358/// Locks are automatically sorted in-place by ascending memory address before acquisition to
359/// preserve the intra-class ordering invariant and prevent AB/BA deadlocks, matching C++
360/// `lockdep::GuardMultiple`.
361#[repr(C)]
362#[pin_data(PinnedDrop)]
363pub struct KMutexGuardMultiple<
364    'a,
365    const N: usize,
366    Class: LockClass,
367    M: RawLock = RawMutex,
368    P: ExclusiveLockPolicy<M> = <M as RawLock>::DefaultPolicy,
369> {
370    mutexes: [Option<&'a KMutex<Class, M>>; N],
371
372    #[pin]
373    states: [P::GuardState; N],
374
375    tokens: [LockToken<'a, Class>; N],
376}
377
378impl<'a, const N: usize, Class: LockClass, M: RawLock, P: ExclusiveLockPolicy<M>>
379    KMutexGuardMultiple<'a, N, Class, M, P>
380{
381    /// Sorts an array of lock references in-place by ascending memory address using insertion sort,
382    /// matching C++ `GuardMultiple::InsertionSortPointers`.
383    #[inline]
384    fn insertion_sort_locks(locks: &mut [&'a KMutex<Class, M>; N]) {
385        let mut i = 1;
386        while i < N {
387            let x = locks[i];
388            let x_addr = x.mutex.as_mut_ptr() as usize;
389            let mut j = i;
390            while j > 0 && (locks[j - 1].mutex.as_mut_ptr() as usize) > x_addr {
391                locks[j] = locks[j - 1];
392                j -= 1;
393            }
394            locks[j] = x;
395            i += 1;
396        }
397    }
398
399    /// Creates a new stack-pinned multi-lock validation guard initialization block.
400    pub fn new(locks: [&'a KMutex<Class, M>; N]) -> impl PinInit<Self, core::convert::Infallible>
401    where
402        P: ExclusiveLockPolicy<M, AcquireArgs = ()>,
403    {
404        Self::new_with_args(locks, ())
405    }
406
407    /// Creates a new stack-pinned multi-lock validation guard initialization block with policy
408    /// acquire arguments.
409    pub fn new_with_args(
410        mut locks: [&'a KMutex<Class, M>; N],
411        args: P::AcquireArgs,
412    ) -> impl PinInit<Self, core::convert::Infallible> {
413        Self::insertion_sort_locks(&mut locks);
414        assert!(
415            N <= 1 || core::mem::size_of::<M>() > 0,
416            "KMutexGuardMultiple requires physical locks with non-zero size"
417        );
418        for i in 1..N {
419            assert_ne!(
420                locks[i - 1].mutex.as_mut_ptr(),
421                locks[i].mutex.as_mut_ptr(),
422                "KMutexGuardMultiple requires all locks to be distinct instances"
423            );
424        }
425        // SAFETY: The closure initializes all fields of `KMutexGuardMultiple` in-place.
426        unsafe {
427            pin_init_from_closure(move |this: *mut Self| -> Result<(), core::convert::Infallible> {
428                let mutexes_ptr =
429                    addr_of_mut!((*this).mutexes) as *mut Option<&'a KMutex<Class, M>>;
430                let states_ptr = addr_of_mut!((*this).states) as *mut P::GuardState;
431
432                for (i, lock) in locks.into_iter().enumerate() {
433                    core::ptr::write(mutexes_ptr.add(i), Some(lock));
434                    let state_addr = states_ptr.add(i);
435                    core::ptr::write(state_addr, P::GuardState::default());
436                    P::acquire(&lock.mutex, &mut *state_addr, args);
437                }
438
439                core::ptr::write(
440                    addr_of_mut!((*this).tokens),
441                    core::array::from_fn(|_| LockToken::new()),
442                );
443
444                Ok(())
445            })
446        }
447    }
448
449    /// Returns shared references to the lock proof tokens.
450    #[inline]
451    pub fn tokens(&self) -> &[LockToken<'a, Class>; N] {
452        assert!(
453            N == 0 || self.mutexes[0].is_some(),
454            "cannot access tokens of a released KMutexGuardMultiple"
455        );
456        &self.tokens
457    }
458
459    /// Returns mutable references to the lock proof tokens.
460    #[inline]
461    pub fn tokens_mut(self: Pin<&mut Self>) -> &mut [LockToken<'a, Class>; N] {
462        // SAFETY: Modifying the non-pinned `tokens` field does not violate pinning invariants.
463        let me = unsafe { self.get_unchecked_mut() };
464        assert!(
465            N == 0 || me.mutexes[0].is_some(),
466            "cannot access tokens of a released KMutexGuardMultiple"
467        );
468        &mut me.tokens
469    }
470
471    /// Explicitly releases all held locks in reverse order before the guard is dropped.
472    ///
473    /// If the locks have already been released, this is a no-op.
474    #[inline]
475    pub fn release(self: Pin<&mut Self>) {
476        // SAFETY: `states` are pinned on the stack and valid. Locks are released in reverse
477        // acquisition order (matching C++ `GuardMultiple` destruction order).
478        unsafe {
479            let me = self.get_unchecked_mut();
480            for i in (0..N).rev() {
481                if let Some(mutex) = me.mutexes[i].take() {
482                    P::release(&mutex.mutex, &mut me.states[i]);
483                }
484            }
485        }
486    }
487}
488
489#[pinned_drop]
490impl<'a, const N: usize, Class: LockClass, M: RawLock, P: ExclusiveLockPolicy<M>> PinnedDrop
491    for KMutexGuardMultiple<'a, N, Class, M, P>
492{
493    fn drop(self: Pin<&mut Self>) {
494        self.release();
495    }
496}
497
498/// Acquires multiple `KMutex` instances of the same class and type in ascending memory address
499/// order, using the default policy.
500#[inline]
501pub fn guard_multiple<'a, const N: usize, Class: LockClass, M: RawLock>(
502    locks: [&'a KMutex<Class, M>; N],
503) -> impl PinInit<KMutexGuardMultiple<'a, N, Class, M>, core::convert::Infallible>
504where
505    <M as RawLock>::DefaultPolicy: ExclusiveLockPolicy<M, AcquireArgs = ()>,
506{
507    KMutexGuardMultiple::new(locks)
508}
509
510/// Acquires multiple `KMutex` instances of the same class and type in ascending memory address
511/// order with arguments, using the default policy.
512#[inline]
513pub fn guard_multiple_with<'a, const N: usize, Class: LockClass, M: RawLock>(
514    locks: [&'a KMutex<Class, M>; N],
515    args: <<M as RawLock>::DefaultPolicy as LockPolicy<M>>::AcquireArgs,
516) -> impl PinInit<KMutexGuardMultiple<'a, N, Class, M>, core::convert::Infallible> {
517    KMutexGuardMultiple::new_with_args(locks, args)
518}
519
520/// Acquires multiple `KMutex` instances of the same class and type in ascending memory address
521/// order, using policy `P`.
522#[inline]
523pub fn guard_multiple_policy<
524    'a,
525    const N: usize,
526    Class: LockClass,
527    M: RawLock,
528    P: ExclusiveLockPolicy<M, AcquireArgs = ()>,
529>(
530    locks: [&'a KMutex<Class, M>; N],
531) -> impl PinInit<KMutexGuardMultiple<'a, N, Class, M, P>, core::convert::Infallible> {
532    KMutexGuardMultiple::new(locks)
533}
534
535/// Acquires multiple `KMutex` instances of the same class and type in ascending memory address
536/// order with policy `P` and arguments.
537#[inline]
538pub fn guard_multiple_policy_with<
539    'a,
540    const N: usize,
541    Class: LockClass,
542    M: RawLock,
543    P: ExclusiveLockPolicy<M>,
544>(
545    locks: [&'a KMutex<Class, M>; N],
546    args: P::AcquireArgs,
547) -> impl PinInit<KMutexGuardMultiple<'a, N, Class, M, P>, core::convert::Infallible> {
548    KMutexGuardMultiple::new_with_args(locks, args)
549}
550
551/// Type tag to indicate aliased lock acquisition.
552pub struct AliasedLock;
553
554/// Acquires an aliased lock on two `KMutex` instances that reference the same underlying lock.
555///
556/// Only `lock1` is physically acquired, but the resulting guard holds proof tokens for both
557/// `Class1` and `Class2`.
558#[inline]
559pub fn aliased_lock<'a, Class1: LockClass, Class2: LockClass, M1: RawLock, M2: RawLock>(
560    lock1: &'a KMutex<Class1, M1>,
561    lock2: &'a KMutex<Class2, M2>,
562) -> impl PinInit<KMutexAliasedGuard<'a, Class1, Class2, M1, M2>, core::convert::Infallible>
563where
564    <M1 as RawLock>::DefaultPolicy: ExclusiveLockPolicy<M1, AcquireArgs = ()>,
565{
566    KMutexAliasedGuard::new(lock1, lock2)
567}
568
569/// Acquires an aliased lock with arguments on two `KMutex` instances that reference the same
570/// underlying lock.
571#[inline]
572pub fn aliased_lock_with<'a, Class1: LockClass, Class2: LockClass, M1: RawLock, M2: RawLock>(
573    lock1: &'a KMutex<Class1, M1>,
574    lock2: &'a KMutex<Class2, M2>,
575    args: <<M1 as RawLock>::DefaultPolicy as LockPolicy<M1>>::AcquireArgs,
576) -> impl PinInit<KMutexAliasedGuard<'a, Class1, Class2, M1, M2>, core::convert::Infallible> {
577    KMutexAliasedGuard::new_with_args(lock1, lock2, args)
578}
579
580/// Acquires an aliased lock with a specific policy on two `KMutex` instances that reference the
581/// same underlying lock.
582#[inline]
583pub fn aliased_lock_policy<
584    'a,
585    Class1: LockClass,
586    Class2: LockClass,
587    M1: RawLock,
588    M2: RawLock,
589    P: ExclusiveLockPolicy<M1, AcquireArgs = ()>,
590>(
591    lock1: &'a KMutex<Class1, M1>,
592    lock2: &'a KMutex<Class2, M2>,
593) -> impl PinInit<KMutexAliasedGuard<'a, Class1, Class2, M1, M2, P>, core::convert::Infallible> {
594    KMutexAliasedGuard::new(lock1, lock2)
595}
596
597/// Acquires an aliased lock with a specific policy and arguments on two `KMutex` instances that
598/// reference the same underlying lock.
599#[inline]
600pub fn aliased_lock_policy_with<
601    'a,
602    Class1: LockClass,
603    Class2: LockClass,
604    M1: RawLock,
605    M2: RawLock,
606    P: ExclusiveLockPolicy<M1>,
607>(
608    lock1: &'a KMutex<Class1, M1>,
609    lock2: &'a KMutex<Class2, M2>,
610    args: P::AcquireArgs,
611) -> impl PinInit<KMutexAliasedGuard<'a, Class1, Class2, M1, M2, P>, core::convert::Infallible> {
612    KMutexAliasedGuard::new_with_args(lock1, lock2, args)
613}
614
615/// A validation guard representing ownership of two aliased locks simultaneously.
616///
617/// Only the first lock is physically acquired, but proof tokens for both lock classes (`Class1` and
618/// `Class2`) are provided.
619#[pin_data]
620pub struct KMutexAliasedGuard<
621    'a,
622    Class1: LockClass,
623    Class2: LockClass,
624    M1: RawLock = RawMutex,
625    M2: RawLock = M1,
626    P: ExclusiveLockPolicy<M1> = <M1 as RawLock>::DefaultPolicy,
627> {
628    #[pin]
629    inner: KMutexGuard<'a, Class1, M1, P>,
630
631    token2: LockToken<'a, Class2>,
632
633    _phantom: PhantomData<&'a KMutex<Class2, M2>>,
634}
635
636impl<'a, Class1: LockClass, Class2: LockClass, M1: RawLock, M2: RawLock, P: ExclusiveLockPolicy<M1>>
637    KMutexAliasedGuard<'a, Class1, Class2, M1, M2, P>
638{
639    /// Creates a new stack-pinned aliased validation guard initialization block.
640    pub fn new(
641        lock1: &'a KMutex<Class1, M1>,
642        lock2: &'a KMutex<Class2, M2>,
643    ) -> impl PinInit<Self, core::convert::Infallible>
644    where
645        P: ExclusiveLockPolicy<M1, AcquireArgs = ()>,
646    {
647        Self::new_with_args(lock1, lock2, ())
648    }
649
650    /// Creates a new stack-pinned aliased validation guard initialization block with policy acquire
651    /// arguments.
652    pub fn new_with_args(
653        lock1: &'a KMutex<Class1, M1>,
654        _lock2: &'a KMutex<Class2, M2>,
655        args: P::AcquireArgs,
656    ) -> impl PinInit<Self, core::convert::Infallible> {
657        assert!(
658            core::mem::size_of::<M1>() > 0 && core::mem::size_of::<M2>() > 0,
659            "AliasedLock requires physical locks with non-zero size"
660        );
661        assert_eq!(
662            lock1.mutex.as_mut_ptr(),
663            _lock2.mutex.as_mut_ptr(),
664            "AliasedLock requires lock1 and lock2 to point to the same physical lock"
665        );
666        pin_init!(Self {
667            inner <- KMutexGuard::new_with_args(lock1, args),
668            // SAFETY: `inner` holds the underlying mutex, which is aliased to represent `Class2`.
669            token2: unsafe { LockToken::new() },
670            _phantom: PhantomData,
671        })
672    }
673
674    /// Returns shared references to both lock proof tokens `(Class1, Class2)`.
675    #[inline]
676    pub fn tokens(&self) -> (&LockToken<'a, Class1>, &LockToken<'a, Class2>) {
677        (self.inner.token(), &self.token2)
678    }
679
680    /// Returns simultaneous mutable references to both proof tokens `(Class1, Class2)`.
681    ///
682    /// # Safety
683    ///
684    /// Both returned tokens are backed by a single physical lock acquisition. The caller must
685    /// ensure that the two mutable tokens are not used to obtain aliasing mutable references to
686    /// the same `KCell` or `KOnceCell` (for example, by passing both tokens to `guard_<lock>_mut`
687    /// on the same struct instance).
688    #[inline]
689    pub unsafe fn tokens_mut(
690        self: Pin<&mut Self>,
691    ) -> (&mut LockToken<'a, Class1>, &mut LockToken<'a, Class2>) {
692        // SAFETY: `get_unchecked_mut` and `Pin::new_unchecked` project the pinned `inner` field.
693        let (inner_pin, token2) = unsafe {
694            let me = self.get_unchecked_mut();
695            (Pin::new_unchecked(&mut me.inner), &mut me.token2)
696        };
697        (inner_pin.token_mut(), token2)
698    }
699
700    /// Returns a shared reference to the primary lock proof token (`Class1`).
701    #[inline]
702    pub fn token(&self) -> &LockToken<'a, Class1> {
703        self.inner.token()
704    }
705
706    /// Returns a mutable reference to the primary lock proof token (`Class1`).
707    #[inline]
708    pub fn token_mut(self: Pin<&mut Self>) -> &mut LockToken<'a, Class1> {
709        self.inner_guard().token_mut()
710    }
711
712    /// Returns a pinned mutable reference to the primary inner `KMutexGuard`.
713    #[inline]
714    pub fn inner_guard(self: Pin<&mut Self>) -> Pin<&mut KMutexGuard<'a, Class1, M1, P>> {
715        // SAFETY: `inner` is structurally pinned inside `Self`.
716        unsafe { Pin::new_unchecked(&mut self.get_unchecked_mut().inner) }
717    }
718
719    /// Explicitly releases the underlying lock before the guard is dropped.
720    ///
721    /// If the lock has already been released, this is a no-op.
722    #[inline]
723    pub fn release(self: Pin<&mut Self>) {
724        self.inner_guard().release();
725    }
726
727    /// Temporarily releases the lock before executing the given callable `f` and then
728    /// re-acquires the lock.
729    #[inline]
730    pub fn call_unlocked<R, F: FnOnce() -> R>(self: Pin<&mut Self>, f: F) -> R {
731        self.inner_guard().call_unlocked(f)
732    }
733
734    /// Calls a closure while temporarily disabling lockdep tracking for the lock held by this
735    /// guard.
736    #[inline]
737    pub fn call_untracked<R, F: FnOnce(&mut LockToken<'a, Class1>) -> R>(
738        self: Pin<&mut Self>,
739        f: F,
740    ) -> R {
741        self.inner_guard().call_untracked(f)
742    }
743}
744
745#[cfg(not(feature = "kernel"))]
746#[cfg(test)]
747mod tests {
748    use super::*;
749    use crate::{KCell, guarded};
750    use lockdep::LockClass;
751    use pin_init::{pin_init, stack_pin_init};
752
753    struct MyClass;
754    impl LockClass for MyClass {
755        const ID: *mut core::ffi::c_void = core::ptr::null_mut();
756    }
757
758    #[pin_init::pin_data]
759    struct MyStruct {
760        #[pin]
761        mu: KMutex<MyClass>,
762        data1: KCell<u32, MyClass>,
763        data2: KCell<i32, MyClass>,
764    }
765
766    #[test]
767    fn test_basic_token_access() {
768        stack_pin_init!(let s = pin_init!(MyStruct {
769            mu <- KMutex::init(),
770            data1: KCell::new(10),
771            data2: KCell::new(-5),
772        }));
773
774        lock!(let mut guard = s.mu.lock());
775
776        // SAFETY: `guard` holds `s.mu`, which guards `s.data1` and `s.data2`.
777        unsafe {
778            assert_eq!(*s.data1.get(guard.token()), 10);
779            assert_eq!(*s.data2.get(guard.token()), -5);
780        }
781        // SAFETY: `guard` holds `s.mu`, which guards `s.data1`.
782        unsafe {
783            let token_mut = guard.as_mut().token_mut();
784            *s.data1.get_mut(token_mut) = 20;
785            assert_eq!(*s.data1.get(guard.token()), 20);
786        }
787    }
788
789    #[guarded]
790    struct MyGuardedStruct {
791        #[mutex]
792        mu: KMutex,
793        #[guarded_by(mu)]
794        data1: u32,
795        #[guarded_by(mu)]
796        data2: i32,
797    }
798
799    #[test]
800    fn test_macro_guarded() {
801        stack_pin_init!(let s = pin_init!(MyGuardedStruct {
802            mu <- KMutex::init(),
803            data1: 100.into(),
804            data2: (-50).into(),
805        }));
806
807        {
808            lock!(let mut guard = s.lock_mu());
809
810            // Safe individual field access
811            assert_eq!(*guard.data1(), 100);
812            assert_eq!(*guard.data2(), -50);
813
814            *guard.as_mut().data1_mut() = 200;
815            assert_eq!(*guard.data1(), 200);
816
817            // Safe disjoint/split access
818            let fields = guard.as_mut().fields_mut();
819            *fields.data1 += 50;
820            *fields.data2 += 50;
821        }
822
823        // Verify fields
824        lock!(let guard = s.lock_mu());
825        assert_eq!(*guard.data1(), 250);
826        assert_eq!(*guard.data2(), 0);
827    }
828
829    #[test]
830    fn test_kmutex_init() {
831        stack_pin_init!(let mu = KMutex::<MyClass>::init());
832        lock!(mu.lock());
833    }
834
835    #[test]
836    fn test_kmutex_debug() {
837        extern crate std;
838        stack_pin_init!(let mu = KMutex::<MyClass>::init());
839        let debug_str = std::format!("{:?}", mu);
840        assert!(debug_str.contains("KMutex"));
841    }
842
843    #[guarded]
844    struct MyMultiGuardedStruct {
845        #[mutex]
846        mu1: KMutex,
847        #[mutex]
848        mu2: KMutex,
849        #[guarded_by(mu1)]
850        data1: u32,
851        #[guarded_by(mu2)]
852        data2: i32,
853    }
854
855    #[test]
856    fn test_macro_multi_guarded() {
857        stack_pin_init!(let s = pin_init!(MyMultiGuardedStruct {
858            mu1 <- KMutex::init(),
859            mu2 <- KMutex::init(),
860            data1: 10.into(),
861            data2: 20.into(),
862        }));
863
864        lock!(let mut guard1 = s.lock_mu1());
865        lock!(let mut guard2 = s.lock_mu2());
866
867        assert_eq!(*guard1.data1(), 10);
868        assert_eq!(*guard2.data2(), 20);
869        *guard1.as_mut().data1_mut() = 15;
870        *guard2.as_mut().data2_mut() = 25;
871        assert_eq!(*guard1.data1(), 15);
872        assert_eq!(*guard2.data2(), 25);
873    }
874
875    #[guarded]
876    struct MyDefaultGuardedStruct {
877        #[mutex]
878        mu: KMutex,
879        #[guarded_by(mu)]
880        data: u32,
881    }
882
883    #[test]
884    fn test_derive_default_guarded() {
885        stack_pin_init!(let s = pin_init!(MyDefaultGuardedStruct {
886            mu <- KMutex::init(),
887            data: 0.into(),
888        }));
889        lock!(let guard = s.lock_mu());
890        assert_eq!(*guard.data(), 0);
891    }
892
893    #[guarded]
894    struct MyGenericGuardedStruct<T> {
895        #[mutex]
896        mu: KMutex,
897        #[guarded_by(mu)]
898        data: T,
899    }
900
901    #[test]
902    fn test_macro_generic_guarded() {
903        stack_pin_init!(let s = pin_init!(MyGenericGuardedStruct::<u32> {
904            mu <- KMutex::init(),
905            data: 0.into(),
906        }));
907        lock!(let mut guard = s.lock_mu());
908        assert_eq!(*guard.data(), 0);
909
910        *guard.as_mut().data_mut() = 42;
911        assert_eq!(*guard.data(), 42);
912
913        let fields = guard.as_mut().fields_mut();
914        *fields.data = 100;
915
916        let fields_shared = guard.fields();
917        assert_eq!(*fields_shared.data, 100);
918    }
919
920    #[guarded]
921    struct MyExplicitParentGuardedStruct {
922        #[mutex]
923        mu: KMutex,
924        #[guarded_by(mu)]
925        data: u32,
926        pub label: &'static str,
927    }
928
929    impl MyExplicitParentGuardedStruct {
930        pub fn has_label(&self) -> bool {
931            !self.label.is_empty()
932        }
933    }
934
935    impl<'a> MyExplicitParentGuardedStructMuGuard<'a> {
936        pub fn process_with_context(mut self: Pin<&mut Self>) {
937            let has_label = self.parent.has_label();
938            let label = self.parent.label;
939            if has_label && label == "apply_update" {
940                let fields = self.as_mut().fields_mut();
941                *fields.data = 100;
942            }
943        }
944    }
945
946    #[test]
947    fn test_macro_guard_explicit_parent_access() {
948        stack_pin_init!(let s = pin_init!(MyExplicitParentGuardedStruct {
949            mu <- KMutex::init(),
950            data: 0.into(),
951            label: "apply_update",
952        }));
953
954        {
955            lock!(let mut guard = s.lock_mu());
956            guard.as_mut().process_with_context();
957        }
958
959        lock!(let guard = s.lock_mu());
960        assert_eq!(*guard.data(), 100);
961    }
962
963    #[test]
964    fn test_call_unlocked() {
965        stack_pin_init!(let s = pin_init!(MyGuardedStruct {
966            mu <- KMutex::init(),
967            data1: 10.into(),
968            data2: 20.into(),
969        }));
970
971        lock!(let mut guard = s.lock_mu());
972        assert_eq!(*guard.data1(), 10);
973
974        let unlocked_result = guard.as_mut().call_unlocked(|| {
975            // During call_unlocked, the lock is temporarily released.
976            42
977        });
978        assert_eq!(unlocked_result, 42);
979
980        // After call_unlocked returns, the lock is held again and fields can be accessed/modified.
981        *guard.as_mut().data1_mut() = 99;
982        assert_eq!(*guard.data1(), 99);
983    }
984
985    #[guarded]
986    struct AliasedTargetStruct {
987        #[mutex(MyGuardedStructMuClass)]
988        mu: KMutex<crate::PhantomMutex>,
989        parent: *const MyGuardedStruct,
990        #[guarded_by(mu)]
991        value: u32,
992    }
993
994    impl AliasedTargetStruct {
995        fn lock(&self) -> &KMutex<MyGuardedStructMuClass> {
996            // SAFETY: `self.parent` points to a live `MyGuardedStruct` in the test scope.
997            unsafe { &(*self.parent).mu }
998        }
999    }
1000
1001    #[test]
1002    fn test_aliased_lock_basic() {
1003        stack_pin_init!(let s = pin_init!(MyGuardedStruct {
1004            mu <- KMutex::init(),
1005            data1: 100.into(),
1006            data2: 200.into(),
1007        }));
1008
1009        let s_ptr: *const MyGuardedStruct = &*s;
1010        stack_pin_init!(let target = pin_init!(AliasedTargetStruct {
1011            mu: KMutex::new(crate::PhantomMutex),
1012            parent: s_ptr,
1013            value: 300.into(),
1014        }));
1015
1016        {
1017            lock!(let mut guard = aliased_lock(&s.mu, target.lock()));
1018
1019            // Access fields of s and target using tokens() pair
1020            let (t1, t2) = guard.tokens();
1021            // SAFETY: `t1` and `t2` are backed by `s.mu`, which guards `s` and `target`.
1022            unsafe {
1023                assert_eq!(*s.guard_mu(t1).data1(), 100);
1024                assert_eq!(*s.guard_mu(t1).data2(), 200);
1025                assert_eq!(*target.guard_mu(t2).value(), 300);
1026            }
1027
1028            // Disjoint simultaneous mutable access to both structs
1029            // SAFETY: `s` and `target` are distinct structs with disjoint `KCell` fields.
1030            unsafe {
1031                let (t1_mut, t2_mut) = guard.as_mut().tokens_mut();
1032                *s.guard_mu_mut(t1_mut).data1_mut() = 101;
1033                *target.guard_mu_mut(t2_mut).value_mut() = 301;
1034            }
1035
1036            // Call unlocked on aliased guard
1037            let res = guard.as_mut().call_unlocked(|| 1234);
1038            assert_eq!(res, 1234);
1039
1040            // Verify mutations persist
1041            let (t1, t2) = guard.tokens();
1042            // SAFETY: `t1` and `t2` are backed by `s.mu`, which guards `s` and `target`.
1043            unsafe {
1044                assert_eq!(*s.guard_mu(t1).data1(), 101);
1045                assert_eq!(*target.guard_mu(t2).value(), 301);
1046            }
1047        }
1048
1049        // Verify state with regular lock
1050        lock!(let guard = s.lock_mu());
1051        assert_eq!(*guard.data1(), 101);
1052    }
1053
1054    #[test]
1055    fn test_aliased_lock_macro_method() {
1056        stack_pin_init!(let s = pin_init!(MyGuardedStruct {
1057            mu <- KMutex::init(),
1058            data1: 10.into(),
1059            data2: 20.into(),
1060        }));
1061
1062        let s_ptr: *const MyGuardedStruct = &*s;
1063        stack_pin_init!(let target = pin_init!(AliasedTargetStruct {
1064            mu: KMutex::new(crate::PhantomMutex),
1065            parent: s_ptr,
1066            value: 30.into(),
1067        }));
1068
1069        {
1070            lock!(let mut guard = s.lock_mu_aliased(target.lock()));
1071            let (_t1, t2) = guard.tokens();
1072            // SAFETY: `target` is guarded by `s.mu` and `_t1_mut` is unused.
1073            unsafe {
1074                assert_eq!(*target.guard_mu(t2).value(), 30);
1075                let (_t1_mut, t2_mut) = guard.as_mut().tokens_mut();
1076                *target.guard_mu_mut(t2_mut).value_mut() = 35;
1077            }
1078        }
1079
1080        lock!(let guard = s.lock_mu());
1081        // SAFETY: `guard` holds `s.mu`, which protects `target`.
1082        unsafe {
1083            assert_eq!(*target.guard_mu(guard.token()).value(), 35);
1084        }
1085    }
1086
1087    #[guarded]
1088    struct FlaggedMutexStruct {
1089        #[mutex(flags = lockdep::LOCK_FLAGS_ACTIVE_LIST_DISABLED)]
1090        seek_lock: KMutex,
1091        #[guarded_by(seek_lock)]
1092        seek: u64,
1093    }
1094
1095    #[test]
1096    fn test_flagged_mutex_struct() {
1097        stack_pin_init!(let s = pin_init!(FlaggedMutexStruct {
1098            seek_lock <- KMutex::init(),
1099            seek: 0.into(),
1100        }));
1101
1102        lock!(let mut guard = s.lock_seek_lock());
1103        assert_eq!(*guard.seek(), 0);
1104        *guard.as_mut().seek_mut() = 1024;
1105        assert_eq!(*guard.seek(), 1024);
1106    }
1107
1108    #[test]
1109    fn test_guard_release() {
1110        stack_pin_init!(let s = pin_init!(MyGuardedStruct {
1111            mu <- KMutex::init(),
1112            data1: 10.into(),
1113            data2: 20.into(),
1114        }));
1115
1116        lock!(let mut guard = s.lock_mu());
1117        *guard.as_mut().data1_mut() = 50;
1118        guard.as_mut().release();
1119        // Calling release again is a no-op.
1120        guard.as_mut().release();
1121
1122        // Since `guard` was explicitly released, we can re-acquire the lock in the same scope.
1123        lock!(let guard2 = s.lock_mu());
1124        assert_eq!(*guard2.data1(), 50);
1125    }
1126
1127    #[test]
1128    fn test_guard_multiple() {
1129        stack_pin_init!(let s1 = pin_init!(MyGuardedStruct {
1130            mu <- KMutex::init(),
1131            data1: 1.into(),
1132            data2: 2.into(),
1133        }));
1134        stack_pin_init!(let s2 = pin_init!(MyGuardedStruct {
1135            mu <- KMutex::init(),
1136            data1: 3.into(),
1137            data2: 4.into(),
1138        }));
1139
1140        // Acquire in [s1, s2] order.
1141        {
1142            lock!(let mut guards = MyGuardedStruct::lock_mu_multiple([&*s1, &*s2]));
1143            let [t1, t2] = guards.as_mut().tokens_mut();
1144            // SAFETY: Both `s1.mu` and `s2.mu` are held by `guards`, and `s1` and `s2` are
1145            // distinct instances.
1146            unsafe {
1147                *s1.guard_mu_mut(t1).data1_mut() = 10;
1148                *s2.guard_mu_mut(t2).data1_mut() = 30;
1149            }
1150        }
1151
1152        // Acquire in reverse [s2, s1] order and release early.
1153        {
1154            lock!(let mut guards = MyGuardedStruct::lock_mu_multiple([&*s2, &*s1]));
1155            let [t1, t2] = guards.tokens();
1156            // SAFETY: Both `s1.mu` and `s2.mu` are held by `guards`.
1157            unsafe {
1158                assert_eq!(*s1.guard_mu(t1).data1(), 10);
1159                assert_eq!(*s2.guard_mu(t2).data1(), 30);
1160            }
1161            guards.as_mut().release();
1162
1163            // Verify both locks were released by re-acquiring them individually.
1164            lock!(let g1 = s1.lock_mu());
1165            lock!(let g2 = s2.lock_mu());
1166            assert_eq!(*g1.data1(), 10);
1167            assert_eq!(*g2.data1(), 30);
1168        }
1169    }
1170
1171    #[test]
1172    #[should_panic(expected = "cannot access token of a released KMutexGuard")]
1173    fn test_released_guard_token_panics() {
1174        stack_pin_init!(let s = pin_init!(MyGuardedStruct {
1175            mu <- KMutex::init(),
1176            data1: 1.into(),
1177            data2: 2.into(),
1178        }));
1179
1180        lock!(let mut guard = s.mu.lock());
1181        guard.as_mut().release();
1182        let _ = guard.token();
1183    }
1184
1185    #[test]
1186    #[should_panic(expected = "cannot access token of a released KMutexGuard")]
1187    fn test_released_guard_field_access_panics() {
1188        stack_pin_init!(let s = pin_init!(MyGuardedStruct {
1189            mu <- KMutex::init(),
1190            data1: 1.into(),
1191            data2: 2.into(),
1192        }));
1193
1194        lock!(let mut guard = s.lock_mu());
1195        guard.as_mut().release();
1196        let _ = guard.as_mut().data1_mut();
1197    }
1198
1199    #[test]
1200    #[should_panic(expected = "cannot access tokens of a released KMutexGuardMultiple")]
1201    fn test_released_guard_multiple_tokens_panics() {
1202        stack_pin_init!(let s1 = pin_init!(MyGuardedStruct {
1203            mu <- KMutex::init(),
1204            data1: 1.into(),
1205            data2: 2.into(),
1206        }));
1207        stack_pin_init!(let s2 = pin_init!(MyGuardedStruct {
1208            mu <- KMutex::init(),
1209            data1: 3.into(),
1210            data2: 4.into(),
1211        }));
1212
1213        lock!(let mut guards = MyGuardedStruct::lock_mu_multiple([&*s1, &*s2]));
1214        guards.as_mut().release();
1215        let _ = guards.tokens();
1216    }
1217
1218    #[test]
1219    #[forbid(unsafe_code)]
1220    fn test_safe_guard_from_guard_and_token_guard_methods() {
1221        stack_pin_init!(let s = pin_init!(MyGuardedStruct {
1222            mu <- KMutex::init(),
1223            data1: 10.into(),
1224            data2: 20.into(),
1225        }));
1226
1227        {
1228            lock!(let mut guard = s.mu.lock());
1229            {
1230                let tg = s.guard_mu_from_guard(&guard);
1231                assert_eq!(*tg.data1(), 10);
1232                assert_eq!(*tg.data2(), 20);
1233                assert!(core::ptr::eq(tg.parent(), &*s));
1234            }
1235            {
1236                let mut tg_mut = s.guard_mu_mut_from_guard(guard.as_mut());
1237                assert!(core::ptr::eq(tg_mut.parent(), &*s));
1238                *tg_mut.reborrow().data1_mut() = 11;
1239                assert_eq!(*tg_mut.as_token_guard().data1(), 11);
1240            }
1241        }
1242
1243        {
1244            lock!(let mut guard = s.lock_mu());
1245            assert!(core::ptr::eq(guard.parent(), &*s));
1246            assert_eq!(*guard.as_token_guard().data1(), 11);
1247            let mut tg_mut = guard.as_mut().as_token_guard_mut();
1248            *tg_mut.reborrow().data2_mut() = 22;
1249            assert_eq!(*tg_mut.as_token_guard().data2(), 22);
1250        }
1251    }
1252
1253    #[test]
1254    #[should_panic(expected = "KMutexGuard does not belong to this struct's mutex")]
1255    fn test_guard_mu_from_guard_mismatched_panics() {
1256        stack_pin_init!(let s1 = pin_init!(MyGuardedStruct {
1257            mu <- KMutex::init(),
1258            data1: 1.into(),
1259            data2: 2.into(),
1260        }));
1261        stack_pin_init!(let s2 = pin_init!(MyGuardedStruct {
1262            mu <- KMutex::init(),
1263            data1: 10.into(),
1264            data2: 20.into(),
1265        }));
1266
1267        lock!(let guard2 = s2.mu.lock());
1268        let _ = s1.guard_mu_from_guard(&guard2);
1269    }
1270
1271    #[test]
1272    #[should_panic(
1273        expected = "AliasedLock requires lock1 and lock2 to point to the same physical lock"
1274    )]
1275    fn test_aliased_lock_mismatched_panics() {
1276        stack_pin_init!(let s1 = pin_init!(MyGuardedStruct {
1277            mu <- KMutex::init(),
1278            data1: 1.into(),
1279            data2: 2.into(),
1280        }));
1281        stack_pin_init!(let s2 = pin_init!(MyGuardedStruct {
1282            mu <- KMutex::init(),
1283            data1: 10.into(),
1284            data2: 20.into(),
1285        }));
1286
1287        lock!(let _guard = aliased_lock(&s1.mu, &s2.mu));
1288    }
1289
1290    #[test]
1291    #[should_panic(expected = "KMutexGuardMultiple requires all locks to be distinct instances")]
1292    fn test_guard_multiple_duplicate_panics() {
1293        stack_pin_init!(let s1 = pin_init!(MyGuardedStruct {
1294            mu <- KMutex::init(),
1295            data1: 1.into(),
1296            data2: 2.into(),
1297        }));
1298
1299        lock!(let _guards = MyGuardedStruct::lock_mu_multiple([&*s1, &*s1]));
1300    }
1301
1302    #[guarded]
1303    struct SharedLockGuardedStruct {
1304        #[mutex]
1305        mu1: KMutex,
1306        #[mutex]
1307        mu2: KMutex,
1308        #[mutex]
1309        mu3: KMutex,
1310        #[guarded_by(mu1)]
1311        only1: u32,
1312        #[guarded_by(mu2)]
1313        only2: u32,
1314        #[guarded_by(mu1, mu2)]
1315        shared12: u32,
1316        #[guarded_by(mu1)]
1317        #[guarded_by(mu2)]
1318        #[pin]
1319        pinned_shared12: u32,
1320        #[guarded_by(mu1, mu2)]
1321        once_shared12: crate::KOnceCell<u32>,
1322        #[guarded_by(mu1, mu2, mu3)]
1323        shared123: u32,
1324    }
1325
1326    impl SharedLockGuardedStruct {
1327        fn init() -> impl pin_init::PinInit<Self, core::convert::Infallible> {
1328            pin_init!(Self {
1329                mu1 <- KMutex::init(),
1330                mu2 <- KMutex::init(),
1331                mu3 <- KMutex::init(),
1332                only1: 1.into(),
1333                only2: 2.into(),
1334                shared12: 12.into(),
1335                pinned_shared12: 120.into(),
1336                once_shared12: crate::KOnceCell::new(),
1337                shared123: 123.into(),
1338            })
1339        }
1340    }
1341
1342    #[test]
1343    #[forbid(unsafe_code)]
1344    fn test_guarded_by_multiple_locks() {
1345        stack_pin_init!(let s = SharedLockGuardedStruct::init());
1346
1347        // Reading `shared12`, `pinned_shared12`, `once_shared12`, and `shared123` is allowed when
1348        // holding only `mu1`.
1349        {
1350            lock!(let mut g1 = s.lock_mu1());
1351            assert_eq!(*g1.only1(), 1);
1352            assert_eq!(*g1.shared12(), 12);
1353            assert_eq!(*g1.pinned_shared12(), 120);
1354            assert_eq!(g1.once_shared12(), None);
1355            assert_eq!(*g1.shared123(), 123);
1356
1357            let f1 = g1.fields();
1358            assert_eq!(*f1.only1, 1);
1359            assert_eq!(*f1.shared12, 12);
1360            assert_eq!(*f1.pinned_shared12, 120);
1361            assert_eq!(f1.once_shared12, None);
1362            assert_eq!(*f1.shared123, 123);
1363
1364            // In `fields_mut()`, `only1` is mutable while multi-lock fields are read-only.
1365            let f1_mut = g1.as_mut().fields_mut();
1366            *f1_mut.only1 = 10;
1367            assert_eq!(*f1_mut.shared12, 12);
1368            assert_eq!(*f1_mut.pinned_shared12, 120);
1369            assert_eq!(f1_mut.once_shared12, None);
1370            assert_eq!(*f1_mut.shared123, 123);
1371        }
1372
1373        // Reading is also allowed when holding only `mu2`.
1374        {
1375            lock!(let g2 = s.lock_mu2());
1376            assert_eq!(*g2.only2(), 2);
1377            assert_eq!(*g2.shared12(), 12);
1378            assert_eq!(*g2.pinned_shared12(), 120);
1379            assert_eq!(g2.once_shared12(), None);
1380            assert_eq!(*g2.shared123(), 123);
1381        }
1382
1383        // Mutating `shared12`, `pinned_shared12`, and `once_shared12` when holding both `mu1` and
1384        // `mu2` (via either guard or token guards).
1385        {
1386            lock!(let mut g1 = s.lock_mu1());
1387            lock!(let mut g2 = s.lock_mu2());
1388
1389            *g1.as_mut().shared12_mut(&mut g2) = 42;
1390            assert_eq!(*g1.shared12(), 42);
1391            assert_eq!(*g2.shared12(), 42);
1392
1393            *g2.as_mut().shared12_mut(&mut g1) = 84;
1394            assert_eq!(*g1.shared12(), 84);
1395
1396            *g1.as_mut().pinned_shared12_mut(&mut g2) = 240;
1397            assert_eq!(*g2.pinned_shared12(), 240);
1398
1399            let mut cell = g2.as_mut().once_shared12_cell(&mut g1);
1400            let val = cell.get_or_init(|| 777);
1401            assert_eq!(*val, 777);
1402            assert_eq!(g1.once_shared12(), Some(&777));
1403            assert_eq!(g2.once_shared12(), Some(&777));
1404
1405            if let Some(v) = g1.as_mut().once_shared12_mut(&mut g2) {
1406                *v = 888;
1407            }
1408            assert_eq!(g2.once_shared12(), Some(&888));
1409
1410            // TokenGuardMut access for standard, pinned, and KOnceCell multi-lock fields.
1411            let mut tg1 = g1.as_mut().as_token_guard_mut();
1412            let mut tg2 = g2.as_mut().as_token_guard_mut();
1413            *tg1.reborrow().shared12_mut(&mut tg2) = 99;
1414            assert_eq!(*tg2.shared12(), 99);
1415            assert_eq!(*tg2.as_token_guard().shared12(), 99);
1416            *tg2.reborrow().shared12_mut(&mut tg1) = 100;
1417            assert_eq!(*tg1.shared12(), 100);
1418            assert_eq!(*tg1.as_token_guard().shared12(), 100);
1419
1420            assert_eq!(*tg1.pinned_shared12(), 240);
1421            *tg1.reborrow().pinned_shared12_mut(&mut tg2) = 360;
1422            assert_eq!(*tg2.pinned_shared12(), 360);
1423            assert_eq!(*tg2.as_token_guard().pinned_shared12(), 360);
1424
1425            assert_eq!(tg1.once_shared12(), Some(&888));
1426            assert_eq!(tg1.as_token_guard().once_shared12(), Some(&888));
1427            if let Some(v) = tg1.reborrow().once_shared12_mut(&mut tg2) {
1428                *v = 999;
1429            }
1430            assert_eq!(tg2.reborrow().once_shared12_cell(&mut tg1).get(), Some(&999));
1431            assert_eq!(tg1.fields().once_shared12, Some(&999));
1432            assert_eq!(tg1.reborrow().fields_mut().once_shared12, Some(&999));
1433            assert_eq!(g1.as_mut().fields_mut().once_shared12, Some(&999));
1434
1435            // Mutating 3-lock field `shared123` requires all 3 locks.
1436            lock!(let mut g3 = s.lock_mu3());
1437            *g1.as_mut().shared123_mut(&mut g2, &mut g3) = 321;
1438            assert_eq!(*g3.shared123(), 321);
1439            *g3.as_mut().shared123_mut(&mut g1, &mut g2) = 654;
1440            assert_eq!(*g1.shared123(), 654);
1441        }
1442    }
1443
1444    #[test]
1445    #[should_panic(expected = "Lock guard belongs to a different struct instance")]
1446    fn test_multi_guarded_wrong_instance_panics() {
1447        stack_pin_init!(let s1 = SharedLockGuardedStruct::init());
1448        stack_pin_init!(let s2 = SharedLockGuardedStruct::init());
1449
1450        lock!(let mut g1 = s1.lock_mu1());
1451        lock!(let mut g2_wrong = s2.lock_mu2());
1452
1453        // Passing `g2_wrong` (locked on `s2`) to `g1` (locked on `s1`) must panic.
1454        *g1.as_mut().shared12_mut(&mut g2_wrong) = 42;
1455    }
1456
1457    #[test]
1458    #[should_panic(expected = "Lock guard belongs to a different struct instance")]
1459    fn test_multi_guarded_token_guard_wrong_instance_panics() {
1460        stack_pin_init!(let s1 = SharedLockGuardedStruct::init());
1461        stack_pin_init!(let s2 = SharedLockGuardedStruct::init());
1462
1463        lock!(let mut g1 = s1.lock_mu1());
1464        lock!(let mut g2_wrong = s2.lock_mu2());
1465
1466        let mut tg1 = g1.as_mut().as_token_guard_mut();
1467        let mut tg2_wrong = g2_wrong.as_mut().as_token_guard_mut();
1468
1469        // Passing `tg2_wrong` (bound to `s2`) to `tg1` (bound to `s1`) must panic.
1470        *tg1.shared12_mut(&mut tg2_wrong) = 42;
1471    }
1472
1473    #[ksync::guarded]
1474    struct GenericMultiLockStruct<'lt, T, const N: usize> {
1475        #[mutex]
1476        mu1: KMutex,
1477        #[mutex]
1478        mu2: KMutex,
1479        #[guarded_by(mu1)]
1480        slice_ref: &'lt [T; N],
1481        #[guarded_by(mu1, mu2)]
1482        shared: [T; N],
1483    }
1484
1485    #[test]
1486    #[forbid(unsafe_code)]
1487    fn test_generic_multi_lock_struct() {
1488        static BACKING: [u32; 2] = [10, 20];
1489        stack_pin_init!(let s = pin_init!(GenericMultiLockStruct::<u32, 2> {
1490            mu1 <- KMutex::init(),
1491            mu2 <- KMutex::init(),
1492            slice_ref: (&BACKING).into(),
1493            shared: [1u32, 2u32].into(),
1494        }));
1495
1496        lock!(let mut g1 = s.lock_mu1());
1497        assert_eq!(*g1.slice_ref(), &[10, 20]);
1498        assert_eq!(*g1.shared(), [1, 2]);
1499        assert_eq!(*g1.fields().shared, [1, 2]);
1500        assert_eq!(*g1.as_mut().fields_mut().shared, [1, 2]);
1501
1502        lock!(let mut g2 = s.lock_mu2());
1503        assert_eq!(*g2.shared(), [1, 2]);
1504        *g1.as_mut().shared_mut(&mut g2) = [30, 40];
1505        assert_eq!(*g2.shared(), [30, 40]);
1506    }
1507
1508    #[ksync::guarded]
1509    struct PhantomMultiLockStruct {
1510        #[mutex(MyGuardedStructMuClass)]
1511        phantom_mu: KMutex<crate::PhantomMutex>,
1512        #[mutex]
1513        local_mu: KMutex,
1514        #[guarded_by(phantom_mu, local_mu)]
1515        shared: u32,
1516    }
1517
1518    #[test]
1519    fn test_phantom_and_custom_class_multi_lock() {
1520        stack_pin_init!(let real_obj = pin_init!(MyGuardedStruct {
1521            mu <- KMutex::init(),
1522            data1: 0.into(),
1523            data2: 0.into(),
1524        }));
1525        stack_pin_init!(let s = pin_init!(PhantomMultiLockStruct {
1526            phantom_mu: KMutex::new(crate::PhantomMutex),
1527            local_mu <- KMutex::init(),
1528            shared: 50.into(),
1529        }));
1530
1531        // SAFETY: `real_obj.mu` is the physical lock synchronizing `s.phantom_mu`.
1532        lock!(let mut pg = unsafe { s.lock_phantom_mu(&real_obj.mu) });
1533        assert_eq!(*pg.shared(), 50);
1534
1535        lock!(let mut lg = s.lock_local_mu());
1536        assert_eq!(*lg.shared(), 50);
1537
1538        *pg.as_mut().shared_mut(&mut lg) = 60;
1539        assert_eq!(*lg.shared(), 60);
1540
1541        *lg.as_mut().shared_mut(&mut pg) = 70;
1542        assert_eq!(*pg.shared(), 70);
1543    }
1544
1545    #[ksync::guarded]
1546    struct UntrackedGuardedStruct {
1547        #[mutex]
1548        mu1: KMutex<crate::RawMutexStorage>,
1549        #[mutex(untracked)]
1550        mu2: KMutex<crate::RawMutexStorage>,
1551        #[guarded_by(mu1)]
1552        val1: u32,
1553        #[guarded_by(mu2)]
1554        val2: u32,
1555    }
1556
1557    #[test]
1558    fn test_untracked_guarded_and_as_mut_ptr() {
1559        stack_pin_init!(let tracked = pin_init!(MyGuardedStruct {
1560            mu <- KMutex::init(),
1561            data1: 1.into(),
1562            data2: 2.into(),
1563        }));
1564        assert!(core::ptr::eq(
1565            <crate::RawMutex as crate::RawLock>::as_mut_ptr(tracked.mu.raw_mutex()),
1566            <crate::RawMutexStorage as crate::RawLock>::as_mut_ptr(&**tracked.mu.raw_mutex()),
1567        ));
1568
1569        stack_pin_init!(let untracked = pin_init!(UntrackedGuardedStruct {
1570            mu1 <- KMutex::init(),
1571            mu2 <- KMutex::init(),
1572            val1: 10.into(),
1573            val2: 20.into(),
1574        }));
1575        assert!(<UntrackedGuardedStructMu1Class as LockClass>::ID.is_null());
1576        assert!(<UntrackedGuardedStructMu2Class as LockClass>::ID.is_null());
1577
1578        lock!(let mut g1 = untracked.lock_mu1());
1579        lock!(let mut g2 = untracked.lock_mu2());
1580        assert_eq!(*g1.val1(), 10);
1581        assert_eq!(*g2.val2(), 20);
1582        *g1.as_mut().val1_mut() = 11;
1583        *g2.as_mut().val2_mut() = 21;
1584        assert_eq!(*g1.val1(), 11);
1585        assert_eq!(*g2.val2(), 21);
1586    }
1587}