Fix and comments addressed

refined fix with only instructions in block

Author: AbhayKanhere

llvm/lib/Transforms/Utils/CodeExtractor.cpp

@@ -627,28 +627,23 @@ bool CodeExtractor::isEligible() const {
         return false;
     }
   }
-  // stacksave as input implies stackrestore in the outlined function
-  // This can confuse prologue epilogue insertion phase
-  // stacksave's uses must not cross outlined function
+  // stacksave as input implies stackrestore in the outlined function.
+  // This can confuse prolog epilog insertion phase.
+  // stacksave's uses must not cross outlined function.
   for (BasicBlock *BB : Blocks) {
-    for (Instruction &II : *BB) {
-      if (IntrinsicInst *Intrin = dyn_cast<IntrinsicInst>(&II)) {
-        if (Intrin->getIntrinsicID() == Intrinsic::stacksave) {
-          for (User *U : Intrin->users())
-            if (!definedInRegion(Blocks, U)) {
-              return false; // stack-restore outside outlined region
-            }
-        }
+    for (Instruction &I : *BB) {
+      IntrinsicInst *II = dyn_cast<IntrinsicInst>(&I);
+      if (!II)
+        continue;
+      bool IsSave = II->getIntrinsicID() == Intrinsic::stacksave;
+      bool IsRestore = II->getIntrinsicID() == Intrinsic::stackrestore;
+      if (IsSave && any_of(II->users(), [&Blks = this->Blocks](User *U) {
+            return !definedInRegion(Blks, U);
+          })) {
+        return false;
       }
-      for (auto &OI : II.operands()) {
-        Value *V = OI;
-        if (IntrinsicInst *II = dyn_cast<IntrinsicInst>(V)) {
-          if (II->getIntrinsicID() == Intrinsic::stacksave) {
-            if (definedInCaller(Blocks, V)) {
-              return false;
-            }
-          }
-        }
+      if (IsRestore && !definedInRegion(Blocks, II->getArgOperand(0))) {
+        return false;
       }
     }
   }

llvm/test/Transforms/HotColdSplit/outline-inner-region-stacktoocomplex.ll

@@ -1,16 +1,36 @@
-; RUN: opt -S -passes=hotcoldsplit -hotcoldsplit-max-params=1 < %s | FileCheck %s
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
+; RUN: opt -S --passes=hotcoldsplit --hotcoldsplit-max-params=1 < %s | FileCheck %s
 
 target datalayout = "E-m:a-p:32:32-i64:64-n32"
-target triple = "powerpc64-ibm-aix7.2.0.0"
 
 define void @foo(i32 %cond) {
-; CHECK-LABEL: define {{.*}}@foo(
-; CHECK:       if.then:
-; CHECK:       if.then1:
-; CHECK:       if.else:
-; CHECK:       call {{.*}}@sink
-; CHECK:       if.end:
-; CHECK:       if.end2:
+; CHECK-LABEL: define void @foo(
+; CHECK-SAME: i32 [[COND:%.*]]) {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    [[COND_ADDR:%.*]] = alloca i32, align 4
+; CHECK-NEXT:    store i32 [[COND]], ptr [[COND_ADDR]], align 4
+; CHECK-NEXT:    [[TMP0:%.*]] = load i32, ptr [[COND_ADDR]], align 4
+; CHECK-NEXT:    [[STKS:%.*]] = tail call ptr @llvm.stacksave.p0()
+; CHECK-NEXT:    [[TOBOOL:%.*]] = icmp ne i32 [[TMP0]], 0
+; CHECK-NEXT:    br i1 [[TOBOOL]], label %[[IF_THEN:.*]], label %[[IF_END2:.*]]
+; CHECK:       [[IF_THEN]]:
+; CHECK-NEXT:    [[TMP1:%.*]] = load i32, ptr [[COND_ADDR]], align 4
+; CHECK-NEXT:    call void @sink(i32 [[TMP0]])
+; CHECK-NEXT:    [[CMP:%.*]] = icmp sgt i32 [[TMP1]], 10
+; CHECK-NEXT:    br i1 [[CMP]], label %[[IF_THEN1:.*]], label %[[IF_ELSE:.*]]
+; CHECK:       [[IF_THEN1]]:
+; CHECK-NEXT:    call void @sideeffect(i32 2)
+; CHECK-NEXT:    br label %[[IF_END:.*]]
+; CHECK:       [[IF_ELSE]]:
+; CHECK-NEXT:    call void @sink(i32 0)
+; CHECK-NEXT:    call void @sideeffect(i32 0)
+; CHECK-NEXT:    call void @llvm.stackrestore.p0(ptr [[STKS]])
+; CHECK-NEXT:    br label %[[IF_END]]
+; CHECK:       [[IF_END]]:
+; CHECK-NEXT:    br label %[[IF_END2]]
+; CHECK:       [[IF_END2]]:
+; CHECK-NEXT:    call void @sideeffect(i32 1)
+; CHECK-NEXT:    ret void
 ;
 entry:
   %cond.addr = alloca i32