[asan] Improve moving of non-instrumented allocas

[kubamracek]

Summary

In r242510, non-instrumented allocas are now moved into the first basic block. This patch limits that to:

1. Only move allocas that were originally in the first basic block.
2. Only move allocas that are present *after* the first instrumented one (i.e. only move allocas up).

A testcase was updated to show behavior in these two cases. Without the patch, an alloca could be moved down, and could cause an invalid IR.

[kcc]

LGTM

[kubamracek]

Updated patch with one more assert (InsBefore is always in the entry block).

[zaks.anna]

LGTM as well. Thanks Kuba!

[atrick]

LGTM!

llvm/trunk/lib/Transforms/Instrumentation/AddressSanitizer.cpp

 //
 //===----------------------------------------------------------------------===//
 
 #include "llvm/Transforms/Instrumentation.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/DepthFirstIterator.h"
+#include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/SmallSet.h"
 #include "llvm/ADT/SmallString.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/Statistic.h"
 #include "llvm/ADT/StringExtras.h"
 #include "llvm/ADT/Triple.h"
 #include "llvm/Analysis/MemoryBuiltins.h"
 #include "llvm/Analysis/TargetLibraryInfo.h"
struct FunctionStackPoisoner : public InstVisitor<FunctionStackPoisoner> {
   AddressSanitizer &ASan;
   DIBuilder DIB;
   LLVMContext *C;
   Type *IntptrTy;
   Type *IntptrPtrTy;
   ShadowMapping Mapping;
 
   SmallVector<AllocaInst *, 16> AllocaVec;
-  SmallVector<AllocaInst *, 16> NonInstrumentedStaticAllocaVec;
+  SmallSetVector<AllocaInst *, 16> NonInstrumentedStaticAllocaVec;
   SmallVector<Instruction *, 8> RetVec;
   unsigned StackAlignment;
 
   Function *AsanStackMallocFunc[kMaxAsanStackMallocSizeClass + 1],
       *AsanStackFreeFunc[kMaxAsanStackMallocSizeClass + 1];
   Function *AsanPoisonStackMemoryFunc, *AsanUnpoisonStackMemoryFunc;
   Function *AsanAllocaPoisonFunc, *AsanAllocasUnpoisonFunc;
 
struct FunctionStackPoisoner : public InstVisitor<FunctionStackPoisoner> {
   //   addr = tmp + 32 (first 32 bytes are for the left redzone).
   // Additional_size is added to make new memory allocation contain not only
   // requested memory, but also left, partial and right redzones.
   void handleDynamicAllocaCall(AllocaInst *AI);
 
   /// \brief Collect Alloca instructions we want (and can) handle.
   void visitAllocaInst(AllocaInst &AI) {
     if (!ASan.isInterestingAlloca(AI)) {
-      if (AI.isStaticAlloca()) NonInstrumentedStaticAllocaVec.push_back(&AI);
+      if (AI.isStaticAlloca()) NonInstrumentedStaticAllocaVec.insert(&AI);
       return;
     }
 
     StackAlignment = std::max(StackAlignment, AI.getAlignment());
     if (ASan.isDynamicAlloca(AI))
       DynamicAllocaVec.push_back(&AI);
     else
       AllocaVec.push_back(&AI);
void FunctionStackPoisoner::poisonStack() {
   DebugLoc EntryDebugLocation;
   if (auto SP = getDISubprogram(&F))
     EntryDebugLocation = DebugLoc::get(SP->getScopeLine(), 0, SP);
 
   Instruction *InsBefore = AllocaVec[0];
   IRBuilder<> IRB(InsBefore);
   IRB.SetCurrentDebugLocation(EntryDebugLocation);
 
-  // Make sure non-instrumented allocas stay in the first basic block.
-  // Otherwise, debug info is broken, because only first-basic-block allocas are
-  // treated as regular stack slots.
-  for (auto *AI : NonInstrumentedStaticAllocaVec) AI->moveBefore(InsBefore);
+  // Make sure non-instrumented allocas stay in the entry block. Otherwise,
+  // debug info is broken, because only entry-block allocas are treated as
+  // regular stack slots.
+  auto InsBeforeB = InsBefore->getParent();
+  assert(InsBeforeB == &F.getEntryBlock());
+  for (BasicBlock::iterator I = InsBefore; I != InsBeforeB->end(); ++I)
+    if (auto *AI = dyn_cast_or_null<AllocaInst>(I))
+      if (NonInstrumentedStaticAllocaVec.count(AI) > 0)
+        AI->moveBefore(InsBefore);
 
   // If we have a call to llvm.localescape, keep it in the entry block.
   if (LocalEscapeCall) LocalEscapeCall->moveBefore(InsBefore);
 
   SmallVector<ASanStackVariableDescription, 16> SVD;
   SVD.reserve(AllocaVec.size());
   for (AllocaInst *AI : AllocaVec) {
     ASanStackVariableDescription D = {AI->getName().data(),

llvm/trunk/test/Instrumentation/AddressSanitizer/debug_info_noninstrumented_alloca.ll

 ; This test checks that non-instrumented allocas stay in the first basic block.
 ; Only first-basic-block allocas are considered stack slots, and moving them
 ; breaks debug info.
 
 ; RUN: opt < %s -asan -asan-module -S | FileCheck %s
 
 target datalayout = "e-m:o-i64:64-f80:128-n8:16:32:64-S128"
 target triple = "x86_64-apple-macosx10.10.0"
 
 define i32 @foo() sanitize_address {
 entry:
+  ; Won't be instrumented because of asan-skip-promotable-allocas.
+  %non_instrumented1 = alloca i32, align 4
+
   ; Regular alloca, will get instrumented (forced by the ptrtoint below).
   %instrumented = alloca i32, align 4
 
   ; Won't be instrumented because of asan-skip-promotable-allocas.
-  %non_instrumented = alloca i32, align 4
-  store i32 0, i32* %non_instrumented, align 4
-  %value = load i32, i32* %non_instrumented, align 4
+  %non_instrumented2 = alloca i32, align 4
+
+  br label %bb0
+
+bb0:
+  ; Won't be instrumented because of asan-skip-promotable-allocas.
+  %non_instrumented3 = alloca i32, align 4
+
+  %ptr = ptrtoint i32* %instrumented to i32
+  br label %bb1
 
-  %ptr = ptrtoint i32* %instrumented to i64
-  ret i32 %value
+bb1:
+  ret i32 %ptr
 }
 
 ; CHECK: entry:
-; CHECK: %non_instrumented = alloca i32, align 4
+; CHECK: %non_instrumented1 = alloca i32, align 4
+; CHECK: %non_instrumented2 = alloca i32, align 4
 ; CHECK: load i32, i32* @__asan_option_detect_stack_use_after_return
+; CHECK: bb0:
+; CHECK: %non_instrumented3 = alloca i32, align 4