[Support] Recycler: Match dealloc size and enforce min size

[optimisan]

Address sanitizer found mismatching deallocation size in Recycler.

[optimisan]

• [Support] Recycler: Match dealloc size and enforce min size #121889 👈 (View in Graphite)
• [Support] Recycler: Implement move constructor #120555 : 1 other dependent PR (#117309 )
• [Support] Recycler: Enforce minimum allocation size #121425
• [CodeGen] LiveRegMatrix: Use allocator through a unique_ptr #120556
• main

This stack of pull requests is managed by Graphite. Learn more about stacking.

[llvmbot]

@llvm/pr-subscribers-llvm-support

Author: Akshat Oke (optimisan)

Changes

Address sanitizer found mismatching deallocation size in Recycler.

---

Full diff: https://github.com/llvm/llvm-project/pull/121889.diff

2 Files Affected:

• (modified) llvm/include/llvm/Support/Recycler.h (+3-1)
• (modified) llvm/unittests/Support/RecyclerTest.cpp (+24)

diff --git a/llvm/include/llvm/Support/Recycler.h b/llvm/include/llvm/Support/Recycler.h
index 693c6559ff2fdc3..e531e235ee78f87 100644
--- a/llvm/include/llvm/Support/Recycler.h
+++ b/llvm/include/llvm/Support/Recycler.h
@@ -72,7 +72,7 @@ class Recycler {
   void clear(AllocatorType &Allocator) {
     while (FreeList) {
       T *t = reinterpret_cast<T *>(pop_val());
-      Allocator.Deallocate(t);
+      Allocator.Deallocate(t, Size, Align);
     }
   }
 
@@ -89,6 +89,8 @@ class Recycler {
                   "Recycler allocation alignment is less than object align!");
     static_assert(sizeof(SubClass) <= Size,
                   "Recycler allocation size is less than object size!");
+    static_assert(Size >= sizeof(FreeNode) &&
+                  "Recycler allocation size must be at least sizeof(FreeNode)");
     return FreeList ? reinterpret_cast<SubClass *>(pop_val())
                     : static_cast<SubClass *>(Allocator.Allocate(Size, Align));
   }
diff --git a/llvm/unittests/Support/RecyclerTest.cpp b/llvm/unittests/Support/RecyclerTest.cpp
index a33506b47ebeae3..696e397d3f10edd 100644
--- a/llvm/unittests/Support/RecyclerTest.cpp
+++ b/llvm/unittests/Support/RecyclerTest.cpp
@@ -14,6 +14,10 @@ using namespace llvm;
 
 namespace {
 
+struct Object1 {
+  char Data[1];
+};
+
 struct Object8 {
   char Data[8];
 };
@@ -22,12 +26,32 @@ class DecoratedMallocAllocator : public MallocAllocator {
 public:
   int DeallocCount = 0;
 
+  void Deallocate(const void *Ptr, size_t Size, size_t Alignment) {
+    DeallocCount++;
+    MallocAllocator::Deallocate(Ptr, Size, Alignment);
+  }
+
   template <typename T> void Deallocate(T *Ptr) {
     DeallocCount++;
     MallocAllocator::Deallocate(Ptr);
   }
 };
 
+TEST(RecyclerTest, RecycleAllocation) {
+  DecoratedMallocAllocator Allocator;
+  // Recycler needs size to be atleast 8 bytes.
+  Recycler<Object1, 8, 8> R;
+  Object1 *A1 = R.Allocate(Allocator);
+  Object1 *A2 = R.Allocate(Allocator);
+  R.Deallocate(Allocator, A2);
+  Object1 *A3 = R.Allocate(Allocator);
+  EXPECT_EQ(A2, A3); // reuse the deallocated object.
+  R.Deallocate(Allocator, A1);
+  R.Deallocate(Allocator, A3);
+  R.clear(Allocator); // Should deallocate A1 and A3.
+  EXPECT_EQ(Allocator.DeallocCount, 2);
+}
+
 TEST(RecyclerTest, MoveConstructor) {
   DecoratedMallocAllocator Allocator;
   Recycler<Object8> R;