Skip to main content

facet_core/impls/alloc/
btreemap.rs

1use alloc::{boxed::Box, collections::BTreeMap};
2
3use crate::{
4    Def, Facet, IterVTable, MapDef, MapVTable, OxPtrMut, OxPtrUninit, PtrConst, PtrMut, PtrUninit,
5    Shape, ShapeBuilder, TypeNameFn, TypeNameOpts, TypeOpsIndirect, TypeParam, VTableIndirect,
6    Variance, VarianceDep, VarianceDesc,
7};
8
9type BTreeMapIterator<'mem, K, V> = alloc::collections::btree_map::Iter<'mem, K, V>;
10
11unsafe extern "C" fn btreemap_init_in_place_with_capacity<K, V>(
12    uninit: PtrUninit,
13    _capacity: usize,
14) -> PtrMut {
15    unsafe { uninit.put(BTreeMap::<K, V>::new()) }
16}
17
18unsafe extern "C" fn btreemap_insert<K: Eq + Ord + 'static, V: 'static>(
19    ptr: PtrMut,
20    key: PtrMut,
21    value: PtrMut,
22) {
23    unsafe {
24        let map = ptr.as_mut::<BTreeMap<K, V>>();
25        let k = key.read::<K>();
26        let v = value.read::<V>();
27        map.insert(k, v);
28    }
29}
30
31unsafe extern "C" fn btreemap_len<K: 'static, V: 'static>(ptr: PtrConst) -> usize {
32    unsafe { ptr.get::<BTreeMap<K, V>>().len() }
33}
34
35unsafe extern "C" fn btreemap_contains_key<K: Eq + Ord + 'static, V: 'static>(
36    ptr: PtrConst,
37    key: PtrConst,
38) -> bool {
39    unsafe { ptr.get::<BTreeMap<K, V>>().contains_key(key.get()) }
40}
41
42unsafe extern "C" fn btreemap_get_value_ptr<K: Eq + Ord + 'static, V: 'static>(
43    ptr: PtrConst,
44    key: PtrConst,
45) -> *const u8 {
46    unsafe {
47        ptr.get::<BTreeMap<K, V>>()
48            .get(key.get())
49            .map_or(core::ptr::null(), |v| v as *const V as *const u8)
50    }
51}
52
53unsafe extern "C" fn btreemap_iter_init<K: 'static, V: 'static>(ptr: PtrConst) -> PtrMut {
54    unsafe {
55        let map = ptr.get::<BTreeMap<K, V>>();
56        let iter: BTreeMapIterator<'_, K, V> = map.iter();
57        let state = Box::new(iter);
58        PtrMut::new(Box::into_raw(state) as *mut u8)
59    }
60}
61
62unsafe fn btreemap_iter_next<K: 'static, V: 'static>(
63    iter_ptr: PtrMut,
64) -> Option<(PtrConst, PtrConst)> {
65    unsafe {
66        let state = iter_ptr.as_mut::<BTreeMapIterator<'static, K, V>>();
67        state.next().map(|(key, value)| {
68            (
69                PtrConst::new(key as *const K),
70                PtrConst::new(value as *const V),
71            )
72        })
73    }
74}
75
76unsafe fn btreemap_iter_next_back<K: 'static, V: 'static>(
77    iter_ptr: PtrMut,
78) -> Option<(PtrConst, PtrConst)> {
79    unsafe {
80        let state = iter_ptr.as_mut::<BTreeMapIterator<'static, K, V>>();
81        state.next_back().map(|(key, value)| {
82            (
83                PtrConst::new(key as *const K),
84                PtrConst::new(value as *const V),
85            )
86        })
87    }
88}
89
90unsafe extern "C" fn btreemap_iter_dealloc<K, V>(iter_ptr: PtrMut) {
91    unsafe {
92        drop(Box::from_raw(
93            iter_ptr.as_ptr::<BTreeMapIterator<'_, K, V>>() as *mut BTreeMapIterator<'_, K, V>,
94        ))
95    }
96}
97
98/// Build a BTreeMap from a contiguous slice of (K, V) pairs.
99unsafe extern "C" fn btreemap_from_pair_slice<K: Eq + Ord + 'static, V: 'static>(
100    uninit: PtrUninit,
101    pairs_ptr: *mut u8,
102    count: usize,
103) -> PtrMut {
104    let pairs = pairs_ptr as *mut (K, V);
105    let iter = (0..count).map(|i| unsafe {
106        let pair_ptr = pairs.add(i);
107        core::ptr::read(pair_ptr)
108    });
109    let map: BTreeMap<K, V> = iter.collect();
110    unsafe { uninit.put(map) }
111}
112
113/// Drop for BTreeMap<K, V>
114unsafe fn btreemap_drop<K: 'static, V: 'static>(ox: OxPtrMut) {
115    unsafe {
116        core::ptr::drop_in_place(ox.as_mut::<BTreeMap<K, V>>());
117    }
118}
119
120/// Default for BTreeMap<K, V>
121unsafe fn btreemap_default<K: 'static, V: 'static>(ox: OxPtrUninit) -> bool {
122    unsafe { ox.put(BTreeMap::<K, V>::new()) };
123    true
124}
125
126// TODO: Debug, Hash, PartialEq, Eq, PartialOrd, Ord, for BTreeMap, BTreeSet
127unsafe impl<'a, K, V> Facet<'a> for BTreeMap<K, V>
128where
129    K: Facet<'a> + core::cmp::Eq + core::cmp::Ord + 'static,
130    V: Facet<'a> + 'static,
131{
132    const SHAPE: &'static crate::Shape = &const {
133        const fn build_map_vtable<K: Eq + Ord + 'static, V: 'static>() -> MapVTable {
134            MapVTable::builder()
135                .init_in_place_with_capacity(btreemap_init_in_place_with_capacity::<K, V>)
136                .insert(btreemap_insert::<K, V>)
137                .len(btreemap_len::<K, V>)
138                .contains_key(btreemap_contains_key::<K, V>)
139                .get_value_ptr(btreemap_get_value_ptr::<K, V>)
140                .iter_vtable(IterVTable {
141                    init_with_value: Some(btreemap_iter_init::<K, V>),
142                    next: btreemap_iter_next::<K, V>,
143                    next_back: Some(btreemap_iter_next_back::<K, V>),
144                    size_hint: None,
145                    dealloc: btreemap_iter_dealloc::<K, V>,
146                })
147                .from_pair_slice(Some(btreemap_from_pair_slice::<K, V>))
148                .pair_stride(core::mem::size_of::<(K, V)>())
149                .value_offset_in_pair(core::mem::offset_of!((K, V), 1))
150                .build()
151        }
152
153        const VTABLE: VTableIndirect = VTableIndirect::EMPTY;
154
155        const fn build_type_name<'a, K: Facet<'a>, V: Facet<'a>>() -> TypeNameFn {
156            fn type_name_impl<'a, K: Facet<'a>, V: Facet<'a>>(
157                _shape: &'static Shape,
158                f: &mut core::fmt::Formatter<'_>,
159                opts: TypeNameOpts,
160            ) -> core::fmt::Result {
161                write!(f, "BTreeMap")?;
162                if let Some(opts) = opts.for_children() {
163                    write!(f, "<")?;
164                    K::SHAPE.write_type_name(f, opts)?;
165                    write!(f, ", ")?;
166                    V::SHAPE.write_type_name(f, opts)?;
167                    write!(f, ">")?;
168                } else {
169                    write!(f, "<…>")?;
170                }
171                Ok(())
172            }
173            type_name_impl::<K, V>
174        }
175
176        ShapeBuilder::for_sized::<Self>("BTreeMap")
177            .module_path("alloc::collections::btree_map")
178            .type_name(build_type_name::<K, V>())
179            .vtable_indirect(&VTABLE)
180            .def(Def::Map(MapDef::new(
181                &const { build_map_vtable::<K, V>() },
182                K::SHAPE,
183                V::SHAPE,
184            )))
185            .type_params(&[
186                TypeParam {
187                    name: "K",
188                    shape: K::SHAPE,
189                },
190                TypeParam {
191                    name: "V",
192                    shape: V::SHAPE,
193                },
194            ])
195            // BTreeMap<K, V> combines K and V variances
196            .variance(VarianceDesc {
197                base: Variance::Bivariant,
198                deps: &const {
199                    [
200                        VarianceDep::covariant(K::SHAPE),
201                        VarianceDep::covariant(V::SHAPE),
202                    ]
203                },
204            })
205            .type_ops_indirect(
206                &const {
207                    TypeOpsIndirect {
208                        drop_in_place: btreemap_drop::<K, V>,
209                        default_in_place: Some(btreemap_default::<K, V>),
210                        clone_into: None,
211                        is_truthy: None,
212                    }
213                },
214            )
215            .build()
216    };
217}