PNaCl: Fix ExpandTls to handle a couple of corner cases involving PHI nodes

ExpandTls's use of replaceUsesOfWith() didn't work for PHI nodes
containing the same Constant twice (which needs to work for same or
differing incoming BasicBlocks).  The same applies to
ExpandTlsConstantExpr.

I noticed this while implementing ExpandConstantExpr.

Fix this and factor out some common code that all three passes can
use.

BUG=https://code.google.com/p/nativeclient/issues/detail?id=2837
TEST=test/Transforms/NaCl/*.ll

Review URL: https://codereview.chromium.org/13128002

8 files changed, 128 insertions, 41 deletions

diff --git a/include/llvm/Transforms/NaCl.h b/include/llvm/Transforms/NaCl.h
index 58166bbf22..0a857653ba 100644
--- a/include/llvm/Transforms/NaCl.h
+++ b/include/llvm/Transforms/NaCl.h
@@ -14,7 +14,10 @@ namespace llvm {
 
 class BasicBlockPass;
 class FunctionPass;
+class Instruction;
 class ModulePass;
+class Use;
+class Value;
 
 FunctionPass *createExpandConstantExprPass();
 ModulePass *createExpandCtorsPass();
@@ -26,6 +29,9 @@ ModulePass *createGlobalCleanupPass();
 ModulePass *createResolveAliasesPass();
 ModulePass *createStripMetadataPass();
 
+Instruction *PhiSafeInsertPt(Use *U);
+void PhiSafeReplaceUses(Use *U, Value *NewVal);
+
 }
 
 #endif
diff --git a/lib/Transforms/NaCl/CMakeLists.txt b/lib/Transforms/NaCl/CMakeLists.txt
index 1ca8e3d59e..5d630bb111 100644
--- a/lib/Transforms/NaCl/CMakeLists.txt
+++ b/lib/Transforms/NaCl/CMakeLists.txt
@@ -4,6 +4,7 @@ add_llvm_library(LLVMTransformsNaCl
   ExpandGetElementPtr.cpp
   ExpandTls.cpp
   ExpandTlsConstantExpr.cpp
+  ExpandUtils.cpp
   ExpandVarArgs.cpp
   GlobalCleanup.cpp
   StripMetadata.cpp
diff --git a/lib/Transforms/NaCl/ExpandConstantExpr.cpp b/lib/Transforms/NaCl/ExpandConstantExpr.cpp
index 5ff47cb00c..2856a9d7e4 100644
--- a/lib/Transforms/NaCl/ExpandConstantExpr.cpp
+++ b/lib/Transforms/NaCl/ExpandConstantExpr.cpp
@@ -63,31 +63,12 @@ static bool expandInstruction(Instruction *Inst) {
     return false;
 
   bool Modified = false;
-  if (PHINode *PN = dyn_cast<PHINode>(Inst)) {
-    // PHI nodes require special treatment.  A incoming block may be
-    // listed twice, but its incoming values must match so they must
-    // be converted only once.
-    std::map<BasicBlock*,Value*> Converted;
-    for (unsigned OpNum = 0; OpNum < Inst->getNumOperands(); OpNum++) {
-      if (ConstantExpr *Expr =
-          dyn_cast<ConstantExpr>(Inst->getOperand(OpNum))) {
-        Modified = true;
-        BasicBlock *Incoming = PN->getIncomingBlock(OpNum);
-        if (Converted.count(Incoming) == 0) {
-          Converted[Incoming] = expandConstantExpr(Incoming->getTerminator(),
-                                                   Expr);
-        }
-        Inst->setOperand(OpNum, Converted[Incoming]);
-      }
-    }
-    return Modified;
-  }
-
   for (unsigned OpNum = 0; OpNum < Inst->getNumOperands(); OpNum++) {
     if (ConstantExpr *Expr =
         dyn_cast<ConstantExpr>(Inst->getOperand(OpNum))) {
       Modified = true;
-      Inst->setOperand(OpNum, expandConstantExpr(Inst, Expr));
+      Use *U = &Inst->getOperandUse(OpNum);
+      PhiSafeReplaceUses(U, expandConstantExpr(PhiSafeInsertPt(U), Expr));
     }
   }
   return Modified;
diff --git a/lib/Transforms/NaCl/ExpandTls.cpp b/lib/Transforms/NaCl/ExpandTls.cpp
index 929b2e0a15..53e1c92a96 100644
--- a/lib/Transforms/NaCl/ExpandTls.cpp
+++ b/lib/Transforms/NaCl/ExpandTls.cpp
@@ -250,15 +250,8 @@ static void rewriteTlsVars(Module &M, std::vector<VarInfo> *TlsVars,
        ++VarInfo) {
     GlobalVariable *Var = VarInfo->TlsVar;
     while (!Var->use_empty()) {
-      Instruction *U = cast<Instruction>(*Var->use_begin());
-      Instruction *InsertPt = U;
-      if (PHINode *PN = dyn_cast<PHINode>(InsertPt)) {
-        // We cannot insert instructions before a PHI node, so insert
-        // before the incoming block's terminator.  Note that if the
-        // terminator is conditional, this could be suboptimal,
-        // because we might be calling ReadTpFunc unnecessarily.
-        InsertPt = PN->getIncomingBlock(Var->use_begin())->getTerminator();
-      }
+      Use *U = &Var->use_begin().getUse();
+      Instruction *InsertPt = PhiSafeInsertPt(U);
       Value *RawThreadPtr = CallInst::Create(ReadTpFunc, "tls_raw", InsertPt);
       Value *TypedThreadPtr = new BitCastInst(RawThreadPtr, TemplatePtrType,
                                               "tls_struct", InsertPt);
@@ -279,7 +272,7 @@ static void rewriteTlsVars(Module &M, std::vector<VarInfo> *TlsVars,
           M.getContext(), APInt(32, VarInfo->TemplateIndex)));
       Value *TlsField = GetElementPtrInst::Create(TypedThreadPtr, Indexes,
                                                   "field", InsertPt);
-      U->replaceUsesOfWith(Var, TlsField);
+      PhiSafeReplaceUses(U, TlsField);
     }
     VarInfo->TlsVar->eraseFromParent();
   }
diff --git a/lib/Transforms/NaCl/ExpandTlsConstantExpr.cpp b/lib/Transforms/NaCl/ExpandTlsConstantExpr.cpp
index 51b21cbf2a..328e5e08e6 100644
--- a/lib/Transforms/NaCl/ExpandTlsConstantExpr.cpp
+++ b/lib/Transforms/NaCl/ExpandTlsConstantExpr.cpp
@@ -78,18 +78,11 @@ static void expandConstExpr(Constant *Expr) {
 
   if (ConstantExpr *CE = dyn_cast<ConstantExpr>(Expr)) {
     while (!Expr->use_empty()) {
-      Instruction *U = cast<Instruction>(*Expr->use_begin());
-      Instruction *InsertPt = U;
-      if (PHINode *PN = dyn_cast<PHINode>(InsertPt)) {
-        // We cannot insert instructions before a PHI node, so insert
-        // before the incoming block's terminator.  This could be
-        // suboptimal if the terminator is a conditional.
-        InsertPt = PN->getIncomingBlock(Expr->use_begin())->getTerminator();
-      }
+      Use *U = &Expr->use_begin().getUse();
       Instruction *NewInst = CE->getAsInstruction();
-      NewInst->insertBefore(InsertPt);
+      NewInst->insertBefore(PhiSafeInsertPt(U));
       NewInst->setName("expanded");
-      U->replaceUsesOfWith(CE, NewInst);
+      PhiSafeReplaceUses(U, NewInst);
     }
   }
 }
diff --git a/lib/Transforms/NaCl/ExpandUtils.cpp b/lib/Transforms/NaCl/ExpandUtils.cpp
new file mode 100644
index 0000000000..0670ff75ce
--- /dev/null
+++ b/lib/Transforms/NaCl/ExpandUtils.cpp
@@ -0,0 +1,40 @@
+//===-- ExpandUtils.cpp - Helper functions for expansion passes -----------===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/IR/BasicBlock.h"
+#include "llvm/IR/Instructions.h"
+#include "llvm/Transforms/NaCl.h"
+
+using namespace llvm;
+
+Instruction *llvm::PhiSafeInsertPt(Use *U) {
+  Instruction *InsertPt = cast<Instruction>(U->getUser());
+  if (PHINode *PN = dyn_cast<PHINode>(InsertPt)) {
+    // We cannot insert instructions before a PHI node, so insert
+    // before the incoming block's terminator.  This could be
+    // suboptimal if the terminator is a conditional.
+    InsertPt = PN->getIncomingBlock(*U)->getTerminator();
+  }
+  return InsertPt;
+}
+
+void llvm::PhiSafeReplaceUses(Use *U, Value *NewVal) {
+  if (PHINode *PN = dyn_cast<PHINode>(U->getUser())) {
+    // A PHI node can have multiple incoming edges from the same
+    // block, in which case all these edges must have the same
+    // incoming value.
+    BasicBlock *BB = PN->getIncomingBlock(*U);
+    for (unsigned I = 0; I < PN->getNumIncomingValues(); ++I) {
+      if (PN->getIncomingBlock(I) == BB)
+        PN->setIncomingValue(I, NewVal);
+    }
+  } else {
+    U->getUser()->replaceUsesOfWith(U->get(), NewVal);
+  }
+}
diff --git a/test/Transforms/NaCl/expand-tls-constexpr.ll b/test/Transforms/NaCl/expand-tls-constexpr.ll
index 06bb8ed830..b7ab253692 100644
--- a/test/Transforms/NaCl/expand-tls-constexpr.ll
+++ b/test/Transforms/NaCl/expand-tls-constexpr.ll
@@ -114,3 +114,39 @@ return:
 ; CHECK: %expanded = getelementptr inbounds i32* @tvar, i32 1
 ; CHECK: return:
 ; CHECK: %result = phi i32* [ %expanded, %entry ], [ null, %else ]
+
+
+; This tests that ExpandTlsConstantExpr correctly handles a PHI node
+; that contains the same ConstantExpr twice.  Using
+; replaceAllUsesWith() is not correct on a PHI node when the new
+; instruction has to be added to an incoming block.
+define i32 @test_converting_phi_twice(i1 %arg) {
+  br i1 %arg, label %iftrue, label %iffalse
+iftrue:
+  br label %exit
+iffalse:
+  br label %exit
+exit:
+  %result = phi i32 [ ptrtoint (i32* @tvar to i32), %iftrue ],
+                    [ ptrtoint (i32* @tvar to i32), %iffalse ]
+  ret i32 %result
+}
+; CHECK: define i32 @test_converting_phi_twice(i1 %arg)
+; CHECK: iftrue:
+; CHECK: %expanded{{.*}} = ptrtoint i32* @tvar to i32
+; CHECK: iffalse:
+; CHECK: %expanded{{.*}} = ptrtoint i32* @tvar to i32
+; CHECK: exit:
+; CHECK: %result = phi i32 [ %expanded1, %iftrue ], [ %expanded, %iffalse ]
+
+
+define i32 @test_converting_phi_multiple_entry(i1 %arg) {
+entry:
+  br i1 %arg, label %done, label %done
+done:
+  %result = phi i32 [ ptrtoint (i32* @tvar to i32), %entry ],
+                    [ ptrtoint (i32* @tvar to i32), %entry ]
+  ret i32 %result
+}
+; CHECK: define i32 @test_converting_phi_multiple_entry(i1 %arg)
+; CHECK: %result = phi i32 [ %expanded, %entry ], [ %expanded, %entry ]
diff --git a/test/Transforms/NaCl/expand-tls-phi.ll b/test/Transforms/NaCl/expand-tls-phi.ll
index 0292a1d633..4aa0a7a32c 100644
--- a/test/Transforms/NaCl/expand-tls-phi.ll
+++ b/test/Transforms/NaCl/expand-tls-phi.ll
@@ -16,8 +16,45 @@ return:
 }
 ; The TLS access gets pushed back into the PHI node's incoming block,
 ; which might be suboptimal but works in all cases.
+; CHECK: define i32* @get_tvar(i1 %cmp) {
 ; CHECK: entry:
 ; CHECK: %field = getelementptr %tls_struct* %tls_struct, i32 -1, i32 0, i32 0
 ; CHECK: else:
 ; CHECK: return:
 ; CHECK: %result = phi i32* [ %field, %entry ], [ null, %else ]
+
+
+; This tests that ExpandTls correctly handles a PHI node that contains
+; the same TLS variable twice.  Using replaceAllUsesWith() is not
+; correct on a PHI node when the new instruction has to be added to an
+; incoming block.
+define i32* @tls_phi_twice(i1 %arg) {
+  br i1 %arg, label %iftrue, label %iffalse
+iftrue:
+  br label %exit
+iffalse:
+  br label %exit
+exit:
+  %result = phi i32* [ @tvar, %iftrue ], [ @tvar, %iffalse ]
+  ret i32* %result
+}
+; CHECK: define i32* @tls_phi_twice(i1 %arg) {
+; CHECK: iftrue:
+; CHECK: %field{{.*}} = getelementptr %tls_struct* %tls_struct{{.*}}, i32 -1, i32 0, i32 0
+; CHECK: iffalse:
+; CHECK: %field{{.*}} = getelementptr %tls_struct* %tls_struct{{.*}}, i32 -1, i32 0, i32 0
+; CHECK: exit:
+; CHECK: %result = phi i32* [ %field{{.*}}, %iftrue ], [ %field{{.*}}, %iffalse ]
+
+
+; In this corner case, ExpandTls must expand out @tvar only once,
+; otherwise it will produce invalid IR.
+define i32* @tls_phi_multiple_entry(i1 %arg) {
+entry:
+  br i1 %arg, label %done, label %done
+done:
+  %result = phi i32* [ @tvar, %entry ], [ @tvar, %entry ]
+  ret i32* %result
+}
+; CHECK: define i32* @tls_phi_multiple_entry(i1 %arg) {
+; CHECK: %result = phi i32* [ %field, %entry ], [ %field, %entry ]