Skip to main content

miden_sdk_alloc/
lib.rs

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