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