1 //! An allocator definition for this embedding.
2 //!
3 //! The Rust standard library and Wasmtime require a memory allocator to be
4 //! configured. For custom embeddings of Wasmtime this might likely already be
5 //! defined elsewhere in the system in which case that should be used. This file
6 //! contains an example implementation using the Rust `dlmalloc` crate using
7 //! memory created by `wasmtime_*` platform symbols. This provides a file that
8 //! manages memory without any extra runtime dependencies, but this is just an
9 //! example.
10 //!
11 //! Allocators in Rust are configured with the `#[global_allocator]` attribute
12 //! and the `GlobalAlloc for T` trait impl. This should be used when hooking
13 //! up to an allocator elsewhere in the system.
14 
15 use alloc::alloc::{GlobalAlloc, Layout};
16 use core::cell::UnsafeCell;
17 use core::ops::{Deref, DerefMut};
18 use core::ptr;
19 use core::sync::atomic::{
20     AtomicBool,
21     Ordering::{Acquire, Release},
22 };
23 use dlmalloc::Dlmalloc;
24 
25 #[global_allocator]
26 static MALLOC: MyGlobalDmalloc = MyGlobalDmalloc {
27     dlmalloc: Mutex::new(Dlmalloc::new_with_allocator(MyAllocator)),
28 };
29 
30 struct MyGlobalDmalloc {
31     dlmalloc: Mutex<Dlmalloc<MyAllocator>>,
32 }
33 
34 struct MyAllocator;
35 
36 unsafe impl GlobalAlloc for MyGlobalDmalloc {
37     unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
38         self.dlmalloc
39             .try_lock()
40             .unwrap()
41             .malloc(layout.size(), layout.align())
42     }
43 
44     unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
45         self.dlmalloc
46             .try_lock()
47             .unwrap()
48             .calloc(layout.size(), layout.align())
49     }
50 
51     unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
52         self.dlmalloc
53             .try_lock()
54             .unwrap()
55             .realloc(ptr, layout.size(), layout.align(), new_size)
56     }
57 
58     unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
59         self.dlmalloc
60             .try_lock()
61             .unwrap()
62             .free(ptr, layout.size(), layout.align())
63     }
64 }
65 
66 const INITIAL_HEAP_SIZE: usize = 64 * 1024;
67 static mut INITIAL_HEAP: [u8; INITIAL_HEAP_SIZE] = [0; INITIAL_HEAP_SIZE];
68 static mut INITIAL_HEAP_ALLOCATED: bool = false;
69 
70 unsafe impl dlmalloc::Allocator for MyAllocator {
71     fn alloc(&self, _size: usize) -> (*mut u8, usize, u32) {
72         unsafe {
73             if INITIAL_HEAP_ALLOCATED {
74                 (ptr::null_mut(), 0, 0)
75             } else {
76                 INITIAL_HEAP_ALLOCATED = true;
77                 ((&raw mut INITIAL_HEAP).cast(), INITIAL_HEAP_SIZE, 0)
78             }
79         }
80     }
81 
82     fn remap(&self, _ptr: *mut u8, _old: usize, _new: usize, _can_move: bool) -> *mut u8 {
83         core::ptr::null_mut()
84     }
85 
86     fn free_part(&self, _ptr: *mut u8, _old: usize, _new: usize) -> bool {
87         false
88     }
89 
90     fn free(&self, _ptr: *mut u8, _size: usize) -> bool {
91         false
92     }
93 
94     fn can_release_part(&self, _flags: u32) -> bool {
95         false
96     }
97 
98     fn allocates_zeros(&self) -> bool {
99         true
100     }
101 
102     fn page_size(&self) -> usize {
103         4096
104     }
105 }
106 
107 // Simple mutex which only supports `try_lock` at this time. This would probably
108 // be replaced with a "real" mutex in a "real" embedding.
109 struct Mutex<T> {
110     data: UnsafeCell<T>,
111     locked: AtomicBool,
112 }
113 
114 unsafe impl<T: Send> Send for Mutex<T> {}
115 unsafe impl<T: Send> Sync for Mutex<T> {}
116 
117 impl<T> Mutex<T> {
118     const fn new(val: T) -> Mutex<T> {
119         Mutex {
120             data: UnsafeCell::new(val),
121             locked: AtomicBool::new(false),
122         }
123     }
124 
125     fn try_lock(&self) -> Option<impl DerefMut<Target = T> + '_> {
126         if self.locked.swap(true, Acquire) {
127             None
128         } else {
129             Some(MutexGuard { lock: self })
130         }
131     }
132 }
133 
134 struct MutexGuard<'a, T> {
135     lock: &'a Mutex<T>,
136 }
137 
138 impl<T> Deref for MutexGuard<'_, T> {
139     type Target = T;
140 
141     fn deref(&self) -> &T {
142         unsafe { &*self.lock.data.get() }
143     }
144 }
145 
146 impl<T> DerefMut for MutexGuard<'_, T> {
147     fn deref_mut(&mut self) -> &mut T {
148         unsafe { &mut *self.lock.data.get() }
149     }
150 }
151 
152 impl<T> Drop for MutexGuard<'_, T> {
153     fn drop(&mut self) {
154         self.lock.locked.store(false, Release);
155     }
156 }
157