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
98unsafe 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
113unsafe 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
120unsafe fn btreemap_default<K: 'static, V: 'static>(ox: OxPtrUninit) -> bool {
122 unsafe { ox.put(BTreeMap::<K, V>::new()) };
123 true
124}
125
126unsafe 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 .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}