sparse_mem/
lib.rs

1/*
2 * Copyright (c) Peter Bjorklund. All rights reserved. https://github.com/piot/sparse-mem
3 * Licensed under the MIT License. See LICENSE in the project root for license information.
4 */
5
6/// Memory layout for a fixed-capacity, generation-tracked sparse array in memory. Designed for the Swamp VM.
7///
8/// Follows Swamp VM Collection standard with capacity and element count first, 2-bytes each.
9/// Aligns to a 4-byte aligned values region, since it is mainly for a 32-bit VM:
10/// ```text
11/// offset 0:   capacity                (u16)
12/// offset 2:   element_count           (u16)
13/// offset 4:   element_size            (u32)
14/// offset 8:   slot_to_id              (u16[capacity])             // reverse lookup
15/// offset 8+2*capacity: id_to_slot     (u16[capacity])             // forward lookup
16/// offset 8+4*capacity: generation     (u16[capacity])             // bumped on alloc/remove
17/// offset 8+6*capacity:     pad (u8[pad])                          // pad to 4-byte boundary
18/// offset 8+6*capacity+pad: values (raw bytes)
19/// ```
20use std::ptr;
21
22/// Compute total bytes needed in memory for a sparse array. Used for code generator to know
23/// how much space to reserve.
24#[must_use]
25pub const fn layout_size(capacity: u16, element_size: u32) -> usize {
26    let cap = capacity as usize;
27    // slot_to_id + id_to_slot: each u16[capacity]
28    let lookup = 2 * cap * size_of::<u16>();
29    // generation: u16[capacity]
30    let generation_size = cap * size_of::<u16>();
31    // bytes before values
32    let before_vals = HEADER_SIZE + lookup + generation_size;
33    // pad to 8-byte alignment
34    let pad = (8 - (before_vals % 8)) % 8;
35    // values: raw bytes per element
36    let vals = cap * element_size as usize;
37    before_vals + pad + vals
38}
39
40/// Alignment requirement for the sparse array.
41/// Report 4 just for the benefit of values
42#[must_use]
43pub const fn alignment() -> usize {
44    8
45}
46
47const SLOT_OFFSET: usize = HEADER_SIZE;
48const HEADER_SIZE: usize = 8;
49
50/// Initialize the sparse array to memory specified by the raw memory pointer.
51/// `base` must point to a region of at least `layout_size(capacity, element_size)` bytes.
52pub unsafe fn init(base: *mut u8, capacity: u16, element_size: u32) {
53    unsafe {
54        ptr::write(base.cast::<u16>(), capacity);
55        ptr::write(base.add(2).cast::<u16>(), 0);
56        ptr::write(base.add(4).cast::<u32>(), element_size);
57        let cap = capacity as usize;
58        let id_offset = SLOT_OFFSET + cap * size_of::<u16>();
59        let generation_offset = id_offset + cap * size_of::<u16>();
60        for i in 0..cap {
61            ptr::write(
62                base.add(SLOT_OFFSET).cast::<u16>().add(i),
63                (cap - 1 - i) as u16,
64            );
65        }
66        for i in 0..cap {
67            ptr::write(base.add(id_offset).cast::<u16>().add(i), u16::MAX);
68        }
69        for i in 0..cap {
70            ptr::write(base.add(generation_offset).cast::<u16>().add(i), 0);
71        }
72    }
73}
74
75/// Allocate a new ID and generation.
76pub unsafe fn allocate(base: *mut u8) -> Option<(u16, u16)> {
77    unsafe {
78        let capacity = *base.cast::<u16>() as usize;
79        let count_ptr = base.add(2).cast::<u16>();
80        let count = *count_ptr as usize;
81        if count >= capacity {
82            return None;
83        }
84
85        let id_offset = SLOT_OFFSET + capacity * size_of::<u16>();
86        let generation_offset = id_offset + capacity * size_of::<u16>();
87
88        let id = *base.add(SLOT_OFFSET).cast::<u16>().add(count);
89        ptr::write(count_ptr, (count as u16).wrapping_add(1));
90        ptr::write(
91            base.add(id_offset).cast::<u16>().add(id as usize),
92            count as u16,
93        );
94
95        let generation_ptr = base.add(generation_offset).cast::<u16>();
96        let new_gen = generation_ptr.add(id as usize).read().wrapping_add(1);
97
98        ptr::write(generation_ptr.add(id as usize), new_gen);
99
100        Some((id, new_gen))
101    }
102}
103
104/// Compute offset of values region (values are aligned to 8 bytes)
105#[must_use]
106pub const fn values_offset(base: *const u8) -> usize {
107    let capacity = unsafe { *base.cast::<u16>() } as usize;
108    let id_offset = SLOT_OFFSET + capacity * size_of::<u16>();
109    let generation_offset = id_offset + capacity * size_of::<u16>();
110    let before_values = generation_offset + capacity * size_of::<u16>();
111    let padding = (8 - (before_values % 8)) % 8;
112    before_values + padding
113}
114
115/// Insert raw bytes at handle id
116/// # Safety
117///
118#[inline]
119pub unsafe fn insert(base: *mut u8, id: u16, src: *const u8) -> bool {
120    unsafe {
121        let capacity = *base.cast::<u16>() as usize;
122
123        // TODO: These should panic!()
124        #[cfg(debug_assertions)]
125        {
126            // BOUNDS CHECK: Ensure id is within capacity
127            if id as usize >= capacity {
128                return false;
129            }
130
131            // LIVENESS CHECK: Ensure the slot is actually allocated
132            // We need to check if this id has a valid slot assignment
133            let id_offset = SLOT_OFFSET + capacity * size_of::<u16>();
134            let slot = *base.add(id_offset).cast::<u16>().add(id as usize);
135            let count = *base.add(2).cast::<u16>() as usize;
136
137            // If slot is u16::MAX or >= count, the slot is not alive
138            if slot == u16::MAX || slot as usize >= count {
139                return false;
140            }
141
142            // Double-check: verify the slot_to_id mapping is consistent
143            let slot_to_id_ptr = base.add(SLOT_OFFSET).cast::<u16>();
144            if *slot_to_id_ptr.add(slot as usize) != id {
145                return false;
146            }
147        }
148
149        let element_size = *base.add(4).cast::<u32>() as usize;
150        let off = values_offset(base) + id as usize * element_size;
151        ptr::copy_nonoverlapping(src, base.add(off), element_size);
152        true
153    }
154}
155
156/// Remove by handle; frees slot and bumps generation. Returns false on mismatching generations.
157/// # Safety
158///
159pub unsafe fn remove(base: *mut u8, id: u16, generation: u16) -> bool {
160    unsafe {
161        let capacity = *base.cast::<u16>() as usize;
162
163        // TODO: These should panic!()
164        #[cfg(debug_assertions)]
165        {
166            // BOUNDS CHECK: Ensure id is within capacity
167            if id as usize >= capacity {
168                return false;
169            }
170
171            if !is_alive(base, id, generation) {
172                return false;
173            }
174        }
175
176        let id_offset = SLOT_OFFSET + capacity * size_of::<u16>();
177        let generation_offset = id_offset + capacity * size_of::<u16>();
178        let count_ptr = base.add(2).cast::<u16>();
179        let count = (*count_ptr) as usize;
180        let last = count - 1;
181        ptr::write(count_ptr, last as u16);
182
183        let slot_ptr = base.add(SLOT_OFFSET).cast::<u16>();
184        let slot_idx = *base.add(id_offset).cast::<u16>().add(id as usize) as usize;
185        let last_id = *slot_ptr.add(last);
186        ptr::write(slot_ptr.add(slot_idx), last_id);
187        ptr::write(
188            base.add(id_offset).cast::<u16>().add(last_id as usize),
189            slot_idx as u16,
190        );
191        ptr::write(slot_ptr.add(last), id);
192        ptr::write(base.add(id_offset).cast::<u16>().add(id as usize), u16::MAX);
193        let gen_ptr = base.add(generation_offset).cast::<u16>();
194        ptr::write(gen_ptr.add(id as usize), generation.wrapping_add(1));
195
196        true
197    }
198}
199
200/// Check handle validity
201/// # Safety
202///
203pub unsafe fn is_alive(base: *mut u8, id: u16, generation: u16) -> bool {
204    unsafe {
205        let capacity = *base.cast::<u16>() as usize;
206
207        // TODO: These should panic!()
208        #[cfg(debug_assertions)]
209        {
210            // BOUNDS CHECK: Ensure id is within capacity
211            if id as usize >= capacity {
212                return false;
213            }
214        }
215
216        let count = *base.add(2).cast::<u16>() as usize;
217        let id_offset = SLOT_OFFSET + capacity * size_of::<u16>();
218        let generation_offset = id_offset + capacity * size_of::<u16>();
219        // Not only check generation, but also check that it slot is among the live ones.
220        // Also makes extra check to look in the other direction as well.
221        let slot = *base.add(id_offset).cast::<u16>().add(id as usize) as usize;
222        let current_generation = *base.add(generation_offset).cast::<u16>().add(id as usize);
223        slot < count
224            && current_generation == generation
225            && *base.add(SLOT_OFFSET).cast::<u16>().add(slot) == id
226    }
227}
228
229/// Get a pointer to the `slot_to_id` array (reverse lookup)
230/// # Safety
231///
232pub const unsafe fn slot_to_id_ptr(base: *mut u8) -> *mut u16 {
233    unsafe { base.add(SLOT_OFFSET).cast::<u16>() }
234}
235
236#[must_use]
237pub const unsafe fn generation_ptr_const(base: *const u8) -> *const u16 {
238    unsafe {
239        let capacity = *base.cast::<u16>() as usize;
240        let id_offset = SLOT_OFFSET + capacity * size_of::<u16>();
241        let generation_offset = id_offset + capacity * size_of::<u16>();
242
243        base.add(generation_offset).cast::<u16>()
244    }
245}
246
247/// Get a pointer to the `slot_to_id` array (reverse lookup)
248/// # Safety
249///
250#[must_use]
251pub const unsafe fn slot_to_id_ptr_const(base: *const u8) -> *const u16 {
252    unsafe { base.add(SLOT_OFFSET).cast::<u16>() }
253}
254
255/// Get a pointer to the `id_to_slot` array (forward lookup)
256/// # Safety
257///
258pub unsafe fn id_to_slot_ptr(base: *mut u8) -> *mut u16 {
259    unsafe {
260        let capacity = *base.cast::<u16>() as usize;
261        base.add(SLOT_OFFSET + capacity * size_of::<u16>())
262            .cast::<u16>()
263    }
264}
265
266#[must_use]
267pub const unsafe fn id_to_slot_ptr_const(base: *const u8) -> *const u16 {
268    unsafe {
269        let capacity = *base.cast::<u16>() as usize;
270        base.add(SLOT_OFFSET + capacity * size_of::<u16>())
271            .cast::<u16>()
272    }
273}
274
275/// Get current element count
276/// # Safety
277///
278#[must_use]
279pub const unsafe fn element_count(base: *const u8) -> u16 {
280    unsafe { *base.add(2).cast::<u16>() }
281}
282
283/// Get current element size
284/// # Safety
285#[must_use]
286pub const unsafe fn element_size(base: *const u8) -> u32 {
287    unsafe { *base.add(4).cast::<u32>() }
288}
289
290/// Safe wrapper for insert that validates the handle
291/// Insert raw bytes at handle id, but only if the handle is valid
292/// Returns true if successful, false if the handle is invalid
293/// # Safety
294pub unsafe fn insert_if_alive(base: *mut u8, id: u16, generation: u16, src: *const u8) -> bool {
295    unsafe {
296        if is_alive(base, id, generation) {
297            insert(base, id, src)
298        } else {
299            false
300        }
301    }
302}
303
304/// Get value pointer for a valid handle
305/// Returns None if the handle is invalid
306/// # Safety
307pub unsafe fn get_value_ptr(base: *mut u8, id: u16, generation: u16) -> Option<*mut u8> {
308    unsafe {
309        if !is_alive(base, id, generation) {
310            return None;
311        }
312
313        let capacity = *base.cast::<u16>() as usize;
314        if id as usize >= capacity {
315            return None;
316        }
317
318        let element_size = *base.add(4).cast::<u32>() as usize;
319        let off = values_offset(base) + id as usize * element_size;
320        Some(base.add(off))
321    }
322}