From adf9928ba42a0a6341d09b99c1ba973b60c7411b Mon Sep 17 00:00:00 2001
From: Brian Gesiak <modocache@gmail.com>
Date: Thu, 15 Feb 2018 19:31:45 +0000
Subject: [PATCH] [Coroutines] Don't move stores for allocator args

Summary:
The behavior described in Coroutines TS `[dcl.fct.def.coroutine]/7`
allows coroutine parameters to be passed into allocator functions.
The instructions to store values into the alloca'd parameters must not
be moved past the frame allocation, otherwise uninitialized values are
passed to the allocator.

Test Plan: `check-llvm`

Reviewers: rsmith, GorNishanov, eric_niebler

Reviewed By: GorNishanov

Subscribers: compnerd, EricWF, llvm-commits

Differential Revision: https://reviews.llvm.org/D43000

git-svn-id: https://llvm.org/svn/llvm-project/llvm/trunk@325285 91177308-0d34-0410-b5e6-96231b3b80d8
---
 lib/Transforms/Coroutines/CoroSplit.cpp        | 17 ++++++-
 test/Transforms/Coroutines/coro-split-alloc.ll | 64 ++++++++++++++++++++++++++
 2 files changed, 80 insertions(+), 1 deletion(-)
 create mode 100644 test/Transforms/Coroutines/coro-split-alloc.ll

diff --git a/lib/Transforms/Coroutines/CoroSplit.cpp b/lib/Transforms/Coroutines/CoroSplit.cpp
index 4a69fbfe435..28c291696ab 100644
--- a/lib/Transforms/Coroutines/CoroSplit.cpp
+++ b/lib/Transforms/Coroutines/CoroSplit.cpp
@@ -654,13 +654,28 @@ getNotRelocatableInstructions(CoroBeginInst *CoroBegin,
   // set.
   do {
     Instruction *Current = Work.pop_back_val();
+    DEBUG(dbgs() << "CoroSplit: Will not relocate: " << *Current << "\n");
     DoNotRelocate.insert(Current);
     for (Value *U : Current->operands()) {
       auto *I = dyn_cast<Instruction>(U);
       if (!I)
         continue;
-      if (isa<AllocaInst>(U))
+
+      if (auto *A = dyn_cast<AllocaInst>(I)) {
+        // Stores to alloca instructions that occur before the coroutine frame
+        // is allocated should not be moved; the stored values may be used by
+        // the coroutine frame allocator. The operands to those stores must also
+        // remain in place.
+        for (const auto &User : A->users())
+          if (auto *SI = dyn_cast<llvm::StoreInst>(User))
+            if (RelocBlocks.count(SI->getParent()) != 0 &&
+                DoNotRelocate.count(SI) == 0) {
+              Work.push_back(SI);
+              DoNotRelocate.insert(SI);
+            }
         continue;
+      }
+
       if (DoNotRelocate.count(I) == 0) {
         Work.push_back(I);
         DoNotRelocate.insert(I);
diff --git a/test/Transforms/Coroutines/coro-split-alloc.ll b/test/Transforms/Coroutines/coro-split-alloc.ll
new file mode 100644
index 00000000000..a0ec0509c77
--- /dev/null
+++ b/test/Transforms/Coroutines/coro-split-alloc.ll
@@ -0,0 +1,64 @@
+; Tests that coro-split passes initialized values to coroutine frame allocator.
+; RUN: opt < %s -coro-split -S | FileCheck %s
+
+define i8* @f(i32 %argument) "coroutine.presplit"="1" {
+entry:
+  %argument.addr = alloca i32, align 4
+  %incremented = add i32 %argument, 1
+  store i32 %incremented, i32* %argument.addr, align 4
+  %id = call token @llvm.coro.id(i32 0, i8* null, i8* null, i8* null)
+  %need.alloc = call i1 @llvm.coro.alloc(token %id)
+  br i1 %need.alloc, label %dyn.alloc, label %begin
+
+dyn.alloc:
+  %size = call i32 @llvm.coro.size.i32()
+  %allocator_argument = load i32, i32* %argument.addr, align 4
+  %alloc = call i8* @custom_alloctor(i32 %size, i32 %allocator_argument)
+  br label %begin
+
+begin:
+  %phi = phi i8* [ null, %entry ], [ %alloc, %dyn.alloc ]
+  %hdl = call i8* @llvm.coro.begin(token %id, i8* %phi)
+  %print_argument = load i32, i32* %argument.addr, align 4
+  call void @print(i32 %print_argument)
+  %0 = call i8 @llvm.coro.suspend(token none, i1 false)
+  switch i8 %0, label %suspend [i8 0, label %resume
+                                i8 1, label %cleanup]
+resume:
+  call void @print(i32 1)
+  br label %cleanup
+
+cleanup:
+  %mem = call i8* @llvm.coro.free(token %id, i8* %hdl)
+  call void @free(i8* %mem)
+  br label %suspend
+suspend:
+  call i1 @llvm.coro.end(i8* %hdl, i1 0)
+  ret i8* %hdl
+}
+
+; CHECK-LABEL: @f(
+; CHECK: %argument.addr = alloca i32
+; CHECK: %incremented = add i32 %argument, 1
+; CHECK-NEXT: store i32 %incremented, i32* %argument.addr
+; CHECK-LABEL: dyn.alloc:
+; CHECK: %allocator_argument = load i32, i32* %argument.addr
+; CHECK: %alloc = call i8* @custom_alloctor(i32 24, i32 %allocator_argument)
+; CHECK-LABEL: begin:
+; CHECK: %print_argument = load i32, i32* %argument.addr
+; CHECK: call void @print(i32 %print_argument)
+
+declare i8* @llvm.coro.free(token, i8*)
+declare i32 @llvm.coro.size.i32()
+declare i8  @llvm.coro.suspend(token, i1)
+declare void @llvm.coro.resume(i8*)
+declare void @llvm.coro.destroy(i8*)
+
+declare token @llvm.coro.id(i32, i8*, i8*, i8*)
+declare i1 @llvm.coro.alloc(token)
+declare i8* @llvm.coro.begin(token, i8*)
+declare i1 @llvm.coro.end(i8*, i1)
+
+declare noalias i8* @custom_alloctor(i32, i32)
+declare void @print(i32)
+declare void @free(i8*)
-- 
2.11.0