libc: improve allocator behavior

test.c

@@ -2,21 +2,77 @@
 #include <stdio.h>
 #include <math.h>
 #include <setjmp.h>
+#include <stdlib.h>
+#include <string.h>
 
-void will_fail(jmp_buf a) {
-    printf("FAIL!!!\n");
-    longjmp(a, 123);
+static void test1(void) {
+#define N 128
+    void *array[N] = {NULL};
+
+    for (size_t i = 0; i < N; ++i) {
+        array[i] = malloc(i * 2 + 1);
+        memset(array[i], 0xA3, i * 2 + 1);
+    }
+
+    for (size_t i = 0; i < N; i += 2) {
+        free(array[i]);
+        array[i] = NULL;
+    }
+
+    for (size_t i = 1; i < N; i += 2) {
+        for (size_t j = 0; j < i * 2 + 1; ++j) {
+            assert(((const unsigned char *) array[i])[j] == 0xA3);
+        }
+    }
+
+    for (size_t i = 0; i < N; i += 2) {
+        array[i] = malloc(i * 3 + 3);
+        memset(array[i], 0x7D, i * 3 + 3);
+    }
+
+    for (size_t i = 1; i < N; i += 2) {
+        for (size_t j = 0; j < i * 2 + 1; ++j) {
+            assert(((const unsigned char *) array[i])[j] == 0xA3);
+        }
+    }
+    for (size_t i = 0; i < N; i += 2) {
+        for (size_t j = 0; j < i * 3 + 3; ++j) {
+            assert(((const unsigned char *) array[i])[j] == 0x7D);
+        }
+    }
+
+    for (size_t i = 0; i < N; ++i) {
+        free(array[i]);
+    }
+#undef N
+}
+
+static void test2(void) {
+#define N 64
+    void *ptr0 = NULL;
+    void *ptr1 = NULL;
+
+    for (size_t i = 0; i < N; ++i) {
+        printf("Round #%zu\n", i);
+        // 2, 4, 6, ...
+        ptr0 = realloc(ptr0, (i + 1) * 2);
+        for (size_t j = 0; j < i * 2; ++j) {
+            assert(((const unsigned char *) ptr0)[j] == 0x12);
+        }
+        memset(ptr0, 0x12, (i + 1) * 2);
+        // 3, 6, 9, ...
+        ptr1 = realloc(ptr1, (i + 1) * 3);
+        for (size_t j = 0; j < i * 3; ++j) {
+            assert(((const unsigned char *) ptr1)[j] == 0x34);
+        }
+        memset(ptr1, 0x34, (i + 1) * 3);
+    }
+#undef N
 }
 
 int main(int argc, const char **argv) {
-    jmp_buf a;
-    int res = setjmp(a);
-
-    printf("setjmp = %d\n", res);
-
-    if (res == 0) {
-        will_fail(a);
-    }
+    test1();
+    test2();
 
     return 0;
 }

userspace/lib/ygglibc/src/allocator.rs

@@ -1,13 +1,18 @@
 use core::{
-    alloc::{GlobalAlloc, Layout},
-    ffi::c_void,
-    ptr::{self, null_mut, NonNull},
+    alloc::{GlobalAlloc, Layout}, cmp::Ordering, ffi::c_void, ptr::{self, null_mut, NonNull}
 };
 
 use libyalloc::{allocator::BucketAllocator, sys::PageProvider};
 use yggdrasil_rt::{mem::MappingSource, sys as syscall};
 
-use crate::{error::EResult, headers::{errno, string::mem::{memcpy, memset}}, sync::Mutex};
+use crate::{
+    error::{EResult, OptionExt},
+    headers::{
+        errno,
+        string::mem::{memcpy, memset},
+    },
+    sync::Mutex,
+};
 
 struct Allocator;
 struct PageProviderImpl;
@@ -45,62 +50,57 @@ impl PageProvider for PageProviderImpl {
     }
 }
 
