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