2 files changed, 66 insertions, 12 deletions

diff --git a/lib/Transforms/Scalar/SROA.cpp b/lib/Transforms/Scalar/SROA.cpp
index a9eb32e..860c90f 100644
--- a/lib/Transforms/Scalar/SROA.cpp
+++ b/lib/Transforms/Scalar/SROA.cpp
@@ -1869,6 +1869,10 @@ class AllocaSliceRewriter : public InstVisitor<AllocaSliceRewriter, bool> {
   Use *OldUse;
   Instruction *OldPtr;
 
+  // Output members carrying state about the result of visiting and rewriting
+  // the slice of the alloca.
+  bool IsUsedByRewrittenSpeculatableInstructions;
+
   // Utility IR builder, whose name prefix is setup for each visited use, and
   // the insertion point is set to point to the user.
   IRBuilderTy IRB;
@@ -1891,7 +1895,8 @@ public:
                         DL.getTypeSizeInBits(NewAI.getAllocatedType()))
                   : 0),
         BeginOffset(), EndOffset(), IsSplittable(), IsSplit(), OldUse(),
-        OldPtr(), IRB(NewAI.getContext(), ConstantFolder()) {
+        OldPtr(), IsUsedByRewrittenSpeculatableInstructions(false),
+        IRB(NewAI.getContext(), ConstantFolder()) {
     if (VecTy) {
       assert((DL.getTypeSizeInBits(ElementTy) % 8) == 0 &&
              "Only multiple-of-8 sized vector elements are viable");
@@ -1923,6 +1928,20 @@ public:
     return CanSROA;
   }
 
+  /// \brief Query whether this slice is used by speculatable instructions after
+  /// rewriting.
+  ///
+  /// These instructions (PHIs and Selects currently) require the alloca slice
+  /// to run back through the rewriter. Thus, they are promotable, but not on
+  /// this iteration. This is distinct from a slice which is unpromotable for
+  /// some other reason, in which case we don't even want to perform the
+  /// speculation. This can be querried at any time and reflects whether (at
+  /// that point) a visit call has rewritten a speculatable instruction on the
+  /// current slice.
+  bool isUsedByRewrittenSpeculatableInstructions() const {
+    return IsUsedByRewrittenSpeculatableInstructions;
+  }
+
 private:
   // Make sure the other visit overloads are visible.
   using Base::visit;
@@ -2553,10 +2572,12 @@ private:
     deleteIfTriviallyDead(OldPtr);
 
     // Check whether we can speculate this PHI node, and if so remember that
-    // fact and return that this alloca remains viable for promotion to an SSA
-    // value.
+    // fact and queue it up for another iteration after the speculation
+    // occurs.
     if (isSafePHIToSpeculate(PN, &DL)) {
       Pass.SpeculatablePHIs.insert(&PN);
+      Pass.Worklist.insert(&NewAI);
+      IsUsedByRewrittenSpeculatableInstructions = true;
       return true;
     }
 
@@ -2581,10 +2602,12 @@ private:
     deleteIfTriviallyDead(OldPtr);
 
     // Check whether we can speculate this select instruction, and if so
-    // remember that fact and return that this alloca remains viable for
-    // promotion to an SSA value.
+    // remember that fact and queue it up for another iteration after the
+    // speculation occurs.
     if (isSafeSelectToSpeculate(SI, &DL)) {
       Pass.SpeculatableSelects.insert(&SI);
+      Pass.Worklist.insert(&NewAI);
+      IsUsedByRewrittenSpeculatableInstructions = true;
       return true;
     }
 
@@ -3064,13 +3087,7 @@ bool SROA::rewritePartition(AllocaInst &AI, AllocaSlices &S,
       std::max<unsigned>(NumUses, MaxUsesPerAllocaPartition);
 #endif
 
-  if (Promotable && (SpeculatablePHIs.size() > SPOldSize ||
-                     SpeculatableSelects.size() > SSOldSize)) {
-    // If we have a promotable alloca except for some unspeculated loads below
-    // PHIs or Selects, iterate once. We will speculate the loads and on the
-    // next iteration rewrite them into a promotable form.
-    Worklist.insert(NewAI);
-  } else if (Promotable) {
+  if (Promotable && !Rewriter.isUsedByRewrittenSpeculatableInstructions()) {
     DEBUG(dbgs() << "  and queuing for promotion\n");
     PromotableAllocas.push_back(NewAI);
   } else if (NewAI != &AI) {
diff --git a/test/Transforms/SROA/phi-and-select.ll b/test/Transforms/SROA/phi-and-select.ll
index 8637578..9f2b191 100644
--- a/test/Transforms/SROA/phi-and-select.ll
+++ b/test/Transforms/SROA/phi-and-select.ll
@@ -427,3 +427,40 @@ if.end:
   ret i64 %result
 ; CHECK-NEXT: ret i64 %[[result]]
 }
+
+define float @PR16687(i64 %x, i1 %flag) {
+; CHECK-LABEL: @PR16687(
+; Check that even when we try to speculate the same phi twice (in two slices)
+; on an otherwise promotable construct, we don't get ahead of ourselves and try
+; to promote one of the slices prior to speculating it.
+
+entry:
+  %a = alloca i64, align 8
+  store i64 %x, i64* %a
+  br i1 %flag, label %then, label %else
+; CHECK-NOT: alloca
+; CHECK-NOT: store
+; CHECK: %[[lo:.*]] = trunc i64 %x to i32
+; CHECK: %[[shift:.*]] = lshr i64 %x, 32
+; CHECK: %[[hi:.*]] = trunc i64 %[[shift]] to i32
+
+then:
+  %a.f = bitcast i64* %a to float*
+  br label %end
+; CHECK: %[[lo_cast:.*]] = bitcast i32 %[[lo]] to float
+
+else:
+  %a.raw = bitcast i64* %a to i8*
+  %a.raw.4 = getelementptr i8* %a.raw, i64 4
+  %a.raw.4.f = bitcast i8* %a.raw.4 to float*
+  br label %end
+; CHECK: %[[hi_cast:.*]] = bitcast i32 %[[hi]] to float
+
+end:
+  %a.phi.f = phi float* [ %a.f, %then ], [ %a.raw.4.f, %else ]
+  %f = load float* %a.phi.f
+  ret float %f
+; CHECK: %[[phi:.*]] = phi float [ %[[lo_cast]], %then ], [ %[[hi_cast]], %else ]
+; CHECK-NOT: load
+; CHECK: ret float %[[phi]]
+}