facet_core/impls_alloc/
btreeset.rs

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