[sancov] Include dynamic stack allocations in stack-depth tracing

The stack-depth instrumentation was being placed only after all
constant-sized stack allocations, and did not take into account
dynamically sized allocations.

Move the instrumentation after the last alloca present in the Entry
Basic Block.

Author: kees

clang/docs/SanitizerCoverage.rst

@@ -421,10 +421,6 @@ prototype is:
   extern "C"
   void __sanitize_cov_stack_depth(void);
 
-Note that, currently, dynamically sized stacks are not tracked by
-instrumentation correctly, as it is inserted too early. This means
-that only constant sized stack allocations are currently tracked.
-
 Gated Trace Callbacks
 =====================
 

llvm/lib/Transforms/Instrumentation/SanitizerCoverage.cpp

@@ -1042,8 +1042,6 @@ void ModuleSanitizerCoverage::InjectCoverageAtBlock(Function &F, BasicBlock &BB,
   if (IsEntryBB) {
     if (auto SP = F.getSubprogram())
       EntryLoc = DILocation::get(SP->getContext(), SP->getScopeLine(), 0, SP);
-    // FIXME: stack-depth does not correctly instrument dynamic allocas.
-    //
     // Keep static allocas and llvm.localescape calls in the entry block.  Even
     // if we aren't splitting the block, it's nice for allocas to be before
     // calls.
@@ -1094,6 +1092,17 @@ void ModuleSanitizerCoverage::InjectCoverageAtBlock(Function &F, BasicBlock &BB,
   }
   if (Options.StackDepth && IsEntryBB && !IsLeafFunc) {
     Module *M = F.getParent();
+
+    // Find an insertion point after last "alloca".
+    llvm::Instruction *InsertBefore = NULL;
+    for (auto &I : BB) {
+      if (llvm::isa<llvm::AllocaInst>(I))
+        InsertBefore = I.getNextNode(); // Move past the "alloca".
+    }
+    // But only use it if we actually found an "alloca".
+    if (InsertBefore)
+      IRB.SetInsertPoint(InsertBefore);
+
     if (Options.StackDepthCallbackMin) {
       // In callback mode, only add call when stack depth reaches minimum.
       const DataLayout &DL = M->getDataLayout();

llvm/test/Instrumentation/SanitizerCoverage/stack-depth-callback.ll

@@ -202,13 +202,13 @@ entry:
 define i32 @alloc0_32xDyn(i32 %input) {
 ; COMMON-LABEL: define i32 @alloc0_32xDyn(i32 %input) {
 ; COMMON-NEXT:  entry:
+; COMMON-NEXT:    [[VAR:%.*]] = alloca i8, i32 %input, align 4
 ; CB1-NEXT:       call void @__sanitizer_cov_stack_depth()
 ; CB8-NEXT:       call void @__sanitizer_cov_stack_depth()
 ; CB16-NEXT:      call void @__sanitizer_cov_stack_depth()
 ; CB32-NEXT:      call void @__sanitizer_cov_stack_depth()
 ; CB64-NEXT:      call void @__sanitizer_cov_stack_depth()
 ; CB128-NEXT:     call void @__sanitizer_cov_stack_depth()
-; COMMON-NEXT:    [[VAR:%.*]] = alloca i8, i32 %input, align 4
 ; COMMON-NEXT:    [[CALL:%.*]] = call i32 @foo()
 ; COMMON-NEXT:    ret i32 [[CALL]]
 entry:
@@ -225,22 +225,21 @@ entry:
 ;   return foo();
 ; }
 define dso_local i32 @dynamic_alloca(i32 noundef %0) #0 {
-; COMMON-LABEL: define dso_local i32 @dynamic_alloca(i32 noundef %0) {
-; COMMON-NEXT:    [[VAR:%.*]] = alloca i32, align 4
-; COMMON-NEXT:    [[VAR:%.*]] = alloca ptr, align 8
-; COMMON-NEXT:    [[VAR:%.*]] = alloca i64, align 8
-; CB1-NEXT:       call void @__sanitizer_cov_stack_depth()
-; CB8-NEXT:       call void @__sanitizer_cov_stack_depth()
-; CB16-NEXT:      call void @__sanitizer_cov_stack_depth()
-; CB32-NEXT:      call void @__sanitizer_cov_stack_depth()
-; CB64-NEXT:      call void @__sanitizer_cov_stack_depth()
-; CB128-NEXT:     call void @__sanitizer_cov_stack_depth()
   %2 = alloca i32, align 4
   %3 = alloca ptr, align 8
   %4 = alloca i64, align 8
   store i32 %0, ptr %2, align 4
   %5 = load i32, ptr %2, align 4
   %6 = zext i32 %5 to i64
+; COMMON-LABEL:   %7 = call ptr @llvm.stacksave
+; COMMON-NEXT:    store ptr %7, ptr %3, align 8
+; COMMON-NEXT:    [[VAR:%.*]] = alloca i32, i64 %6, align 16
+; CB1-NEXT:       call void @__sanitizer_cov_stack_depth()
+; CB8-NEXT:       call void @__sanitizer_cov_stack_depth()
+; CB16-NEXT:      call void @__sanitizer_cov_stack_depth()
+; CB32-NEXT:      call void @__sanitizer_cov_stack_depth()
+; CB64-NEXT:      call void @__sanitizer_cov_stack_depth()
+; CB128-NEXT:     call void @__sanitizer_cov_stack_depth()
   %7 = call ptr @llvm.stacksave.p0()
   store ptr %7, ptr %3, align 8
   %8 = alloca i32, i64 %6, align 16