From e31e9451ba966b7c614c999d5e1a3c6806467494 Mon Sep 17 00:00:00 2001
From: Akshat Oke <Akshat.Oke@amd.com>
Date: Tue, 7 Jan 2025 06:10:39 +0000
Subject: [PATCH] [Support] Recycler: Match dealloc size and enforce min size

---
 llvm/include/llvm/Support/Recycler.h    |  4 +++-
 llvm/unittests/Support/RecyclerTest.cpp | 24 ++++++++++++++++++++++++
 2 files changed, 27 insertions(+), 1 deletion(-)

diff --git a/llvm/include/llvm/Support/Recycler.h b/llvm/include/llvm/Support/Recycler.h
index 693c6559ff2fdc..e531e235ee78f8 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 a33506b47ebeae..696e397d3f10ed 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;