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 // Hand-copied from `crates/wasmtime/src/runtime/vm/sys/custom/capi.rs`.
67 const PROT_READ: u32 = 1 << 0;
68 const PROT_WRITE: u32 = 1 << 1;
69 extern "C" {
70     fn wasmtime_mmap_new(size: usize, prot_flags: u32, ret: &mut *mut u8) -> i32;
71     fn wasmtime_page_size() -> usize;
72     fn wasmtime_munmap(ptr: *mut u8, size: usize) -> i32;
73 }
74 
75 unsafe impl dlmalloc::Allocator for MyAllocator {
76     fn alloc(&self, size: usize) -> (*mut u8, usize, u32) {
77         unsafe {
78             let mut ptr = ptr::null_mut();
79             let rc = wasmtime_mmap_new(size, PROT_READ | PROT_WRITE, &mut ptr);
80             if rc != 0 {
81                 (ptr::null_mut(), 0, 0)
82             } else {
83                 (ptr, size, 0)
84             }
85         }
86     }
87 
88     fn remap(&self, _ptr: *mut u8, _old: usize, _new: usize, _can_move: bool) -> *mut u8 {
89         core::ptr::null_mut()
90     }
91 
92     fn free_part(&self, _ptr: *mut u8, _old: usize, _new: usize) -> bool {
93         false
94     }
95 
96     fn free(&self, ptr: *mut u8, size: usize) -> bool {
97         unsafe {
98             wasmtime_munmap(ptr, size);
99             true
100         }
101     }
102 
103     fn can_release_part(&self, _flags: u32) -> bool {
104         false
105     }
106 
107     fn allocates_zeros(&self) -> bool {
108         true
109     }
110 
111     fn page_size(&self) -> usize {
112         unsafe { wasmtime_page_size() }
113     }
114 }
115 
116 // Simple mutex which only supports `try_lock` at this time. This would probably
117 // be replaced with a "real" mutex in a "real" embedding.
118 struct Mutex<T> {
119     data: UnsafeCell<T>,
120     locked: AtomicBool,
121 }
122 
123 unsafe impl<T: Send> Send for Mutex<T> {}
124 unsafe impl<T: Send> Sync for Mutex<T> {}
125 
126 impl<T> Mutex<T> {
127     const fn new(val: T) -> Mutex<T> {
128         Mutex {
129             data: UnsafeCell::new(val),
130             locked: AtomicBool::new(false),
131         }
132     }
133 
134     fn try_lock(&self) -> Option<impl DerefMut<Target = T> + '_> {
135         if self.locked.swap(true, Acquire) {
136             None
137         } else {
138             Some(MutexGuard { lock: self })
139         }
140     }
141 }
142 
143 struct MutexGuard<'a, T> {
144     lock: &'a Mutex<T>,
145 }
146 
147 impl<T> Deref for MutexGuard<'_, T> {
148     type Target = T;
149 
150     fn deref(&self) -> &T {
151         unsafe { &*self.lock.data.get() }
152     }
153 }
154 
155 impl<T> DerefMut for MutexGuard<'_, T> {
156     fn deref_mut(&mut self) -> &mut T {
157         unsafe { &mut *self.lock.data.get() }
158     }
159 }
160 
161 impl<T> Drop for MutexGuard<'_, T> {
162     fn drop(&mut self) {
163         self.lock.locked.store(false, Release);
164     }
165 }
166