From 85ec68d69bca5ab9b150efc6dacb03c1696ea865 Mon Sep 17 00:00:00 2001
From: Liqiang Tao <taolq@outlook.com>
Date: Wed, 20 Sep 2023 19:12:34 +0800
Subject: [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
---
 llvm/lib/IR/Function.cpp | 24 ++++++++++++------------
 1 file changed, 12 insertions(+), 12 deletions(-)

(limited to 'llvm/lib/IR/Function.cpp')

diff --git a/llvm/lib/IR/Function.cpp b/llvm/lib/IR/Function.cpp
index b1b8404..658aa67 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);
   }
 
-- 
cgit v1.1