miden_sdk_alloc/
lib.rs

1#![no_std]
2#![deny(warnings)]
3
4extern crate alloc;
5
6use alloc::alloc::{GlobalAlloc, Layout};
7use core::{
8    ptr::null_mut,
9    sync::atomic::{AtomicPtr, Ordering},
10};
11
12/// We assume the Wasm page size for purposes of initializing the heap
13#[cfg(target_family = "wasm")]
14const PAGE_SIZE: usize = 2usize.pow(16);
15
16/// We require all allocations to be minimally word-aligned, i.e. 32 byte alignment
17const MIN_ALIGN: usize = 32;
18
19/// The linear memory heap must not spill over into the region reserved for procedure
20/// locals, which begins at 2^30 in Miden's address space.
21const HEAP_END: *mut u8 = (2usize.pow(30) / 4) as *mut u8;
22
23/// A very simple allocator for Miden SDK-based programs.
24///
25/// This allocator does not free memory, it simply grows the heap until it runs out of available
26/// space for further allocations.
27pub struct BumpAlloc {
28    /// The address at which the available heap begins
29    top: AtomicPtr<u8>,
30}
31
32impl Default for BumpAlloc {
33    fn default() -> Self {
34        Self::new()
35    }
36}
37
38impl BumpAlloc {
39    /// Create a new instance of this allocator
40    ///
41    /// NOTE: Only one instance of this allocator should ever be used at a time, as it is
42    /// allocating from the global heap, not from memory reserved for itself.
43    pub const fn new() -> Self {
44        Self {
45            top: AtomicPtr::new(null_mut()),
46        }
47    }
48
49    /// Initialize the allocator, if it has not yet been initialized
50    #[cfg(target_family = "wasm")]
51    fn maybe_init(&self) {
52        let top = self.top.load(Ordering::Relaxed);
53        if top.is_null() {
54            let base = unsafe { heap_base() };
55            let size = core::arch::wasm32::memory_size(0);
56            self.top.store(unsafe { base.byte_add(size * PAGE_SIZE) }, Ordering::Relaxed);
57        }
58        // TODO: Once treeify issue is fixed, switch to this implementation
59        /*
60        let _ = self.top.fetch_update(Ordering::Relaxed, Ordering::Relaxed, |top| {
61            if top.is_null() {
62                let base = unsafe { heap_base() };
63                let size = core::arch::wasm32::memory_size(0);
64                Some(unsafe { base.byte_add(size * PAGE_SIZE) })
65            } else {
66                None
67            }
68        });
69        */
70    }
71
72    #[cfg(not(target_family = "wasm"))]
73    fn maybe_init(&self) {}
74}
75
76unsafe impl GlobalAlloc for BumpAlloc {
77    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
78        // Force allocations to be at minimally word-aligned. This is wasteful of memory, but
79        // we don't need to be particularly conservative with memory anyway, as most, if not all,
80        // Miden programs will be relatively short-lived. This makes interop at the Rust/Miden
81        // call boundary less expensive, as we can typically pass pointers directly to Miden,
82        // whereas without this alignment guarantee, we would have to set up temporary buffers for
83        // Miden code to write to, and then copy out of that buffer to whatever Rust type, e.g.
84        // `Vec`, we actually want.
85        //
86        // NOTE: This cannot fail, because we're always meeting minimum alignment requirements
87        let layout = layout
88            .align_to(core::cmp::max(layout.align(), MIN_ALIGN))
89            .unwrap()
90            .pad_to_align();
91        let size = layout.size();
92        let align = layout.align();
93
94        self.maybe_init();
95
96        let top = self.top.load(Ordering::Relaxed);
97        let available = HEAP_END.byte_offset_from(top) as usize;
98        if available >= size {
99            self.top.store(top.byte_add(size), Ordering::Relaxed);
100            unsafe { top.byte_offset(align as isize) }
101        } else {
102            null_mut()
103        }
104
105        // TODO: Once treeify issue is fixed, switch to this implementation
106        /*
107        match self.top.fetch_update(Ordering::Relaxed, Ordering::Relaxed, |top| {
108            let available = HEAP_END.byte_offset_from(top) as usize;
109            if available < size {
110                None
111            } else {
112                Some(top.byte_add(size))
113            }
114        }) {
115            Ok(prev_top) => {
116                unsafe { prev_top.byte_offset(align as isize) }
117            }
118            Err(_) => null_mut(),
119        }
120         */
121    }
122
123    unsafe fn dealloc(&self, _ptr: *mut u8, _layout: Layout) {}
124}
125
126#[cfg(target_family = "wasm")]
127#[link(wasm_import_module = "miden:core-intrinsics/intrinsics-mem@1.0.0")]
128extern "C" {
129    #[link_name = "heap-base"]
130    fn heap_base() -> *mut u8;
131}