frozen_hashbrown/
frozen.rs

1use core::{alloc::Layout, ptr::NonNull};
2use std::fmt::Debug;
3
4pub const RANDOM_STATE_TYPE_NAME: &str = "std::collections::hash::map::RandomState";
5pub const GLOBAL_ALLOC_TYPE_NAME: &str = "alloc::alloc::Global";
6
7#[derive(Clone)]
8pub struct FrozenHashMap<S = RandomState> {
9    pub table_layout: TableLayout,
10    pub hashmap: HashMap<S>,
11    pub memory: Vec<u8>,
12}
13
14impl<S: Debug> Debug for FrozenHashMap<S> {
15    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
16        f.debug_struct("FrozenHashMap")
17            .field("table_layout", &self.table_layout)
18            .field("hashmap", &self.hashmap)
19            .field(
20                "memory",
21                &format!("<binary data of size {}>", self.memory.len()),
22            )
23            .finish()
24    }
25}
26
27#[derive(Debug, Clone)]
28pub struct HashMap<S = RandomState> {
29    pub hash_builder: S,
30    pub table: RawTable,
31}
32
33#[derive(Debug, Clone)]
34pub struct RandomState {
35    pub k0: u64,
36    pub k1: u64,
37}
38
39#[derive(Debug, Clone)]
40pub struct RawTable {
41    pub table: RawTableInner,
42}
43
44#[derive(Debug, Clone)]
45pub struct RawTableInner {
46    pub bucket_mask: usize,
47    pub ctrl: NonNull<u8>,
48    pub growth_left: usize,
49    pub items: usize,
50}
51
52#[derive(Debug, Copy, Clone)]
53pub struct TableLayout {
54    pub size: usize,
55    pub ctrl_align: usize,
56}
57
58impl TableLayout {
59    pub fn new(layout: Layout) -> Self {
60        Self {
61            size: layout.size(),
62            ctrl_align: if layout.align() > crate::Group::WIDTH {
63                layout.align()
64            } else {
65                crate::Group::WIDTH
66            },
67        }
68    }
69
70    pub fn calculate_layout_for(&self, buckets: usize) -> Option<(Layout, usize)> {
71        assert!(buckets.is_power_of_two());
72
73        let TableLayout { size, ctrl_align } = *self;
74        // Manual layout calculation since Layout methods are not yet stable.
75        let ctrl_offset =
76            size.checked_mul(buckets)?.checked_add(ctrl_align - 1)? & !(ctrl_align - 1);
77        let len = ctrl_offset.checked_add(buckets + crate::Group::WIDTH)?;
78
79        Some((
80            unsafe { Layout::from_size_align_unchecked(len, ctrl_align) },
81            ctrl_offset,
82        ))
83    }
84}
85
86impl RawTableInner {
87    pub fn allocation(&self, table_layout: &TableLayout) -> Option<(*const u8, Layout)> {
88        if self.is_empty_singleton() {
89            None
90        } else {
91            let (layout, ctrl_offset) = table_layout.calculate_layout_for(self.buckets())?;
92            Some((unsafe { self.ctrl.as_ptr().sub(ctrl_offset) }, layout))
93        }
94    }
95
96    pub fn reallocation(&self, table_layout: &TableLayout) -> Option<(usize, Layout)> {
97        if self.is_empty_singleton() {
98            None
99        } else {
100            let (layout, ctrl_offset) = table_layout.calculate_layout_for(self.buckets())?;
101            Some((ctrl_offset, layout))
102        }
103    }
104
105    fn buckets(&self) -> usize {
106        self.bucket_mask + 1
107    }
108
109    fn is_empty_singleton(&self) -> bool {
110        self.bucket_mask == 0
111    }
112}
113
114impl FrozenHashMap<RandomState> {
115    pub fn construct<K, V>(hashmap: &std::collections::HashMap<K, V>) -> Self {
116        Self::construct_with(
117            unsafe {
118                core::slice::from_raw_parts(
119                    std::mem::transmute(hashmap as *const _),
120                    std::mem::size_of::<std::collections::HashMap<K, V>>(),
121                )
122            },
123            TableLayout::new(Layout::new::<(K, V)>()),
124        )
125    }
126
127    pub fn construct_with(hashmap: &[u8], table_layout: TableLayout) -> Self {
128        assert_eq!(std::mem::size_of::<HashMap<RandomState>>(), hashmap.len());
129        let hashmap: HashMap<RandomState> =
130            unsafe { std::ptr::read_unaligned(hashmap.as_ptr() as *const _) };
131        let memory = if let Some((location, layout)) = hashmap.table.table.allocation(&table_layout)
132        {
133            let location: &[u8] =
134                unsafe { core::slice::from_raw_parts(location as *const u8, layout.size()) };
135            location.to_vec()
136        } else {
137            vec![]
138        };
139        Self {
140            table_layout,
141            hashmap,
142            memory,
143        }
144    }
145
146    pub fn reconstruct<K, V>(&mut self) -> Option<&std::collections::HashMap<K, V>> {
147        assert_eq!(
148            std::mem::size_of::<HashMap<RandomState>>(),
149            std::mem::size_of::<std::collections::HashMap<K, V>>()
150        );
151        if self.memory.is_empty() {
152            return None;
153        }
154        if let Some((offset, layout)) = self.hashmap.table.table.reallocation(&self.table_layout) {
155            assert_eq!(layout.size(), self.memory.len());
156            let address = self.memory.as_ptr() as usize + offset;
157            if address == 0 {
158                return None;
159            }
160            self.hashmap.table.table.ctrl = unsafe { NonNull::new_unchecked(address as *mut u8) };
161            unsafe {
162                // this is the crazy part
163                Some(std::mem::transmute(&self.hashmap))
164            }
165        } else {
166            None
167        }
168    }
169
170    pub fn store(&self) -> Vec<u8> {
171        let mut bytes = Vec::new();
172        bytes.extend_from_slice(unsafe {
173            core::slice::from_raw_parts(
174                std::mem::transmute(&self.table_layout as *const _),
175                std::mem::size_of::<TableLayout>(),
176            )
177        });
178        bytes.extend_from_slice(unsafe {
179            core::slice::from_raw_parts(
180                std::mem::transmute(&self.hashmap as *const _),
181                std::mem::size_of::<HashMap<RandomState>>(),
182            )
183        });
184        bytes.extend_from_slice(&self.memory.len().to_ne_bytes());
185        bytes.extend_from_slice(&self.memory);
186        bytes
187    }
188
189    /// None means failed to load
190    pub fn load(bytes: &[u8]) -> Option<Self> {
191        let mut cursor = 0;
192        let chunk = std::mem::size_of::<TableLayout>();
193        if cursor + chunk > bytes.len() {
194            return None;
195        }
196        let table_layout: TableLayout =
197            unsafe { std::ptr::read_unaligned(bytes.as_ptr() as *const _) };
198        cursor += chunk;
199        let chunk = std::mem::size_of::<HashMap<RandomState>>();
200        if cursor + chunk > bytes.len() {
201            return None;
202        }
203        let hashmap: HashMap<RandomState> =
204            unsafe { std::ptr::read_unaligned(bytes.as_ptr().add(cursor) as *const _) };
205        cursor += chunk;
206        let chunk = 8;
207        if cursor + chunk > bytes.len() {
208            return None;
209        }
210        let ll = [
211            bytes[cursor],
212            bytes[cursor + 1],
213            bytes[cursor + 2],
214            bytes[cursor + 3],
215            bytes[cursor + 4],
216            bytes[cursor + 5],
217            bytes[cursor + 6],
218            bytes[cursor + 7],
219        ];
220        cursor += chunk;
221        let length = usize::from_ne_bytes(ll);
222        if cursor + length != bytes.len() {
223            return None;
224        }
225        let memory = bytes[cursor..].to_vec();
226        Some(Self {
227            table_layout,
228            hashmap,
229            memory,
230        })
231    }
232
233    pub fn len(&self) -> usize {
234        self.hashmap.len()
235    }
236}
237
238impl<S> HashMap<S> {
239    pub fn len(&self) -> usize {
240        self.table.table.items
241    }
242}
frozen_hashbrown/frozen.rs

frozen_hashbrown/
frozen.rs
