gazebo/
cell.rs

1/*
2 * Copyright (c) Meta Platforms, Inc. and affiliates.
3 *
4 * This source code is licensed under both the MIT license found in the
5 * LICENSE-MIT file in the root directory of this source tree and the Apache
6 * License, Version 2.0 found in the LICENSE-APACHE file in the root directory
7 * of this source tree.
8 */
9
10//! Additions to the [`Ref`](Ref) mechanism.
11
12// We used to implement `ARef` as an enum of `{Ptr(&'a T), Ref(Ref<'a, T>)}`.
13// That works, but consumes 3 words and requires a branch on every access of the underlying
14// pointer. We can optimise that by relying on the underlying details of `Ref`, which is
15// (currently) defined as (after a bit of inlining):
16//
17// ```
18// pub struct Ref<'a, T: ?Sized + 'a> {
19//    value: &'a T,
20//    borrow: &'a Cell<isize>,
21// }
22// ```
23//
24// Because the pointer must always be non-null, we can switch that out for:
25//
26// ```
27// pub struct ARef<'a, T: ?Sized + 'a> {
28//    value: &'a T,
29//    borrow: Option<&'a Cell<isize>>,
30// }
31// ```
32//
33// And use `None` to represent the `Ptr` case. We do that with transmute trickery,
34// but write some good tests that will break if the representation changes,
35// and if necessary we can always switch to the enum representation.
36
37use std::cell::BorrowError;
38use std::cell::Cell;
39use std::cell::Ref;
40use std::cell::RefCell;
41use std::cmp::Ordering;
42use std::fmt;
43use std::fmt::Display;
44use std::hash::Hash;
45use std::hash::Hasher;
46use std::mem;
47use std::ops::Deref;
48
49use crate::cast;
50
51/// A [`Ref`](Ref) that might not actually be borrowed.
52/// Either a `Ptr` (a normal & style reference), or a `Ref` (like from
53/// [`RefCell`](std::cell::RefCell)), but exposes all the methods available on [`Ref`](Ref).
54#[derive(Debug)]
55pub struct ARef<'a, T: ?Sized + 'a> {
56    value: &'a T,
57    borrow: Option<&'a Cell<isize>>,
58}
59
60impl<T: ?Sized> Drop for ARef<'_, T> {
61    fn drop(&mut self) {
62        if self.borrow.is_some() {
63            let me: ARef<T> = ARef {
64                value: self.value,
65                borrow: self.borrow,
66            };
67            // The transmute forgets me, so I won't get called recursively
68            let them: Ref<T> = unsafe { cast::transmute_unchecked(me) };
69            // But we can now drop on the Ref
70            mem::drop(them);
71        }
72    }
73}
74
75impl<T: ?Sized> Deref for ARef<'_, T> {
76    type Target = T;
77
78    fn deref(&self) -> &T {
79        self.value
80    }
81}
82
83impl<'a, T: ?Sized + 'a> ARef<'a, T> {
84    /// Create a new [`ARef`] from a pointer.
85    pub fn new_ptr(x: &'a T) -> Self {
86        ARef {
87            value: x,
88            borrow: None,
89        }
90    }
91
92    /// Create a new [`ARef`] from a reference.
93    pub fn new_ref(x: Ref<'a, T>) -> Self {
94        // This is safe if the representation is the same as we expect,
95        // which we check for in a test below.
96        // Unfortunately, we can't directly transmute between Ref<T> and ARef<T>
97        // as the type T is generic. So we have to transmute between an intermediate (usize, usize)
98        let v: ARef<T> = unsafe { cast::transmute_unchecked(x) };
99        debug_assert!(v.borrow.is_some());
100        v
101    }
102
103    /// See [`Ref.clone`](Ref::clone). Not a self method since that interferes with the [`Deref`](Deref).
104    #[allow(clippy::should_implement_trait)]
105    pub fn clone(orig: &Self) -> Self {
106        if orig.borrow.is_none() {
107            ARef::new_ptr(orig.value)
108        } else {
109            let orig: &Ref<T> = unsafe { cast::ptr(orig) };
110            Self::new_ref(Ref::clone(orig))
111        }
112    }
113
114    /// See [`Ref.map`](Ref::map). Not a self method since that interferes with the [`Deref`](Deref).
115    pub fn map<U: ?Sized, F>(orig: ARef<'a, T>, f: F) -> ARef<'a, U>
116    where
117        F: FnOnce(&T) -> &U,
118    {
119        // The `map` implementation for Ref doesn't touch the borrow, so we just use the pointer.
120        let res = ARef {
121            value: f(orig.value),
122            borrow: orig.borrow,
123        };
124        // We have to make sure we don't accidentally free the original value, since its drop will change
125        // the borrow flag.
126        #[allow(clippy::mem_forget)]
127        mem::forget(orig);
128        res
129    }
130
131    /// See [`Ref.map_split`](Ref::map_split). Not a self method since that interferes with the
132    /// [`Deref`](Deref).
133    pub fn map_split<U: ?Sized, V: ?Sized, F>(orig: ARef<'a, T>, f: F) -> (ARef<'a, U>, ARef<'a, V>)
134    where
135        F: FnOnce(&T) -> (&U, &V),
136    {
137        if orig.borrow.is_none() {
138            let (a, b) = f(orig.value);
139            (ARef::new_ptr(a), ARef::new_ptr(b))
140        } else {
141            let orig: Ref<T> = unsafe { cast::transmute_unchecked(orig) };
142            let (a, b) = Ref::map_split(orig, f);
143            (ARef::new_ref(a), ARef::new_ref(b))
144        }
145    }
146
147    /// See [`Ref.filter_map`](Ref::filter_map). Not a self method since that interferes with the
148    /// [`Deref`](Deref).
149    pub fn filter_map<U: ?Sized, F>(orig: ARef<'a, T>, f: F) -> Result<ARef<'a, U>, Self>
150    where
151        F: FnOnce(&T) -> Option<&U>,
152    {
153        match f(orig.value) {
154            Some(value) => {
155                let res = Ok(ARef {
156                    value,
157                    borrow: orig.borrow,
158                });
159                // We have to make sure we don't accidentally free the original value, since its drop will change
160                // the borrow flag.
161                #[allow(clippy::mem_forget)]
162                mem::forget(orig);
163                res
164            }
165            None => Err(orig),
166        }
167    }
168}
169
170// `Ref` doesn't have many traits on it. I don't really know why - I think that's an oversight.
171// & references do have many traits on them. Therefore, when being "either" we choose to do as many
172// implementations as we can.
173
174impl<T: Display + ?Sized> Display for ARef<'_, T> {
175    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
176        ARef::deref(self).fmt(f)
177    }
178}
179
180impl<T: Hash + ?Sized> Hash for ARef<'_, T> {
181    fn hash<H: Hasher>(&self, state: &mut H) {
182        ARef::deref(self).hash(state)
183    }
184}
185
186impl<A: PartialEq<B> + ?Sized, B: ?Sized> PartialEq<ARef<'_, B>> for ARef<'_, A> {
187    fn eq(&self, other: &ARef<'_, B>) -> bool {
188        ARef::deref(self).eq(ARef::deref(other))
189    }
190}
191
192impl<A: Eq + ?Sized> Eq for ARef<'_, A> {}
193
194impl<A: PartialOrd<B> + ?Sized, B: ?Sized> PartialOrd<ARef<'_, B>> for ARef<'_, A> {
195    fn partial_cmp(&self, other: &ARef<'_, B>) -> Option<Ordering> {
196        ARef::deref(self).partial_cmp(ARef::deref(other))
197    }
198}
199
200impl<A: Ord + ?Sized> Ord for ARef<'_, A> {
201    fn cmp(&self, other: &Self) -> Ordering {
202        ARef::deref(self).cmp(ARef::deref(other))
203    }
204}
205
206/// Obtain an [`ARef`] from either a normal pointer or a [`RefCell`](std::cell::RefCell).
207pub trait AsARef<T: ?Sized> {
208    /// Get an [`ARef`] pointing at this type.
209    fn as_aref(this: &Self) -> ARef<T>;
210    /// Try and get an [`ARef`] pointing at this type. Returns an [`Err`] if
211    /// the type `Self` is a [`RefCell`] which is already mutably borrowed.
212    fn try_as_aref(this: &Self) -> Result<ARef<T>, BorrowError>;
213    /// Return the underlying [`RefCell`] if `Self` is one, otherwise [`None`].
214    fn as_ref_cell(this: &Self) -> Option<&RefCell<T>>;
215}
216
217impl<T: ?Sized> AsARef<T> for T {
218    fn as_aref(this: &Self) -> ARef<T> {
219        ARef::new_ptr(this)
220    }
221    fn try_as_aref(this: &Self) -> Result<ARef<T>, BorrowError> {
222        Ok(ARef::new_ptr(this))
223    }
224    fn as_ref_cell(_this: &Self) -> Option<&RefCell<T>> {
225        None
226    }
227}
228
229impl<T: ?Sized> AsARef<T> for RefCell<T> {
230    fn as_aref(this: &Self) -> ARef<T> {
231        ARef::new_ref(this.borrow())
232    }
233    fn try_as_aref(this: &Self) -> Result<ARef<T>, BorrowError> {
234        Ok(ARef::new_ref(this.try_borrow()?))
235    }
236    fn as_ref_cell(this: &Self) -> Option<&RefCell<T>> {
237        Some(this)
238    }
239}
240
241#[cfg(test)]
242mod tests {
243    use std::cell::RefCell;
244    use std::mem;
245
246    use super::*;
247    use crate::cast;
248
249    #[test]
250    fn test_from_ref_docs() {
251        let c = RefCell::new((5, 'b'));
252        let b1: ARef<(u32, char)> = ARef::new_ref(c.borrow());
253        let b2: ARef<u32> = ARef::map(b1, |t| &t.0);
254        assert_eq!(*b2, 5);
255
256        let cell = RefCell::new([1, 2, 3, 4]);
257        let borrow = ARef::new_ref(cell.borrow());
258        let (begin, end) = ARef::map_split(borrow, |slice| slice.split_at(2));
259        assert_eq!(*begin, [1, 2]);
260        assert_eq!(*end, [3, 4]);
261    }
262
263    #[test]
264    fn test_borrow_guards() {
265        let c = RefCell::new(5);
266        assert!(c.try_borrow_mut().is_ok());
267        let r1 = ARef::new_ref(c.borrow());
268        assert!(c.try_borrow_mut().is_err());
269        let r2 = c.borrow();
270        assert!(c.try_borrow_mut().is_err());
271        mem::drop(r1);
272        assert!(c.try_borrow_mut().is_err());
273        mem::drop(r2);
274        assert!(c.try_borrow_mut().is_ok());
275    }
276
277    #[test]
278    fn test_pointer_basics() {
279        let c = "test".to_owned();
280        let p = ARef::new_ptr(&c);
281        let p2 = ARef::map(p, |x| &x[1..3]);
282        assert_eq!(&*p2, "es");
283    }
284
285    #[test]
286    fn test_ref_map_dropping() {
287        let c = RefCell::new("test".to_owned());
288        let p = ARef::new_ref(c.borrow());
289        let p = ARef::map(p, |x| &x[1..3]);
290        assert_eq!(&*p, "es");
291        mem::drop(p);
292        assert!(c.try_borrow_mut().is_ok());
293    }
294
295    #[test]
296    fn test_ref_filter_map_dropping() {
297        let c = RefCell::new("test".to_owned());
298        let p = ARef::new_ref(c.borrow());
299        let p = ARef::filter_map(p, |x| Some(&x[1..3])).unwrap();
300        assert_eq!(&*p, "es");
301        mem::drop(p);
302        assert!(c.try_borrow_mut().is_ok());
303    }
304
305    #[test]
306    /// Test that the representation of ARef is what we expect
307    fn test_ref_as_expected() {
308        let orig = RefCell::new("test".to_owned());
309        let p = orig.borrow();
310        let p2 = Ref::clone(&p);
311        let (pointer, cell): (usize, usize) = unsafe { mem::transmute(p) };
312        // We expect the first to be a pointer to the underlying string
313        assert_eq!(pointer, cast::ptr_to_usize(Ref::deref(&p2)));
314        // We want to make sure the second is never zero
315        assert_ne!(cell, 0);
316
317        // Put it back as it was, to make sure our test doesn't leak memory
318        let _ignore: Ref<String> = unsafe { mem::transmute((pointer, cell)) };
319    }
320
321    #[test]
322    fn test_as_aref() {
323        fn get_str(x: &impl AsARef<String>) -> ARef<str> {
324            ARef::map(AsARef::as_aref(x), |x| x.as_str())
325        }
326
327        let a = RefCell::new("hello".to_owned());
328        let b = "world".to_owned();
329        assert_eq!(&*get_str(&a), "hello");
330        assert_eq!(&*get_str(&b), "world");
331    }
332}
gazebo/cell.rs

gazebo/
cell.rs
