From 2bad3d13926f5de344bb85ec58de6abf5e5d3475 Mon Sep 17 00:00:00 2001
From: Andrei Damian <andreidaamian@gmail.com>
Date: Tue, 4 Jun 2024 18:49:13 +0300
Subject: [PATCH] Make deleting on LinkedList aware of the allocator

---
 library/alloc/src/collections/linked_list.rs  |  4 +-
 .../src/collections/linked_list/tests.rs      | 39 +++++++++++++++++++
 2 files changed, 41 insertions(+), 2 deletions(-)

diff --git a/library/alloc/src/collections/linked_list.rs b/library/alloc/src/collections/linked_list.rs
index 1c90c171a155b..01510a6140512 100644
--- a/library/alloc/src/collections/linked_list.rs
+++ b/library/alloc/src/collections/linked_list.rs
@@ -1705,7 +1705,7 @@ impl<'a, T, A: Allocator> CursorMut<'a, T, A> {
         unsafe {
             self.current = unlinked_node.as_ref().next;
             self.list.unlink_node(unlinked_node);
-            let unlinked_node = Box::from_raw(unlinked_node.as_ptr());
+            let unlinked_node = Box::from_raw_in(unlinked_node.as_ptr(), &self.list.alloc);
             Some(unlinked_node.element)
         }
     }
@@ -1946,7 +1946,7 @@ where
                 if (self.pred)(&mut node.as_mut().element) {
                     // `unlink_node` is okay with aliasing `element` references.
                     self.list.unlink_node(node);
-                    return Some(Box::from_raw(node.as_ptr()).element);
+                    return Some(Box::from_raw_in(node.as_ptr(), &self.list.alloc).element);
                 }
             }
         }
diff --git a/library/alloc/src/collections/linked_list/tests.rs b/library/alloc/src/collections/linked_list/tests.rs
index 8dcd59d12d927..d3744c5a9d0c9 100644
--- a/library/alloc/src/collections/linked_list/tests.rs
+++ b/library/alloc/src/collections/linked_list/tests.rs
@@ -1164,3 +1164,42 @@ fn test_drop_panic() {
 
     assert_eq!(unsafe { DROPS }, 8);
 }
+
+#[test]
+fn test_allocator() {
+    use core::alloc::AllocError;
+    use core::alloc::Allocator;
+    use core::alloc::Layout;
+    use core::cell::Cell;
+
+    struct A {
+        has_allocated: Cell<bool>,
+        has_deallocated: Cell<bool>,
+    }
+
+    unsafe impl Allocator for A {
+        fn allocate(&self, layout: Layout) -> Result<NonNull<[u8]>, AllocError> {
+            assert!(!self.has_allocated.get());
+            self.has_allocated.set(true);
+
+            Global.allocate(layout)
+        }
+
+        unsafe fn deallocate(&self, ptr: NonNull<u8>, layout: Layout) {
+            assert!(!self.has_deallocated.get());
+            self.has_deallocated.set(true);
+
+            unsafe { Global.deallocate(ptr, layout) }
+        }
+    }
+
+    let alloc = &A { has_allocated: Cell::new(false), has_deallocated: Cell::new(false) };
+    {
+        let mut list = LinkedList::new_in(alloc);
+        list.push_back(5u32);
+        list.remove(0);
+    }
+
+    assert!(alloc.has_allocated.get());
+    assert!(alloc.has_deallocated.get());
+}