bring back WeeAlloc; README edit; make the test compile on non-wasm t…

…argets; unsafe fn
OffchainLabs · Dec 4, 2023 · 3953ba2 · 3953ba2
1 parent e3fa50c
commit 3953ba2
Show file tree

Hide file tree

Showing 6 changed files with 75 additions and 17 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/mini-alloc/Cargo.toml b/mini-alloc/Cargo.toml
@@ -15,3 +15,6 @@ wasm-bindgen-test = "0.3.0"
 
 [dependencies]
 cfg-if = "1.0.0"
+
+[target.'cfg(not(target_arch = "wasm32"))'.dependencies]
+wee_alloc = "0.4.5"
diff --git a/mini-alloc/README.md b/mini-alloc/README.md
@@ -13,10 +13,12 @@ thread.
 Also, `core::arch::wasm32::memory_grow` must never be called by any code outside
 this crate.
 
+On targets other than wasm32, `MiniAlloc` simply forwards to the allocator from
+another crate, `wee_alloc::WeeAlloc`.
+
 Use it like this:
 
 ```rust
-#[cfg(target_arch = "wasm32")]
 #[global_allocator]
-static ALLOC: mini_alloc::MiniAlloc = mini_alloc::MiniAlloc;
+static ALLOC: mini_alloc::MiniAlloc = mini_alloc::MiniAlloc::INIT;
 ```
diff --git a/mini-alloc/src/imp.rs b/mini-alloc/src/imp.rs
@@ -12,6 +12,10 @@ pub struct MiniAlloc;
 /// This is not a valid implementation of [`Sync`] but is ok in single-threaded WASM.
 unsafe impl Sync for MiniAlloc {}
 
+impl MiniAlloc {
+    pub const INIT: Self = MiniAlloc;
+}
+
 unsafe impl GlobalAlloc for MiniAlloc {
     unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
         alloc_impl(layout).unwrap_or(core::ptr::null_mut())
@@ -35,12 +39,12 @@ const PAGE_SIZE: usize = 1 << 16;
 /// instructions than the positive offset.
 static mut STATE: Option<(NonZero, usize)> = None;
 
-fn alloc_impl(layout: Layout) -> Option<*mut u8> {
-    let (neg_offset, neg_bound) = unsafe { &mut STATE }.get_or_insert_with(|| {
-        let heap_base = unsafe { &__heap_base } as *const u8 as usize;
+unsafe fn alloc_impl(layout: Layout) -> Option<*mut u8> {
+    let (neg_offset, neg_bound) = STATE.get_or_insert_with(|| {
+        let heap_base = &__heap_base as *const u8 as usize;
         let bound = PAGE_SIZE * wasm32::memory_size(0);
         (
-            unsafe { NonZero::new_unchecked(heap_base.wrapping_neg()) },
+            NonZero::new_unchecked(heap_base.wrapping_neg()),
             bound.wrapping_neg(),
         )
     });
@@ -55,7 +59,7 @@ fn alloc_impl(layout: Layout) -> Option<*mut u8> {
         }
         *neg_bound -= PAGE_SIZE * pages_needed;
     }
-    *neg_offset = unsafe { NonZero::new_unchecked(next_neg_offset) };
+    *neg_offset = NonZero::new_unchecked(next_neg_offset);
     Some(neg_aligned.wrapping_neg() as *mut u8)
 }
 

diff --git a/mini-alloc/src/lib.rs b/mini-alloc/src/lib.rs
@@ -9,5 +9,7 @@ cfg_if! {
     if #[cfg(target_arch = "wasm32")] {
         mod imp;
         pub use imp::MiniAlloc;
+    } else {
+        pub use wee_alloc::WeeAlloc as MiniAlloc;
     }
 }
diff --git a/mini-alloc/tests/misc.rs b/mini-alloc/tests/misc.rs
@@ -5,8 +5,6 @@
 
 extern crate alloc;
 
-use alloc::vec::Vec;
-
 #[cfg(target_arch = "wasm32")]
 use wasm_bindgen_test::*;
 
@@ -17,6 +15,8 @@ static ALLOC: mini_alloc::MiniAlloc = mini_alloc::MiniAlloc;
 #[cfg(target_arch = "wasm32")]
 #[wasm_bindgen_test]
 fn vec_test() {
+    use alloc::vec::Vec;
+
     let p1 = Vec::<u8>::with_capacity(700);
     let p2 = Vec::<u8>::with_capacity(65536);
     let p3 = Vec::<u8>::with_capacity(700000);