Make use of mspaces

src/platform/allocator/dlmalloc.rs

@@ -2,15 +2,20 @@ use core::{
     alloc::{GlobalAlloc, Layout},
     sync::atomic::{AtomicUsize, Ordering},
 };
+use crate::ALLOCATOR;
 
 use super::types::*;
 
 extern "C" {
     fn create_mspace(capacity: size_t, locked: c_int) -> usize;
-    fn dlmalloc(bytes: size_t) -> *mut c_void;
-    fn dlmemalign(alignment: size_t, bytes: size_t) -> *mut c_void;
-    fn dlrealloc(oldmem: *mut c_void, bytes: size_t) -> *mut c_void;
-    fn dlfree(mem: *mut c_void);
+    fn mspace_malloc(msp: usize, bytes: size_t) -> *mut c_void;
+    fn mspace_memalign(msp: usize, alignment: size_t, bytes: size_t) -> *mut c_void;
+    fn mspace_realloc(msp: usize, oldmem: *mut c_void, bytes: size_t) -> *mut c_void;
+    fn mspace_free(msp: usize, mem: *mut c_void);
+    //fn dlmalloc(bytes: size_t) -> *mut c_void;
+    //fn dlmemalign(alignment: size_t, bytes: size_t) -> *mut c_void;
+    //fn dlrealloc(oldmem: *mut c_void, bytes: size_t) -> *mut c_void;
+    //fn dlfree(mem: *mut c_void);
 }
 
 pub struct Allocator {
@@ -40,19 +45,19 @@ unsafe impl<'a> GlobalAlloc for Allocator {
 }
 
 pub unsafe fn alloc(size: usize) -> *mut c_void {
-    dlmalloc(size)
+    mspace_malloc(ALLOCATOR.get_book_keeper(), size)
 }
 
 pub unsafe fn alloc_align(size: usize, alignment: usize) -> *mut c_void {
-    dlmemalign(alignment, size)
+    mspace_memalign(ALLOCATOR.get_book_keeper(), alignment, size)
 }
 
 pub unsafe fn realloc(ptr: *mut c_void, size: size_t) -> *mut c_void {
-    dlrealloc(ptr, size)
+    mspace_realloc(ALLOCATOR.get_book_keeper(), ptr, size)
 }
 
 pub unsafe fn free(ptr: *mut c_void) {
-    dlfree(ptr)
+    mspace_free(ALLOCATOR.get_book_keeper(), ptr)
 }
 
 pub fn new_mspace() -> usize {

src/start.rs

@@ -67,7 +67,26 @@ static INIT_ARRAY: [extern "C" fn(); 1] = [init_array];
 
 static mut init_complete: bool = false;
 
+fn alloc_init() {
+    unsafe{
+        if let Some(tcb) = ld_so::tcb::Tcb::current() {
+            if tcb.mspace != 0 {
+                ALLOCATOR.set_book_keeper(tcb.mspace);
+            } else if ALLOCATOR.get_book_keeper() == 0 {
+                ALLOCATOR.set_book_keeper(new_mspace());
+            }
+        } else if ALLOCATOR.get_book_keeper() == 0 {
+            ALLOCATOR.set_book_keeper(new_mspace());
+        }
+    }
+}
+
 extern "C" fn init_array() {
+    // The thing is that we cannot guarantee if
+    // init_array runs first or if relibc_start runs first
+    // Still whoever gets to run first must initialize rust
+    // memory allocator before doing anything else.
+    alloc_init();
     io_init();
     unsafe { init_complete = true };
 }
@@ -95,11 +114,7 @@ pub unsafe extern "C" fn relibc_start(sp: &'static Stack) -> ! {
     }
     // Step 1 setup the right allocator...
     // if any memory rust based memory allocation happen before this step .. we are doomed.
-    if let Some(tcb) = ld_so::tcb::Tcb::current() {
-        ALLOCATOR.set_book_keeper(tcb.mspace);
-    } else {
-        ALLOCATOR.set_book_keeper(new_mspace());
-    }
+    alloc_init();
 
     // Ensure correct host system before executing more system calls
     relibc_verify_host();