fix(alloc): padding node correct size, add more tests

marekvospel · Oct 25, 2024 · 1e5ff50 · 1e5ff50
1 parent c9b7315
commit 1e5ff50
Show file tree

Hide file tree

Showing 2 changed files with 78 additions and 13 deletions.
diff --git a/crates/allocator/src/allocator.rs b/crates/allocator/src/allocator.rs
@@ -30,7 +30,7 @@ unsafe impl<'a> GlobalAlloc for LinkedListAllocator {
                 - size_of::<LinkedAllocatorNode>();
             let mut empty_node = None;
 
-            // section isn't at the end, add padding
+            // section isn't at the end, add padding that is at least LinkedAllocatorNode big
             if aligned_addr
                 != node.end_address() + 1 - layout.size() - size_of::<LinkedAllocatorNode>()
             {
@@ -41,12 +41,15 @@ unsafe impl<'a> GlobalAlloc for LinkedListAllocator {
                         - size_of::<LinkedAllocatorNode>();
 
                 empty_node = Some(
-                    (node.end_address() + 1 - size_of::<LinkedAllocatorNode>())
+                    (aligned_addr + layout.size() + size_of::<LinkedAllocatorNode>())
                         as *mut LinkedAllocatorNode,
                 );
                 let empty_node = &mut *empty_node.unwrap();
                 empty_node.free = true;
-                empty_node.size = 0;
+                empty_node.size = (node.end_address() + 1)
+                    - (aligned_addr + layout.size() + 2 * size_of::<LinkedAllocatorNode>());
+
+                assert_eq!(empty_node.end_address(), node.end_address());
                 empty_node.next = node.next;
             }
 
@@ -134,7 +137,7 @@ impl LinkedListAllocator {
         self.nodes = Some(node as *mut _);
     }
 
-    pub(crate) unsafe fn nodes_mut(&self) -> Option<*mut LinkedAllocatorNode> {
+    pub unsafe fn nodes_mut(&self) -> Option<*mut LinkedAllocatorNode> {
         self.nodes.map(|n| n as *mut _)
     }
 }
diff --git a/crates/allocator/tests/alloc.rs b/crates/allocator/tests/alloc.rs
@@ -36,27 +36,89 @@ fn should_initialize_allocator_once() {
 }
 
 #[test]
-fn should_merge_with_previous_node() {
+#[should_panic]
+fn should_fail_double_initing() {
     const SIZE: usize = 256;
     let mut buf = [0u8; SIZE];
-    let node = create_node(&mut buf[0] as *const _ as usize, SIZE);
+    let mut buf2 = [0u8; SIZE];
+    let node = unsafe { &mut *create_node(&mut buf[0] as *const _ as usize, SIZE) };
+    let node2 = unsafe { &mut *create_node(&mut buf2[0] as *const _ as usize, SIZE) };
 
-    let init_node = unsafe { &mut *node };
-    let node = unsafe { &mut *node };
+    let mut allocator = LinkedListAllocator::new();
+    allocator.init(node);
+    allocator.init(node2);
+}
+
+#[test]
+#[should_panic]
+fn should_fail_not_initialized() {
+    let allocator = LinkedListAllocator::new();
+
+    unsafe {
+        allocator.alloc(Layout::from_size_align_unchecked(32, 1));
+    }
+}
+
+#[test]
+fn should_merge_with_previous_node() {
+    const SIZE: usize = 256;
+    let mut buf = [0u8; SIZE];
+    let node = unsafe { &mut *create_node(&mut buf[0] as *const _ as usize, SIZE) };
 
     let mut allocator = LinkedListAllocator::new();
-    allocator.init(init_node);
+    allocator.init(node);
 
+    unsafe {
+        let ptr = allocator.alloc(Layout::from_size_align_unchecked(32, 1));
+        allocator.dealloc(ptr, Layout::from_size_align_unchecked(32, 1));
+    }
+
+    let node = unsafe { &*(allocator.nodes_mut().unwrap()) };
     assert!(node.free);
-    assert_eq!(node.size, 256 - size_of::<LinkedAllocatorNode>());
+    assert_eq!(node.size, SIZE - size_of::<LinkedAllocatorNode>());
     assert_eq!(node.next, None);
+}
+
+#[test]
+fn should_merge_with_next_node() {
+    const SIZE: usize = 256;
+    let mut buf = [0u8; SIZE];
+    let node = create_node(&mut buf[0] as *const _ as usize, SIZE);
+
+    let node = unsafe { &mut *node };
+
+    let mut allocator = LinkedListAllocator::new();
+    allocator.init(node);
 
     unsafe {
-        let ptr = allocator.alloc(Layout::from_size_align_unchecked(32, 1));
-        allocator.dealloc(ptr, Layout::from_size_align_unchecked(32, 1));
+        let ptr1 = allocator.alloc(Layout::from_size_align_unchecked(32, 32));
+        let ptr2 = allocator.alloc(Layout::from_size_align_unchecked(5, 1));
+
+        print_nodes(&allocator);
+
+        // Start, ptr2, ptr1, reserved
+        {
+            let node = &*(allocator.nodes_mut().unwrap());
+            assert_eq!(node.as_iter().count(), 4);
+        }
+
+        allocator.dealloc(ptr1, Layout::from_size_align_unchecked(32, 32));
+
+        print_nodes(&allocator);
+
+        // Start, ptr2, free
+        {
+            let node = &*(allocator.nodes_mut().unwrap());
+            assert_eq!(node.as_iter().count(), 3);
+        }
+
+        allocator.dealloc(ptr2, Layout::from_size_align_unchecked(5, 1));
     }
 
+    // Free
+    let node = unsafe { &*(allocator.nodes_mut().unwrap()) };
+
     assert!(node.free);
-    assert_eq!(node.size, 256 - size_of::<LinkedAllocatorNode>());
+    assert_eq!(node.size, SIZE - size_of::<LinkedAllocatorNode>());
     assert_eq!(node.next, None);
 }