-const MALLOC_HEADER_SIZE: usize = 16;
-
 unsafe fn get_allocation(ptr: NonNull<u8>) -> (NonNull<u8>, Layout) {
-    assert!(usize::from(ptr.addr()) > MALLOC_HEADER_SIZE);
-    let real_ptr = ptr.sub(MALLOC_HEADER_SIZE);
-    let size = *real_ptr.cast::<usize>().as_ref();
-    let layout = Layout::from_size_align(size, 16).unwrap();
+    assert!(usize::from(ptr.addr()) >= 16);
+    let real_ptr = ptr.sub(16);
+    let size = real_ptr.cast::<usize>().read();
+    let align = real_ptr.cast::<usize>().add(1).read();
+    assert_eq!(align, 16);
+    let layout = Layout::from_size_align(size, align).unwrap();
     (real_ptr, layout)
 }
 
 pub fn c_alloc(size: usize, mut align: usize, zero: bool) -> EResult<NonNull<c_void>> {
-    if align < 16 {
-        align = 16;
-    }
-    if align > 16 {
-        todo!()
-    }
-    let size = size + MALLOC_HEADER_SIZE;
-    let layout = Layout::from_size_align(size, align).unwrap();
-
-    let ptr = match YALLOC.lock().allocate(layout) {
-        Some(value) => value,
-        None => return EResult::Err(errno::ENOMEM),
+    // allocation layout:
+    // [ header ] [ ... ]
+    // header: 16 bytes
+    let align = match align {
+        0..=16 => 16,
+        _ => panic!("TODO: handle larger aligns: align={}", align),
     };
-
-    if zero {
-        unsafe {
-            memset(ptr.as_ptr().cast(), 0, size);
-        }
-    }
+    let size = size + 16;
+    let layout = Layout::from_size_align(size, align).expect("Couldn't setup malloc() layout");
+    let ptr = YALLOC.lock().allocate(layout).e_ok_or(errno::ENOMEM)?;
 
     unsafe {
-        // Write the size right below the pointer
         ptr.cast::<usize>().write(size);
+        ptr.cast::<usize>().add(1).write(align);
     }
 
-    unsafe { EResult::Ok(ptr.cast::<c_void>().add(MALLOC_HEADER_SIZE)) }
+    EResult::Ok(unsafe { ptr.add(16).cast() })
 }
 
 pub unsafe fn c_realloc(old_ptr: Option<NonNull<c_void>>, size: usize) -> EResult<NonNull<c_void>> {
     match old_ptr {
         Some(old_ptr) => {
-            // TODO libyalloc realloc
             let (real_old_ptr, old_layout) = get_allocation(old_ptr.cast());
-            let new_ptr = c_alloc(size, 16, false)?;
+            assert!(old_layout.size() >= 16);
+            let old_size = old_layout.size() - 16;
 
-            memcpy(
-                new_ptr.cast().as_ptr(),
-                old_ptr.cast().as_ptr(),
-                old_layout.size() - MALLOC_HEADER_SIZE,
-            );
+            if size == old_size {
+                return EResult::Ok(old_ptr);
+            }
+
+            let new_ptr = c_alloc(size, old_layout.align(), false)?;
+
+            // Copy no more memory than was allocated and will be allocated
+            let memcpy_size = core::cmp::min(size, old_size);
+            memcpy(new_ptr.as_ptr(), old_ptr.as_ptr(), memcpy_size);
 
             YALLOC.lock().free(real_old_ptr, old_layout);
 
             EResult::Ok(new_ptr)
-        }
+        },
         None => c_alloc(size, 16, false),
     }
 }
@@ -110,14 +110,6 @@ pub unsafe fn c_free(ptr: NonNull<c_void>) {
     YALLOC.lock().free(real_ptr, layout);
 }
 
-pub unsafe fn malloc(size: usize) -> EResult<NonNull<c_void>> {
-    c_alloc(size, 16, false)
-}
-
-pub unsafe fn free(ptr: NonNull<c_void>) {
-    c_free(ptr)
-}
-
 #[global_allocator]
 static ALLOCATOR: Allocator = Allocator;
 

userspace/lib/ygglibc/src/headers/setjmp/x86_64.rs

@@ -28,6 +28,12 @@ unsafe extern "C" fn setjmp(buf: __jmp_buf) -> c_int {
     )
 }
 
+#[no_mangle]
+#[naked]
+unsafe extern "C" fn _setjmp(buf: __jmp_buf) -> c_int {
+    core::arch::naked_asm!("jmp setjmp", options(att_syntax))
+}
+
 #[no_mangle]
 #[naked]
 unsafe extern "C" fn longjmp(buf: __jmp_buf, val: c_int) -> c_int {
@@ -51,3 +57,9 @@ unsafe extern "C" fn longjmp(buf: __jmp_buf, val: c_int) -> c_int {
         ret
     "#, options(att_syntax))
 }
+
+#[no_mangle]
+#[naked]
+unsafe extern "C" fn _longjmp(buf: __jmp_buf, val: c_int) -> c_int {
+    core::arch::naked_asm!("jmp longjmp", options(att_syntax))
+}

userspace/lib/ygglibc/src/headers/string/str.rs

@@ -5,7 +5,7 @@ use core::{
 };
 
 use crate::{
-    allocator::{self, malloc},
+    allocator,
     error::CPtrResult,
     headers::{
         errno::{self, Errno},
@@ -179,7 +179,7 @@ unsafe extern "C" fn strncpy(dst: *mut c_char, src: *const c_char, n: usize) ->
 unsafe extern "C" fn strndup(s: *const c_char, n: usize) -> CPtrResult<c_char> {
     let src = NonNull::new(s.cast_mut()).unwrap();
     let len = strnlen(s, n);
-    let dst = malloc(len + 1)?.cast();
+    let dst = allocator::c_alloc(len + 1, 1, false)?.cast();
     dst.copy_from_nonoverlapping(src, len);
     dst.add(len).write(0);
     CPtrResult::success(dst)