Fix use-after-free in CodeGenPrepare

[beccadax]

Summary

If OptimizeExtractBits() encountered a shift instruction with no operands at all, it would erase the instruction, but still return false. This previously didn’t matter because its caller would always return after processing the instruction, but https://reviews.llvm.org/D63233 changed the function’s caller to fall through if it returned false, which would then cause a use-after-free detectable by ASAN.

This change makes OptimizeExtractBits return true if it removes a shift instruction with no users, terminating processing of the instruction.

[lebedev.ri]

Please file a blocker bug for https://bugs.llvm.org/show_bug.cgi?id=42474, this sounds 9.0-worthy.
This looks obviously correct (C), but wait a bit in case anyone else wants to comment.

[spatel]

LGTM - in D63233, I added a TODO that I've marked here for reference, but this appears to be a minimal/safe fix.

llvm/lib/CodeGen/CodeGenPrepare.cpp

7019

Alternative to this patch: act on that TODO.

[beccadax]

I've filed a bug at https://bugs.llvm.org/show_bug.cgi?id=43021.

[beccadax]

llvm/lib/CodeGen/CodeGenPrepare.cpp

7019

Even if I did act on the TODO, I think we would still need the fix in this patch to avoid calling optimizeShiftInst() with the freed instruction.

[beccadax]

@spatel @lebedev.ri I don't have commit access, so could one of you commit this when you think it's had enough time for comments?

llvm/lib/CodeGen/CodeGenPrepare.cpp

if (!InsertedShift) {
 
       MadeChange = true;
     }
 
     // Replace a use of the shift with a use of the new shift.
     TheUse = InsertedShift;
   }
 
-  // If we removed all uses, nuke the shift.
+  // If we removed all uses, or there are none, nuke the shift.
   if (ShiftI->use_empty()) {
     salvageDebugInfo(*ShiftI);
     ShiftI->eraseFromParent();
+    MadeChange = true;
   }
 
   return MadeChange;
 }
 
 /// If counting leading or trailing zeros is an expensive operation and a zero
 /// input is defined, add a check for zero to avoid calling the intrinsic.
 ///
bool CodeGenPrepare::optimizeInst(Instruction *I, bool &ModifiedDT) {
   }
 
   BinaryOperator *BinOp = dyn_cast<BinaryOperator>(I);
 
   if (BinOp && (BinOp->getOpcode() == Instruction::And) &&
       EnableAndCmpSinking && TLI)
     return sinkAndCmp0Expression(BinOp, *TLI, InsertedInsts);
 
   // TODO: Move this into the switch on opcode - it handles shifts already.
   if (BinOp && (BinOp->getOpcode() == Instruction::AShr ||
                 BinOp->getOpcode() == Instruction::LShr)) {
     ConstantInt *CI = dyn_cast<ConstantInt>(BinOp->getOperand(1));
     if (TLI && CI && TLI->hasExtractBitsInsn())
       if (OptimizeExtractBits(BinOp, CI, *TLI, *DL))
         return true;
   }
 

llvm/test/Transforms/CodeGenPrepare/sink-shift-and-trunc.ll

 ; REQUIRES: aarch64-registered-target
-; RUN: opt -codegenprepare -mtriple=arm64-apple=ios -S -o - %s | FileCheck %s
+; RUN: opt -codegenprepare -mtriple=arm64-apple-ios -S -o - %s | FileCheck %s
 
 @first_ones = external global [65536 x i8]
 
 define i32 @fct19(i64 %arg1) #0 !dbg !6 {
 ; CHECK-LABEL: @fct19
 entry:
   %x.sroa.1.0.extract.shift = lshr i64 %arg1, 16, !dbg !35
   %x.sroa.1.0.extract.trunc = trunc i64 %x.sroa.1.0.extract.shift to i16, !dbg !36
if.then17:                                        ; preds = %if.end13
   %add23 = add nsw i32 %conv22, 32, !dbg !60
   br label %return, !dbg !61
 
 return:                                           ; preds = %if.then17, %if.end13, %if.then7, %if.then
   %retval.0 = phi i32 [ %conv, %if.then ], [ %add, %if.then7 ], [ %add23, %if.then17 ], [ 64, %if.end13 ], !dbg !62
   ret i32 %retval.0, !dbg !63
 }
 
+define i32 @shift_unused(i32 %a) {
+; CHECK-LABEL: @shift_unused
+
+  %as = lshr i32 %a, 3
+  br label %BB2
+
+; CHECK-NEXT: BB2:
+BB2:
+  ; CodeGenPrepare was erasing the unused lshr instruction, but then further
+  ; processing the instruction after it was freed. If this bug is still present,
+  ; this test will always crash in an LLVM built with ASAN enabled, and may
+  ; crash even if ASAN is not enabled.
+
+; CHECK-NEXT: ret i32
+  ret i32 %a
+}
+
 ; CHECK: [[shift1_loc]] = !DILocation(line: 1
 ; CHECK: [[trunc1_loc]] = !DILocation(line: 2
 ; CHECK: [[shift2_loc]] = !DILocation(line: 3
 
 declare void @llvm.dbg.value(metadata, metadata, metadata) #1
 
 attributes #0 = { nounwind readonly ssp }
 attributes #1 = { nounwind readnone speculatable }