[llvm] [Support] Recycler: Enforce minimum allocation size (PR #121425)

[Akshat Oke via llvm-commits]

https://github.com/optimisan updated https://github.com/llvm/llvm-project/pull/121425

>From 5e6aefd9fcbc41e06f5b665c244e8e1c85ad23f6 Mon Sep 17 00:00:00 2001
From: Akshat Oke <Akshat.Oke at amd.com>
Date: Wed, 1 Jan 2025 06:33:26 +0000
Subject: [PATCH 1/2] [Support] Recycler: Enforce minimum allocation size

Recycler uses reinterpret_cast to an internal structure of size 8.
Invalid write occurs if Recycler is used for objects with sizes less
than 8.
---
 llvm/include/llvm/Support/Recycler.h    |  2 ++
 llvm/unittests/Support/CMakeLists.txt   |  1 +
 llvm/unittests/Support/RecyclerTest.cpp | 46 +++++++++++++++++++++++++
 3 files changed, 49 insertions(+)
 create mode 100644 llvm/unittests/Support/RecyclerTest.cpp

diff --git a/llvm/include/llvm/Support/Recycler.h b/llvm/include/llvm/Support/Recycler.h
index bbd9ae321ae30c..8cc882ea5fa058 100644
--- a/llvm/include/llvm/Support/Recycler.h
+++ b/llvm/include/llvm/Support/Recycler.h
@@ -85,6 +85,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 size must be atleast 8");
     return FreeList ? reinterpret_cast<SubClass *>(pop_val())
                     : static_cast<SubClass *>(Allocator.Allocate(Size, Align));
   }
diff --git a/llvm/unittests/Support/CMakeLists.txt b/llvm/unittests/Support/CMakeLists.txt
index d64f89847aa8e7..6de81658264420 100644
--- a/llvm/unittests/Support/CMakeLists.txt
+++ b/llvm/unittests/Support/CMakeLists.txt
@@ -69,6 +69,7 @@ add_llvm_unittest(SupportTests
   PerThreadBumpPtrAllocatorTest.cpp
   ProcessTest.cpp
   ProgramTest.cpp
+  RecyclerTest.cpp
   RegexTest.cpp
   ReverseIterationTest.cpp
   ReplaceFileTest.cpp
diff --git a/llvm/unittests/Support/RecyclerTest.cpp b/llvm/unittests/Support/RecyclerTest.cpp
new file mode 100644
index 00000000000000..8cd763c0b83f8a
--- /dev/null
+++ b/llvm/unittests/Support/RecyclerTest.cpp
@@ -0,0 +1,46 @@
+//===--- unittest/Support/RecyclerTest.cpp --------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/Support/Recycler.h"
+#include "llvm/Support/AllocatorBase.h"
+#include "gtest/gtest.h"
+
+using namespace llvm;
+
+namespace {
+
+struct Object1 {
+  char Data[1];
+};
+
+class DecoratedMallocAllocator : public MallocAllocator {
+public:
+  int DeallocCount = 0;
+
+  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);
+}
+
+} // end anonymous namespace

>From 8dc68d35da89e8926abddd2ac1437578d80bbc4b Mon Sep 17 00:00:00 2001
From: Akshat Oke <Akshat.Oke at amd.com>
Date: Fri, 3 Jan 2025 10:22:04 +0530
Subject: [PATCH 2/2] Update llvm/include/llvm/Support/Recycler.h

Co-authored-by: Matt Arsenault <Matthew.Arsenault at amd.com>
---
 llvm/include/llvm/Support/Recycler.h | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/llvm/include/llvm/Support/Recycler.h b/llvm/include/llvm/Support/Recycler.h
index 8cc882ea5fa058..1b9ee2084148d0 100644
--- a/llvm/include/llvm/Support/Recycler.h
+++ b/llvm/include/llvm/Support/Recycler.h
@@ -86,7 +86,7 @@ class Recycler {
     static_assert(sizeof(SubClass) <= Size,
                   "Recycler allocation size is less than object size!");
     static_assert(Size >= sizeof(FreeNode) &&
-                  "Recycler size must be atleast 8");
+                  "Recycler size must be at least sizeof(FreeNode)");
     return FreeList ? reinterpret_cast<SubClass *>(pop_val())
                     : static_cast<SubClass *>(Allocator.Allocate(Size, Align));
   }