Skip to main content

facet_maybe_mut/
maybe_mut.rs

1use derive_more::{Deref, DerefMut, From};
2use facet::{Def, PointerFlags, PtrConst, PtrMut, ReadLockResult, Shape, WriteLockResult};
3use facet_reflect::{Peek, Poke};
4
5/// Some reference to a type that implements [`Facet`](facet::Facet) that may be
6/// `mut` or not.
7#[derive(From)]
8#[repr(C)]
9pub enum MaybeMut<'mem, 'facet> {
10    Not(Peek<'mem, 'facet>),
11    Mut(Poke<'mem, 'facet>),
12}
13
14impl<'mem, 'facet> MaybeMut<'mem, 'facet> {
15    /// Returns a readonly/immutable version of the inner type
16    pub fn as_peek(&'mem self) -> Peek<'mem, 'facet> {
17        match self {
18            Self::Not(peek) => *peek,
19            Self::Mut(poke) => poke.as_peek(),
20        }
21    }
22
23    pub fn into_peek(self) -> Peek<'mem, 'facet> {
24        match self {
25            MaybeMut::Not(n) => n,
26            MaybeMut::Mut(m) => m.into_peek(),
27        }
28    }
29
30    /// Returns the [`Shape`] of the underlying type
31    ///
32    /// The [`Shape`] is the same for [`Mut`](Self::Mut) and [`Not`](Self::Not)
33    pub fn shape(&self) -> &'static Shape {
34        self.as_peek().shape()
35    }
36}
37
38#[derive(Debug, thiserror::Error)]
39#[error("{kind}")]
40pub struct MakeLockError<'mem, 'facet> {
41    pub unchanged: Peek<'mem, 'facet>,
42    pub kind: MakeLockErrorKind,
43}
44
45#[derive(Debug, thiserror::Error)]
46pub enum MakeLockErrorKind {
47    /// The underlying type is not a type that we can lock from a `&T` to `&mut T` (but it is more complicated...)
48    #[error("type cannot be locked")]
49    NotLockable,
50    /// The underlying type could be locked but the provided lock method in the
51    /// vtable returned an error.
52    #[error("locking of type failed")]
53    LockFailure,
54    /// A weak pointer where the upgrade function returned None.
55    ///
56    /// There exist no strong references (no instances of an `Arc`)
57    #[error("could not upgrade weak pointer, no strong references exist")]
58    NotUpgradable,
59}
60
61/// Depending on whether this is a read or write lock, `P` will be either
62/// [`PtrConst`] or [`PtrMut`](facet::PtrMut). This enum makes `P` dynamic
63#[derive(From)]
64pub(crate) enum LockGuardType {
65    Write(WriteLockResult),
66    Read(ReadLockResult),
67    /// A Weak that has been upgraded to an Arc or Rc
68    /// The downgrade is handled directly in the [`Drop`] implementation of this
69    /// [`LockGuardType`]
70    Upgrade {
71        /// The [`Shape`] of the strong pointer that can be used later to
72        /// downgrade again
73        strong_shape: &'static Shape,
74        /// The pointer to the allocated Arc or Rc
75        allocation: PtrMut,
76    },
77}
78
79impl LockGuardType {
80    /// Safety
81    ///
82    /// For an Weak guard, this calls the `BorrowFn` of the (currently existing)
83    /// strong shape to obtain the data pointer.
84    /// This is the raw pointer returned from the lock which is already
85    /// available via [`Guard`]. Creating a new [`Peek`] or [`Poke`] from this
86    /// [`PtrConst`] is UB.
87    pub fn data_const(&self) -> PtrConst {
88        match self {
89            Self::Write(w) => w.data_const(),
90            Self::Read(r) => *r.data(),
91            Self::Upgrade {
92                strong_shape,
93                allocation,
94            } => {
95                let borrow_fn = strong_shape
96                    .def
97                    .into_pointer()
98                    .expect("only pointer types get this lock type")
99                    .vtable
100                    .borrow_fn
101                    .expect("all strong pointers have a borrow function");
102                // SAFETY: allocation is the pointer of the strong type (Arc or Rc)
103                unsafe { borrow_fn(allocation.as_const()) }
104            }
105        }
106    }
107}
108
109impl Drop for LockGuardType {
110    fn drop(&mut self) {
111        if let Self::Upgrade {
112            strong_shape,
113            allocation,
114        } = self
115        {
116            // dropping the strong pointer automatically decreases reference
117            // count
118            // SAFETY: we cant just deallocate but need to run the actual drop
119            // implementation as well since the drop impl of the strong pointer decreases strong pointer count
120            unsafe {
121                strong_shape.call_drop_in_place(*allocation);
122            }
123            // SAFETY: the allocation was created using Shape::alloc
124            unsafe {
125                strong_shape
126                    .deallocate_mut(*allocation)
127                    .expect("strong pointer is sized");
128            }
129        }
130    }
131}
132
133/// Contains the guard, the data ptr, and drop vtable to free the lock
134///
135/// # Note
136///
137/// The contained [`MaybeMut`] is NOT guaranteed to be [`Mut`](MaybeMut::Mut)
138///
139/// For example, RwLock also needs a lock and guard for a read.
140///
141#[derive(Deref, DerefMut)]
142pub struct Guard<'lock_mem, 'facet> {
143    /// Dropping the guard handles freeing the lock
144    ///
145    /// If this is empty, the `data` can be accessed directly and there is no
146    /// lock that must be freeed
147    ///
148    /// SAFETY: The pointer inside the [`LockGuardType`] MUST NOT be used
149    /// since the data is already (mutable) available via `data`
150    ///
151    /// The drop order must be the reverse of this [`Vec`] (pop until empty)
152    /// in order to guarantee the locks are released in the correct order.
153    guards: Vec<LockGuardType>,
154    /// This [`MaybeMut`] contains the [`Peek`] or [`Poke`] of the most inner,
155    /// non lockable type.
156    #[deref]
157    #[deref_mut]
158    data: MaybeMut<'lock_mem, 'facet>,
159}
160
161impl<'mem, 'facet> MaybeMut<'mem, 'facet> {
162    /// Try to turn [`MaybeMut::Not`] into [`MaybeMut::Mut`]
163    ///
164    /// The returned [`MaybeMut`] may contain a different [`Shape`].
165    /// Which exact [`Shape`] it is, depends on what the input type was.
166    ///
167    /// One edge case is if you pass a `&mut Arc<RwLock<String>` the type will
168    /// not be changed to `&mut String`. But if you pass a `&Arc<RwLock<String>`
169    /// due to locking etc, it will be a `&mut String`.
170    ///
171    /// If the underlying type is something that can be write locked,
172    /// for example an `RwLock` or `Mutex`, this method creates a lock on it.
173    ///
174    /// If we already have [`MaybeMut::Mut`] this is a no-op.
175    ///
176    /// If we have [`MaybeMut::Not`] and the [`Shape`] of
177    /// `T` does not contain a [`PointerDef`](facet::PointerDef) which
178    /// has a vtable with a `write_fn` we can call with `&T`, this method
179    /// returns [`Err(MaybeMut::Not)`](Err). In this case, besides the lookup,
180    /// it is also a no-op.
181    ///
182    /// # Note
183    ///
184    /// It is very important that you drop the [`Guard`] as soon as possible
185    /// to free the lock
186    ///
187    /// [`Shape`]: facet::Shape
188    pub fn write<'lock>(self) -> Result<Guard<'lock, 'facet>, MakeLockError<'mem, 'facet>>
189    where
190        'mem: 'lock,
191    {
192        match self {
193            // if we already have a mut this is a no op
194            MaybeMut::Mut(v) => {
195                // but only if this is a type that is not a smart pointer that can be locked
196                if let Def::Pointer(p) = v.as_peek().innermost_peek().shape().def
197                // restrict downgrading to Peek only if there is a a lock _somewhere_
198                    && (p.flags.contains(PointerFlags::LOCK) || p.flags.contains(PointerFlags::WEAK))
199                {
200                    Self::Not(v.into_peek()).write()
201                } else {
202                    Ok(Guard {
203                        guards: Vec::new(),
204                        data: v.into(),
205                    })
206                }
207            }
208            // this is where it gets interesting
209            MaybeMut::Not(v) => {
210                // SAFETY: v.innermost_peek() unwraps all transparent wrappers like Arc or Rc until something that needs
211                // locking is reached which is all we care about
212                // FIXME: naively using innermost_peek is a bad idea i think.
213                // for example, in the UI if there is a NonZero<u32> this will peek
214                // up tu u32. Then, we will perhaps display an editable u32 which
215                // can be set to zero. Now what? We broke it boys
216                let v = v.innermost_peek();
217                // the shape of the pointer type (if it is one) but derefence smart pointers that can so without locking
218                // e.g. Arc<T> AND also &T
219                let shape = v.shape();
220                let def = shape.def;
221
222                // short cirucit if it is not a pointer. in these
223                // cases we wont be able to reach something like
224                // RwLock or Mutex
225                let Def::Pointer(pointer) = def else {
226                    return Err(MakeLockError {
227                        unchanged: v,
228                        kind: MakeLockErrorKind::NotLockable,
229                    });
230                };
231
232                // we dont care if we lock it (Mutex) or write lock it (RwLock)
233                let lock_fn =
234                    pointer
235                        .vtable
236                        .write_fn
237                        .or(pointer.vtable.lock_fn)
238                        .ok_or(MakeLockError {
239                            unchanged: v,
240                            kind: MakeLockErrorKind::NotLockable,
241                        });
242
243                // SAFETY: v.innermost_peek() unwraps all transparent wrappers like Arc or Rc until something that needs
244                // locking is reached which is also the same type we get the lock_fn from
245                let (mut guards, mut value): (Vec<LockGuardType>, MaybeMut<'lock, 'facet>) =
246                    match lock_fn {
247                        Ok(lock_fn) => {
248                            let res = unsafe { lock_fn(v.data()) };
249                            let Ok(lock) = res else {
250                                return Err(MakeLockError {
251                                    unchanged: v,
252                                    kind: MakeLockErrorKind::LockFailure,
253                                });
254                            };
255                            // SAFETY: creates access via the PtrMut returned from locking
256                            // the smart pointer. 'mem outlives 'lock this means
257                            // the returned SmartPointer<'mem> also outlives the mutable Poke<'lock>
258                            let poke: Poke<'lock, 'facet> = unsafe {
259                                Poke::from_raw_parts(
260                                    // if the input type was Arc<RwLock<String>> this willbe
261                                    // a pointer to a String
262                                    *lock.data(),
263                                    shape
264                                        .inner
265                                        .expect("a smart pointer always has an inner shape"),
266                                )
267                            };
268
269                            (vec![lock.into()], MaybeMut::Mut(poke))
270                        }
271                        // TODO: write upgrade Weak<RwLock>??
272                        // try it as an upgrade instead
273                        Err(MakeLockError {
274                            unchanged,
275                            kind: MakeLockErrorKind::NotLockable,
276                        }) if let Def::Pointer(pointer) = unchanged.shape().def
277                            && let Some(upgrade_fn) = pointer.vtable.upgrade_into_fn
278                            && let Some(strong_shape) =
279                                def.into_pointer().ok().and_then(|x| x.strong()) =>
280                        {
281                            // if the strong shape is unsized, the Facet implementation of the type is wrong.
282                            let strong = strong_shape
283                                .allocate()
284                                .expect("strong pointer is always sized");
285
286                            // SAFETY: turning this peek into a PtrMut is okay,. because
287                            // the upgrade function only needs &self
288                            // in theory, the upgrade_fn signature could take a PtrConst as well.
289                            let ptr = unsafe { v.data().into_mut() };
290                            // SAFETY: Facet implementation of Weak garantees strong is the correct
291                            // shape of the strong part for this Weak
292                            let guard = unsafe { upgrade_fn(ptr, strong) }
293                                .map(|strong_instance| LockGuardType::Upgrade {
294                                    strong_shape,
295                                    allocation: strong_instance,
296                                })
297                                .ok_or(MakeLockError {
298                                    kind: MakeLockErrorKind::NotUpgradable,
299                                    unchanged: v,
300                                })?;
301                            // SAFETY: creates access via the PtrMut returned from locking
302                            // the smart pointer. 'mem outlives 'lock this means
303                            // the returned mutable Poke<'lock> lives shorter than 'mem
304                            let peek: Peek<'lock, 'facet> = unsafe {
305                                Peek::unchecked_new(
306                                    guard.data_const(),
307                                    shape
308                                        .inner
309                                        .expect("a smart pointer always has an inner shape"),
310                                )
311                            };
312
313                            (vec![guard], MaybeMut::Not(peek.innermost_peek()))
314                        }
315                        Err(e) => {
316                            return Err(e);
317                        }
318                    };
319
320                // unwrap remaining inner pointer types
321                // -> all types that have an inner type and are a pointer
322                while let Some(_inner) = value.as_peek().shape().inner
323                    && let Def::Pointer(def) = value.as_peek().shape().def
324                    // lock gets locked in the next write call
325                    && (def.flags.contains(PointerFlags::LOCK) ||
326                    // weak gets upgraded at the next write call
327                    def.flags.contains(PointerFlags::WEAK) ||
328                    // atomics just get unwrapped in the next write call
329                    def.flags.contains(PointerFlags::ATOMIC))
330                {
331                    // SAFETY: we synthesize a Peek with the outer 'mem lifetime so we
332                    // can re-enter `write`. The fabricated lifetime never escapes:
333                    //
334                    // * On success, the recursive call returns `Guard<'lock>`. Its
335                    //   `data` is bounded by 'lock and its guards are moved into our
336                    //   `guards` Vec, so the parent lock keeps the pointer live for
337                    //   as long as the returned `Guard` exists.
338                    // * On failure, we MUST NOT propagate the inner
339                    //   `MakeLockError::unchanged` since that Peek carries the
340                    //   fabricated 'mem lifetime while actually pointing into
341                    //   lock-protected memory that we are about to release as our
342                    //   `guards` Vec drops on early return. Instead we substitute
343                    //   the original outer Peek `v`, which is genuinely valid for
344                    //   'mem because it comes straight from the function input.
345                    let shorter_peek: Peek<'mem, 'facet> =
346                        unsafe { Peek::unchecked_new(value.as_peek().data(), value.shape()) };
347                    let shorter_maybe: MaybeMut<'mem, 'facet> = MaybeMut::Not(shorter_peek);
348                    let guard: Guard<'lock, 'facet> = match shorter_maybe.write() {
349                        Ok(g) => g,
350                        Err(e) => {
351                            // Peek v needs no Guards, drop them explicitly to free
352                            // locks (they would be dropped by the return anyways)
353                            drop(guards);
354                            // Discard `e.unchanged` (would dangle once `guards`
355                            // drops on return); substitute the original outer Peek
356                            // which is actually valid for 'mem.
357                            return Err(MakeLockError {
358                                unchanged: v,
359                                kind: e.kind,
360                            });
361                        }
362                    };
363                    // SAFETY: the values must be moved to the new vec not cloned and NOT dropped.
364                    // this would lead to a double unlock later
365                    guards.extend(guard.guards);
366                    // SAFETY: set the new value to the inner most available value
367                    value = guard.data;
368                }
369                Ok(Guard {
370                    data: value,
371                    guards,
372                })
373            }
374        }
375    }
376
377    /// Returns a [`Guard`] with a lock that is sufficent for reading.
378    ///
379    /// In case of `RwLock` it is locked to read. If it is a `Mutex`, it must
380    /// be exclusively locked to write but we only consider it being read which
381    /// is safe
382    pub fn read<'lock>(self) -> Result<Guard<'lock, 'facet>, MakeLockError<'mem, 'facet>>
383    where
384        'mem: 'lock,
385    {
386        let peek = self.into_peek();
387        // unwrap smart pointers
388        // this will deref Arcs but not Weaks, Weaks are handled special as a guard that automatically downgrades them on Drop
389        let v = peek.innermost_peek();
390        // the shape of the pointer type (if it is one) but derefence smart pointers that can so without locking
391        // e.g. Arc<T>
392        let shape = v.shape();
393        let def = shape.def;
394
395        // short cirucit if it is not a pointer. in these
396        // cases we wont be able to reach something like
397        // RwLock or Mutex
398        // In this case, we just return the reference to the underlying type
399        let Def::Pointer(pointer) = def else {
400            return Ok(Guard {
401                guards: Vec::new(),
402                data: MaybeMut::Not(v),
403            });
404        };
405
406        // we dont care if we lock it (Mutex) or read lock it (RwLock) or just upgrade it so it can be dereferenced
407        let res: Result<LockGuardType, _> = if let Some(read_fn) = pointer.vtable.read_fn {
408            unsafe { read_fn(v.data()) }.map(Into::into)
409        } else if let Some(lock_fn) = pointer.vtable.lock_fn {
410            unsafe { lock_fn(v.data()) }.map(Into::into)
411            // handle weak pointers and try to upgrade them
412        } else if let Some(upgrade_fn) = pointer.vtable.upgrade_into_fn
413            && let Some(strong_shape) = def.into_pointer().ok().and_then(|x| x.strong())
414        {
415            // if the strong shape is unsized, the Facet implementation of the type is wrong.
416            let strong = strong_shape
417                .allocate()
418                .expect("strong pointer is always sized");
419
420            // SAFETY: turning this peek into a PtrMut is okay, because
421            // the upgrade function only needs &self
422            // in theory, the upgrade_fn signature could take a PtrConst as well.
423            let ptr = unsafe { v.data().into_mut() };
424            // SAFETY: Facet implementation of Weak garantees strong is the correct
425            // shape of the strong part for this Weak
426            Ok(unsafe { upgrade_fn(ptr, strong) }
427                .map(|strong_instance| LockGuardType::Upgrade {
428                    strong_shape,
429                    allocation: strong_instance,
430                })
431                .ok_or(MakeLockError {
432                    kind: MakeLockErrorKind::NotUpgradable,
433                    unchanged: v,
434                })?)
435        } else {
436            return Err(MakeLockError {
437                unchanged: v,
438                kind: MakeLockErrorKind::NotLockable,
439            });
440        };
441
442        let Ok(lock) = res else {
443            return Err(MakeLockError {
444                unchanged: v,
445                kind: MakeLockErrorKind::LockFailure,
446            });
447        };
448        // SAFETY: creates access via the PtrMut returned from locking
449        // the smart pointer. 'mem outlives 'lock this means
450        // the returned mutable Poke<'lock> lives shorter than 'mem
451        let peek: Peek<'lock, 'facet> = unsafe {
452            Peek::unchecked_new(
453                lock.data_const(),
454                shape
455                    .inner
456                    .expect("a smart pointer always has an inner shape"),
457            )
458        };
459        let (mut guards, mut value): (Vec<LockGuardType>, MaybeMut<'lock, 'facet>) =
460            (vec![lock], MaybeMut::Not(peek.innermost_peek()));
461
462        // unwrap remaining inner pointer types
463        // -> all types that have an inner type and are a pointer
464        while let Some(_inner) = value.as_peek().shape().inner
465            && let Def::Pointer(def) = value.as_peek().shape().def
466            // lock gets read-locked in the next read call
467            && (def.flags.contains(PointerFlags::LOCK) ||
468            // weak gets upgraded at the next read call
469            def.flags.contains(PointerFlags::WEAK) ||
470            // atomics just get unwrapped in the next read call
471            def.flags.contains(PointerFlags::ATOMIC))
472        {
473            // SAFETY: we synthesize a Peek with the outer 'mem lifetime so we
474            // can re-enter `read`. The fabricated lifetime never escapes:
475            //
476            // * On success, the recursive call returns `Guard<'lock>`. Its
477            //   `data` is bounded by 'lock and its guards are moved into our
478            //   `guards` Vec, so the parent lock keeps the pointer live for
479            //   as long as the returned `Guard` exists.
480            // * On failure, we MUST NOT propagate the inner
481            //   `MakeLockError::unchanged` since that Peek carries the
482            //   fabricated 'mem lifetime while actually pointing into
483            //   lock-protected memory that we are about to release as our
484            //   `guards` Vec drops on early return. Instead we substitute
485            //   the original outer Peek `v`, which is genuinely valid for
486            //   'mem because it comes straight from the function input.
487            let shorter_peek: Peek<'mem, 'facet> =
488                unsafe { Peek::unchecked_new(value.as_peek().data(), value.shape()) };
489            let shorter_maybe: MaybeMut<'mem, 'facet> = MaybeMut::Not(shorter_peek);
490            let guard: Guard<'lock, 'facet> = match shorter_maybe.read() {
491                Ok(g) => g,
492                Err(e) => {
493                    // Peek v needs no Guards, drop them explicitly to free
494                    // locks (they would be dropped by the return anyways)
495                    drop(guards);
496                    // Discard `e.unchanged` (would dangle once `guards`
497                    // drops on return); substitute the original outer Peek
498                    // which is actually valid for 'mem.
499                    return Err(MakeLockError {
500                        unchanged: v,
501                        kind: e.kind,
502                    });
503                }
504            };
505            // SAFETY: the values must be moved to the new vec not cloned and NOT dropped.
506            // this would lead to a double unlock later
507            guards.extend(guard.guards);
508            // SAFETY: set the new value to the inner most available value
509            value = guard.data;
510        }
511        Ok(Guard {
512            data: value,
513            guards,
514        })
515    }
516}
517
518#[cfg(test)]
519mod tests {
520    use facet::{Def, Facet, KnownPointer};
521    use facet_reflect::Peek;
522
523    #[derive(Debug, Facet)]
524    struct Foo {
525        value: String,
526    }
527
528    #[facet_testhelpers::test]
529    fn shared_reference() {
530        let a = Foo {
531            value: "aaaa".to_string(),
532        };
533        println!("{:#?}", <&Foo as Facet<'_>>::SHAPE.def);
534        assert!(
535            matches!(<&Foo as Facet<'_>>::SHAPE.def, Def::Pointer(p) if p.known == Some(KnownPointer::SharedReference))
536        );
537        let ref_a: &Foo = &a;
538        let ref_ref_a: &&Foo = &ref_a;
539        let peek = Peek::new(ref_ref_a);
540        println!("{:#?}", peek.shape().def); // `Undefined`
541        assert!(
542            matches!(peek.shape().def, Def::Pointer(p) if p.known == Some(KnownPointer::SharedReference))
543        );
544    }
545}