From b5fbbdd2025fb04dc8ee1fc892d6575a9a3e2bcd Mon Sep 17 00:00:00 2001
From: Owen Anderson <resistor@mac.com>
Date: Fri, 17 Sep 2021 18:28:58 +0000
Subject: [PATCH] Teach InstCombine to eliminate malloc-realloc-free triplets.

Reviewed By: majnemer

Differential Revision: https://reviews.llvm.org/D109988
---
 .../InstCombine/InstructionCombining.cpp      | 16 +++++++++++++
 .../InstCombine/alloc-realloc-free.ll         | 24 +++++++++++++++++++
 .../Transforms/InstCombine/realloc-free.ll    | 19 +++++++++++++++
 3 files changed, 59 insertions(+)
 create mode 100644 llvm/test/Transforms/InstCombine/alloc-realloc-free.ll
 create mode 100644 llvm/test/Transforms/InstCombine/realloc-free.ll

diff --git a/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp b/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp
index 249915bceebb..d8651a77bc8f 100644
--- a/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp
@@ -2596,6 +2596,13 @@ static bool isAllocSiteRemovable(Instruction *AI,
           Users.emplace_back(I);
           continue;
         }
+
+        if (isReallocLikeFn(I, TLI, true)) {
+          Users.emplace_back(I);
+          Worklist.push_back(I);
+          continue;
+        }
+
         return false;
 
       case Instruction::Store: {
@@ -2808,6 +2815,15 @@ Instruction *InstCombinerImpl::visitFree(CallInst &FI) {
   if (isa<ConstantPointerNull>(Op))
     return eraseInstFromFunction(FI);
 
+  // If we had free(realloc(...)) with no intervening uses, then eliminate the
+  // realloc() entirely.
+  if (CallInst *CI = dyn_cast<CallInst>(Op)) {
+    if (CI->hasOneUse() && isReallocLikeFn(CI, &TLI, true)) {
+      return eraseInstFromFunction(
+          *replaceInstUsesWith(*CI, CI->getOperand(0)));
+    }
+  }
+
   // If we optimize for code size, try to move the call to free before the null
   // test so that simplify cfg can remove the empty block and dead code
   // elimination the branch. I.e., helps to turn something like:
diff --git a/llvm/test/Transforms/InstCombine/alloc-realloc-free.ll b/llvm/test/Transforms/InstCombine/alloc-realloc-free.ll
new file mode 100644
index 000000000000..8f3cf87adcaa
--- /dev/null
+++ b/llvm/test/Transforms/InstCombine/alloc-realloc-free.ll
@@ -0,0 +1,24 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --scrub-attributes
+; RUN: opt -S -instcombine < %s | FileCheck %s
+
+define dso_local void @test() local_unnamed_addr #0 {
+; CHECK-LABEL: @test(
+; CHECK-NEXT:    ret void
+;
+  %1 = tail call noalias align 16 dereferenceable_or_null(4) i8* @malloc(i64 4) #4
+  %2 = tail call align 16 dereferenceable_or_null(6) i8* @realloc(i8* %1, i64 6) #4
+  tail call void @free(i8* %2) #4
+  ret void
+}
+
+declare dso_local noalias noundef i8* @malloc(i64 noundef) local_unnamed_addr #1
+declare dso_local noalias noundef i8* @realloc(i8* nocapture, i64 noundef) local_unnamed_addr #2
+declare dso_local void @free(i8* nocapture noundef) local_unnamed_addr #2
+declare void @llvm.dbg.value(metadata, metadata, metadata) #3
+
+attributes #0 = { mustprogress nounwind uwtable willreturn }
+attributes #1 = { inaccessiblememonly mustprogress nofree nounwind willreturn }
+attributes #2 = { inaccessiblemem_or_argmemonly mustprogress nounwind willreturn }
+attributes #3 = { nofree nosync nounwind readnone speculatable willreturn }
+attributes #4 = { nounwind }
+
diff --git a/llvm/test/Transforms/InstCombine/realloc-free.ll b/llvm/test/Transforms/InstCombine/realloc-free.ll
new file mode 100644
index 000000000000..64311c9ee188
--- /dev/null
+++ b/llvm/test/Transforms/InstCombine/realloc-free.ll
@@ -0,0 +1,19 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --scrub-attributes
+; RUN: opt -S -instcombine < %s | FileCheck %s
+
+define dso_local void @_Z3fooPv(i8* nocapture %0) local_unnamed_addr #0 {
+; CHECK-LABEL: @_Z3fooPv(
+; CHECK-NEXT:    tail call void @free(i8* [[TMP0:%.*]])
+; CHECK-NEXT:    ret void
+;
+  %2 = tail call align 16 dereferenceable_or_null(6) i8* @realloc(i8* %0, i64 6) #2
+  tail call void @free(i8* %2) #2
+  ret void
+}
+
+declare dso_local noalias noundef i8* @realloc(i8* nocapture, i64 noundef) local_unnamed_addr #1
+declare dso_local void @free(i8* nocapture noundef) local_unnamed_addr #1
+
+attributes #0 = { mustprogress nounwind uwtable willreturn }
+attributes #1 = { inaccessiblemem_or_argmemonly mustprogress nounwind willreturn }
+attributes #2 = { nounwind }