facet_core/impls_alloc/
btreeset.rs

1use core::ptr::NonNull;
2
3use alloc::boxed::Box;
4use alloc::collections::BTreeSet;
5
6use crate::ptr::{PtrConst, PtrMut};
7
8use crate::{
9    Def, Facet, IterVTable, MarkerTraits, SetDef, SetVTable, Shape, Type, TypeParam, UserType,
10    ValueVTable,
11};
12
13type BTreeSetIterator<'mem, T> = alloc::collections::btree_set::Iter<'mem, T>;
14
15unsafe impl<'a, T> Facet<'a> for BTreeSet<T>
16where
17    T: Facet<'a> + core::cmp::Eq + core::cmp::Ord,
18{
19    const SHAPE: &'static Shape = &const {
20        Shape::builder_for_sized::<Self>()
21            .vtable({
22                ValueVTable::builder::<Self>()
23                    .marker_traits({
24                        MarkerTraits::SEND
25                            .union(MarkerTraits::SYNC)
26                            .union(MarkerTraits::EQ)
27                            .union(MarkerTraits::UNPIN)
28                            .intersection(T::SHAPE.vtable.marker_traits())
29                    })
30                    .type_name(|f, opts| {
31                        write!(f, "{}<", Self::SHAPE.type_identifier)?;
32                        if let Some(opts) = opts.for_children() {
33                            T::SHAPE.vtable.type_name()(f, opts)?;
34                        } else {
35                            write!(f, "…")?;
36                        }
37                        write!(f, ">")
38                    })
39                    .default_in_place({
40                        Some(|target| unsafe { target.put(Self::default()).into() })
41                    })
42                    .build()
43            })
44            .type_identifier("BTreeSet")
45            .type_params(&[TypeParam {
46                name: "T",
47                shape: T::SHAPE,
48            }])
49            .ty(Type::User(UserType::Opaque))
50            .def(Def::Set(
51                SetDef::builder()
52                    .t(T::SHAPE)
53                    .vtable(
54                        &const {
55                            SetVTable::builder()
56                                .init_in_place_with_capacity(|uninit, _capacity| unsafe {
57                                    uninit.put(Self::new())
58                                })
59                                .insert(|ptr, item| unsafe {
60                                    let set = ptr.as_mut::<BTreeSet<T>>();
61                                    let item = item.read::<T>();
62                                    set.insert(item)
63                                })
64                                .len(|ptr| unsafe {
65                                    let set = ptr.get::<BTreeSet<T>>();
66                                    set.len()
67                                })
68                                .contains(|ptr, item| unsafe {
69                                    let set = ptr.get::<BTreeSet<T>>();
70                                    set.contains(item.get())
71                                })
72                                .iter_vtable(
73                                    IterVTable::builder()
74                                        .init_with_value(|ptr| {
75                                            let set = unsafe { ptr.get::<BTreeSet<T>>() };
76                                            let iter: BTreeSetIterator<'_, T> = set.iter();
77                                            let iter_state = Box::new(iter);
78                                            PtrMut::new(unsafe {
79                                                NonNull::new_unchecked(
80                                                    Box::into_raw(iter_state) as *mut u8
81                                                )
82                                            })
83                                        })
84                                        .next(|iter_ptr| {
85                                            let state = unsafe {
86                                                iter_ptr.as_mut::<BTreeSetIterator<'_, T>>()
87                                            };
88                                            state
89                                                .next()
90                                                .map(|value| PtrConst::new(NonNull::from(value)))
91                                        })
92                                        .next_back(|iter_ptr| {
93                                            let state = unsafe {
94                                                iter_ptr.as_mut::<BTreeSetIterator<'_, T>>()
95                                            };
96                                            state
97                                                .next_back()
98                                                .map(|value| PtrConst::new(NonNull::from(value)))
99                                        })
100                                        .dealloc(|iter_ptr| unsafe {
101                                            drop(Box::from_raw(
102                                                iter_ptr.as_ptr::<BTreeSetIterator<'_, T>>()
103                                                    as *mut BTreeSetIterator<'_, T>,
104                                            ));
105                                        })
106                                        .build(),
107                                )
108                                .build()
109                        },
110                    )
111                    .build(),
112            ))
113            .build()
114    };
115}
116
117#[cfg(test)]
118mod tests {
119    use alloc::collections::BTreeSet;
120    use alloc::string::String;
121    use alloc::vec::Vec;
122
123    use super::*;
124
125    #[test]
126    fn test_btreesetset_type_params() {
127        let [type_param_1] = <BTreeSet<i32>>::SHAPE.type_params else {
128            panic!("BTreeSet<T> should have 1 type param")
129        };
130        assert_eq!(type_param_1.shape(), i32::SHAPE);
131    }
132
133    #[test]
134    fn test_btreeset_vtable_1_new_insert_iter_drop() {
135        facet_testhelpers::setup();
136
137        let btreeset_shape = <BTreeSet<String>>::SHAPE;
138        let btreeset_def = btreeset_shape
139            .def
140            .into_set()
141            .expect("BTreeSet<T> should have a set definition");
142
143        // Allocate memory for the BTreeSet
144        let btreeset_uninit_ptr = btreeset_shape.allocate().unwrap();
145
146        // Create the BTreeSet
147        let btreeset_ptr =
148            unsafe { (btreeset_def.vtable.init_in_place_with_capacity_fn)(btreeset_uninit_ptr, 0) };
149
150        // The BTreeSet is empty, so ensure its length is 0
151        let btreeset_actual_length =
152            unsafe { (btreeset_def.vtable.len_fn)(btreeset_ptr.as_const()) };
153        assert_eq!(btreeset_actual_length, 0);
154
155        // 5 sample values to insert
156        let strings = ["foo", "bar", "bazz", "fizzbuzz", "fifth thing"];
157
158        // Insert the 5 values into the BTreeSet
159        let mut btreeset_length = 0;
160        for string in strings {
161            // Create the value
162            let mut new_value = string.to_string();
163
164            // Insert the value
165            let did_insert = unsafe {
166                (btreeset_def.vtable.insert_fn)(
167                    btreeset_ptr,
168                    PtrMut::new(NonNull::from(&mut new_value)),
169                )
170            };
171
172            // The value now belongs to the BTreeSet, so forget it
173            core::mem::forget(new_value);
174
175            assert!(did_insert, "expected value to be inserted in the BTreeSet");
176
177            // Ensure the BTreeSet's length increased by 1
178            btreeset_length += 1;
179            let btreeset_actual_length =
180                unsafe { (btreeset_def.vtable.len_fn)(btreeset_ptr.as_const()) };
181            assert_eq!(btreeset_actual_length, btreeset_length);
182        }
183
184        // Insert the same 5 values again, ensuring they are deduplicated
185        for string in strings {
186            // Create the value
187            let mut new_value = string.to_string();
188
189            // Try to insert the value
190            let did_insert = unsafe {
191                (btreeset_def.vtable.insert_fn)(
192                    btreeset_ptr,
193                    PtrMut::new(NonNull::from(&mut new_value)),
194                )
195            };
196
197            // The value now belongs to the BTreeSet, so forget it
198            core::mem::forget(new_value);
199
200            assert!(
201                !did_insert,
202                "expected value to not be inserted in the BTreeSet"
203            );
204
205            // Ensure the BTreeSet's length did not increase
206            let btreeset_actual_length =
207                unsafe { (btreeset_def.vtable.len_fn)(btreeset_ptr.as_const()) };
208            assert_eq!(btreeset_actual_length, btreeset_length);
209        }
210
211        // Create a new iterator over the BTreeSet
212        let iter_init_with_value_fn = btreeset_def.vtable.iter_vtable.init_with_value.unwrap();
213        let btreeset_iter_ptr = unsafe { iter_init_with_value_fn(btreeset_ptr.as_const()) };
214
215        // Collect all the items from the BTreeSet's iterator
216        let mut iter_items = Vec::<&str>::new();
217        loop {
218            // Get the next item from the iterator
219            let item_ptr = unsafe { (btreeset_def.vtable.iter_vtable.next)(btreeset_iter_ptr) };
220            let Some(item_ptr) = item_ptr else {
221                break;
222            };
223
224            let item = unsafe { item_ptr.get::<String>() };
225
226            // Add the item into the list of items returned from the iterator
227            iter_items.push(&**item);
228        }
229
230        // Deallocate the iterator
231        unsafe {
232            (btreeset_def.vtable.iter_vtable.dealloc)(btreeset_iter_ptr);
233        }
234
235        // BTrees iterate in sorted order, so ensure the iterator returned
236        // each item in order
237        let mut strings_sorted = strings.to_vec();
238        strings_sorted.sort();
239        assert_eq!(iter_items, strings_sorted);
240
241        // Get the function pointer for dropping the BTreeSet
242        let drop_fn = btreeset_shape
243            .vtable
244            .drop_in_place
245            .expect("BTreeSet<T> should have drop_in_place");
246
247        // Drop the BTreeSet in place
248        unsafe { drop_fn(btreeset_ptr) };
249
250        // Deallocate the memory
251        unsafe { btreeset_shape.deallocate_mut(btreeset_ptr).unwrap() };
252    }
253}