[llvm] 85ec68d - [IR] Fix a memory leak if Function::dropAllReferences() is followed by setHungoffOperand

[Liqiang Tao via llvm-commits]

Author: Liqiang Tao
Date: 2023-09-20T19:13:28+08:00
New Revision: 85ec68d69bca5ab9b150efc6dacb03c1696ea865

URL: https://github.com/llvm/llvm-project/commit/85ec68d69bca5ab9b150efc6dacb03c1696ea865
DIFF: https://github.com/llvm/llvm-project/commit/85ec68d69bca5ab9b150efc6dacb03c1696ea865.diff

LOG: [IR] Fix a memory leak if Function::dropAllReferences() is followed by setHungoffOperand

This patch fixes a memory leak if Function::dropAllReferences() is followed by setHungoffOperand (e.g. setPersonality)
If NumUserOperands changes from 3 to 0 before calling allocHungoffUselist() to allocate memory,
the memory leaks which are allocated when NumUserOperands is changed from 0 to 3.
e.g.
```
llvm::Function* func = ...;
func->setPersonalityFn(foo);  // (1). call allocHungoffUselist() to allocate memory for uses
func->deleteBody();  // (2). call dropAllReferences(), and it changes NumUserOperands from 3 to 0
// (3). at this point, NumUserOperands is 0, the next line will allocate memory by allocHungoffUselist()
func->setPersonalityFn(bar);  // (4). call allocHungoffUselist(), so memory allocated in (1) leaks.
```

Reviewed By: dexonsmith, MaskRay

Differential Revision: https://reviews.llvm.org/D156618

Added: 
    

Modified: 
    llvm/include/llvm/IR/Function.h
    llvm/lib/IR/Function.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/include/llvm/IR/Function.h b/llvm/include/llvm/IR/Function.h
index 0feb536620e5990..692ef7535f0e6b1 100644
--- a/llvm/include/llvm/IR/Function.h
+++ b/llvm/include/llvm/IR/Function.h
@@ -116,6 +116,8 @@ class LLVM_EXTERNAL_VISIBILITY Function : public GlobalObject,
 
   void clearArguments();
 
+  void deleteBodyImpl(bool ShouldDrop);
+
   /// Function ctor - If the (optional) Module argument is specified, the
   /// function is automatically inserted into the end of the function list for
   /// the module.
@@ -667,7 +669,7 @@ class LLVM_EXTERNAL_VISIBILITY Function : public GlobalObject,
   /// the linkage to external.
   ///
   void deleteBody() {
-    dropAllReferences();
+    deleteBodyImpl(/*ShouldDrop=*/false);
     setLinkage(ExternalLinkage);
   }
 
@@ -882,7 +884,9 @@ class LLVM_EXTERNAL_VISIBILITY Function : public GlobalObject,
   /// function, dropping all references deletes the entire body of the function,
   /// including any contained basic blocks.
   ///
-  void dropAllReferences();
+  void dropAllReferences() {
+    deleteBodyImpl(/*ShouldDrop=*/true);
+  }
 
   /// hasAddressTaken - returns true if there are any uses of this function
   /// other than direct calls or invokes to it, or blockaddress expressions.

diff  --git a/llvm/lib/IR/Function.cpp b/llvm/lib/IR/Function.cpp
index b1b8404157c3b2d..658aa67a38c937f 100644
--- a/llvm/lib/IR/Function.cpp
+++ b/llvm/lib/IR/Function.cpp
@@ -518,15 +518,7 @@ void Function::stealArgumentListFrom(Function &Src) {
   Src.setValueSubclassData(Src.getSubclassDataFromValue() | (1 << 0));
 }
 
-// dropAllReferences() - This function causes all the subinstructions to "let
-// go" of all references that they are maintaining.  This allows one to
-// 'delete' a whole class at a time, even though there may be circular
-// references... first all references are dropped, and all use counts go to
-// zero.  Then everything is deleted for real.  Note that no operations are
-// valid on an object that has "dropped all references", except operator
-// delete.
-//
-void Function::dropAllReferences() {
+void Function::deleteBodyImpl(bool ShouldDrop) {
   setIsMaterializable(false);
 
   for (BasicBlock &BB : *this)
@@ -537,10 +529,18 @@ void Function::dropAllReferences() {
   while (!BasicBlocks.empty())
     BasicBlocks.begin()->eraseFromParent();
 
-  // Drop uses of any optional data (real or placeholder).
   if (getNumOperands()) {
-    User::dropAllReferences();
-    setNumHungOffUseOperands(0);
+    if (ShouldDrop) {
+      // Drop uses of any optional data (real or placeholder).
+      User::dropAllReferences();
+      setNumHungOffUseOperands(0);
+    } else {
+      // The code needs to match Function::allocHungoffUselist().
+      auto *CPN = ConstantPointerNull::get(PointerType::get(getContext(), 0));
+      Op<0>().set(CPN);
+      Op<1>().set(CPN);
+      Op<2>().set(CPN);
+    }
     setValueSubclassData(getSubclassDataFromValue() & ~0xe);
   }