PNaCl: Extend ExpandStructRegs to handle "insertvalue" instructions

Previously, the pass could expand out struct_load+extractvalue
combinations (via ReplaceUsesOfStructWithFields()).  This change makes
the pass more general by converting struct_loads to
scalar_load+insertvalue combinations and then by expanding out
insertvalue+extractvalue combinations.

This means the pass can now handle the insertvalue instructions that
are sometimes generated by the SROA pass.  (Clang compiles ScummVM's
use of C++ method pointers to struct loads+stores which SROA generates
insertvalue instructions from.)

To make the pass more general, we also extend it to be able to handle
phi nodes of struct type.  These are split into
extractvalue+scalar_phi+insertvalue combinations.

However, the pass is not yet fully general, because it doesn't handle:
 * nested struct types
 * array types

Change ExpandArithWithOverflow to rely on ExpandStructRegs'
insertvalue handling rather than the less general
ReplaceUsesOfStructWithFields() helper function, which can now be
removed.

SplitUpStore() no longer needs to handle Constants, because this case
is handled by ExpandExtractValue().

BUG=https://code.google.com/p/nativeclient/issues/detail?id=3476
TEST=*.ll tests + PNaCl toolchain trybots

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

7 files changed, 213 insertions, 78 deletions

diff --git a/include/llvm/Transforms/NaCl.h b/include/llvm/Transforms/NaCl.h
index e493b48c03..8632c47706 100644
--- a/include/llvm/Transforms/NaCl.h
+++ b/include/llvm/Transforms/NaCl.h
@@ -60,12 +60,6 @@ Instruction *CopyDebug(Instruction *NewInst, Instruction *Original);
 // different.
 Function *RecreateFunction(Function *Func, FunctionType *NewType);
 
-// Given a value of struct type, StructVal, this replaces all uses of
-// StructVal with the given struct fields.  This involves replacing
-// extractvalue instructions that refer to StructVal.
-void ReplaceUsesOfStructWithFields(Value *StructVal,
-                                   const SmallVectorImpl<Value *> &Fields);
-
 }
 
 #endif
diff --git a/lib/Transforms/NaCl/ExpandArithWithOverflow.cpp b/lib/Transforms/NaCl/ExpandArithWithOverflow.cpp
index 1c879d524c..7113c839d2 100644
--- a/lib/Transforms/NaCl/ExpandArithWithOverflow.cpp
+++ b/lib/Transforms/NaCl/ExpandArithWithOverflow.cpp
@@ -63,6 +63,15 @@ static uint64_t UintTypeMax(unsigned Bits) {
   return (((uint64_t) 1) << Bits) - 1;
 }
 
+static Value *CreateInsertValue(Value *StructVal, unsigned Index,
+                                Value *Field, Instruction *BasedOn) {
+  SmallVector<unsigned, 1> EVIndexes;
+  EVIndexes.push_back(Index);
+  return CopyDebug(InsertValueInst::Create(
+                       StructVal, Field, EVIndexes,
+                       BasedOn->getName() + ".insert", BasedOn), BasedOn);
+}
+
 static bool ExpandOpForIntSize(Module *M, unsigned Bits, bool Mul) {
   IntegerType *IntTy = IntegerType::get(M->getContext(), Bits);
   SmallVector<Type *, 1> Types;
@@ -94,10 +103,9 @@ static bool ExpandOpForIntSize(Module *M, unsigned Bits, bool Mul) {
                          "*.with.overflow must be a constant");
     }
 
-    SmallVector<Value *, 2> Fields;
-    Fields.push_back(BinaryOperator::Create(
+    Value *ArithResult = BinaryOperator::Create(
         (Mul ? Instruction::Mul : Instruction::Add), VariableArg, ConstantArg,
-        Call->getName() + ".arith", Call));
+        Call->getName() + ".arith", Call);
 
     uint64_t ArgMax;
     if (Mul) {
@@ -105,10 +113,15 @@ static bool ExpandOpForIntSize(Module *M, unsigned Bits, bool Mul) {
     } else {
       ArgMax = UintTypeMax(Bits) - ConstantArg->getZExtValue();
     }
-    Fields.push_back(new ICmpInst(
+    Value *OverflowResult = new ICmpInst(
         Call, CmpInst::ICMP_UGT, VariableArg, ConstantInt::get(IntTy, ArgMax),
-        Call->getName() + ".overflow"));
-    ReplaceUsesOfStructWithFields(Call, Fields);
+        Call->getName() + ".overflow");
+
+    // Construct the struct result.
+    Value *NewStruct = UndefValue::get(Call->getType());
+    NewStruct = CreateInsertValue(NewStruct, 0, ArithResult, Call);
+    NewStruct = CreateInsertValue(NewStruct, 1, OverflowResult, Call);
+    Call->replaceAllUsesWith(NewStruct);
     Call->eraseFromParent();
   }
   Intrinsic->eraseFromParent();
diff --git a/lib/Transforms/NaCl/ExpandStructRegs.cpp b/lib/Transforms/NaCl/ExpandStructRegs.cpp
index eacad1f7f4..ac83d33e78 100644
--- a/lib/Transforms/NaCl/ExpandStructRegs.cpp
+++ b/lib/Transforms/NaCl/ExpandStructRegs.cpp
@@ -12,17 +12,17 @@
 // structs with separate loads and stores of the structs' fields.  The
 // motivation is to omit struct types from PNaCl's stable ABI.
 //
-// ExpandStructRegs does not handle all possible uses of struct
-// values.  It is only intended to handle the uses that Clang
-// generates.  Clang generates struct loads and stores, along with
+// ExpandStructRegs does not yet handle all possible uses of struct
+// values.  It is intended to handle the uses that Clang and the SROA
+// pass generate.  Clang generates struct loads and stores, along with
 // extractvalue instructions, in its implementation of C++ method
-// pointers.
+// pointers, and the SROA pass sometimes converts this code to using
+// insertvalue instructions too.
 //
 // ExpandStructRegs does not handle:
 //
-//  * The insertvalue instruction, which does not appear to be
-//    generated anywhere.
-//  * PHI nodes of struct type.
+//  * Nested struct types.
+//  * Array types.
 //  * Function types containing arguments or return values of struct
 //    type without the "byval" or "sret" attributes.  Since by-value
 //    struct-passing generally uses "byval"/"sret", this does not
@@ -62,6 +62,41 @@ char ExpandStructRegs::ID = 0;
 INITIALIZE_PASS(ExpandStructRegs, "expand-struct-regs",
                 "Expand out variables with struct types", false, false)
 
+static void SplitUpPHINode(PHINode *Phi) {
+  StructType *STy = cast<StructType>(Phi->getType());
+
+  Value *NewStruct = UndefValue::get(STy);
+  Instruction *NewStructInsertPt = Phi->getParent()->getFirstInsertionPt();
+
+  // Create a separate PHINode for each struct field.
+  for (unsigned Index = 0; Index < STy->getNumElements(); ++Index) {
+    SmallVector<unsigned, 1> EVIndexes;
+    EVIndexes.push_back(Index);
+
+    PHINode *NewPhi = PHINode::Create(
+        STy->getElementType(Index), Phi->getNumIncomingValues(),
+        Phi->getName() + ".index", Phi);
+    CopyDebug(NewPhi, Phi);
+    for (unsigned PhiIndex = 0; PhiIndex < Phi->getNumIncomingValues();
+         ++PhiIndex) {
+      BasicBlock *IncomingBB = Phi->getIncomingBlock(PhiIndex);
+      Value *EV = CopyDebug(
+          ExtractValueInst::Create(
+              Phi->getIncomingValue(PhiIndex), EVIndexes,
+              Phi->getName() + ".extract", IncomingBB->getTerminator()), Phi);
+      NewPhi->addIncoming(EV, IncomingBB);
+    }
+
+    // Reconstruct the original struct value.
+    NewStruct = CopyDebug(
+        InsertValueInst::Create(NewStruct, NewPhi, EVIndexes,
+                                Phi->getName() + ".insert", NewStructInsertPt),
+        Phi);
+  }
+  Phi->replaceAllUsesWith(NewStruct);
+  Phi->eraseFromParent();
+}
+
 template <class InstType>
 static void ProcessLoadOrStoreAttrs(InstType *Dest, InstType *Src) {
   CopyDebug(Dest, Src);
@@ -77,7 +112,7 @@ static void ProcessLoadOrStoreAttrs(InstType *Dest, InstType *Src) {
   Dest->setAlignment(1);
 }
 
-static void ExpandStore(StoreInst *Store) {
+static void SplitUpStore(StoreInst *Store) {
   StructType *STy = cast<StructType>(Store->getValueOperand()->getType());
   // Create a separate store instruction for each struct field.
   for (unsigned Index = 0; Index < STy->getNumElements(); ++Index) {
@@ -90,23 +125,19 @@ static void ExpandStore(StoreInst *Store) {
                                Store), Store);
     SmallVector<unsigned, 1> EVIndexes;
     EVIndexes.push_back(Index);
-    Value *Field;
-    if (Constant *C = dyn_cast<Constant>(Store->getValueOperand())) {
-      Field = ConstantExpr::getExtractValue(C, EVIndexes);
-    } else {
-      Field = ExtractValueInst::Create(
-          Store->getValueOperand(), EVIndexes, "", Store);
-    }
+    Value *Field = ExtractValueInst::Create(Store->getValueOperand(),
+                                            EVIndexes, "", Store);
     StoreInst *NewStore = new StoreInst(Field, GEP, Store);
     ProcessLoadOrStoreAttrs(NewStore, Store);
   }
   Store->eraseFromParent();
 }
 
-static void ExpandLoad(LoadInst *Load) {
+static void SplitUpLoad(LoadInst *Load) {
   StructType *STy = cast<StructType>(Load->getType());
+  Value *NewStruct = UndefValue::get(STy);
+
   // Create a separate load instruction for each struct field.
-  SmallVector<Value *, 5> Fields;
   for (unsigned Index = 0; Index < STy->getNumElements(); ++Index) {
     SmallVector<Value *, 2> Indexes;
     Indexes.push_back(ConstantInt::get(Load->getContext(), APInt(32, 0)));
@@ -116,49 +147,109 @@ static void ExpandLoad(LoadInst *Load) {
                                   Load->getName() + ".index", Load), Load);
     LoadInst *NewLoad = new LoadInst(GEP, Load->getName() + ".field", Load);
     ProcessLoadOrStoreAttrs(NewLoad, Load);
-    Fields.push_back(NewLoad);
+
+    // Reconstruct the struct value.
+    SmallVector<unsigned, 1> EVIndexes;
+    EVIndexes.push_back(Index);
+    NewStruct = CopyDebug(
+        InsertValueInst::Create(NewStruct, NewLoad, EVIndexes,
+                                Load->getName() + ".insert", Load), Load);
   }
-  ReplaceUsesOfStructWithFields(Load, Fields);
+  Load->replaceAllUsesWith(NewStruct);
   Load->eraseFromParent();
 }
 
+static void ExpandExtractValue(ExtractValueInst *EV) {
+  // Search for the insertvalue instruction that inserts the struct
+  // field referenced by this extractvalue instruction.
+  Value *StructVal = EV->getAggregateOperand();
+  Value *ResultField;
+  for (;;) {
+    if (InsertValueInst *IV = dyn_cast<InsertValueInst>(StructVal)) {
+      if (EV->getNumIndices() != 1 || IV->getNumIndices() != 1) {
+        errs() << "Value: " << *EV << "\n";
+        errs() << "Value: " << *IV << "\n";
+        report_fatal_error("ExpandStructRegs does not handle nested structs");
+      }
+      if (EV->getIndices()[0] == IV->getIndices()[0]) {
+        ResultField = IV->getInsertedValueOperand();
+        break;
+      }
+      // No match.  Try the next struct value in the chain.
+      StructVal = IV->getAggregateOperand();
+    } else if (Constant *C = dyn_cast<Constant>(StructVal)) {
+      ResultField = ConstantExpr::getExtractValue(C, EV->getIndices());
+      break;
+    } else {
+      errs() << "Value: " << *StructVal << "\n";
+      report_fatal_error("Unrecognized struct value");
+    }
+  }
+  EV->replaceAllUsesWith(ResultField);
+  EV->eraseFromParent();
+}
+
 bool ExpandStructRegs::runOnFunction(Function &Func) {
   bool Changed = false;
 
-  // It is not safe to iterate through the basic block while
-  // deleting extractvalue instructions, so make a copy of the
-  // instructions we will operate on first.
-  SmallVector<StoreInst *, 10> Stores;
-  SmallVector<LoadInst *, 10> Loads;
+  // Split up aggregate loads, stores and phi nodes into operations on
+  // scalar types.  This inserts extractvalue and insertvalue
+  // instructions which we will expand out later.
   for (Function::iterator BB = Func.begin(), E = Func.end();
        BB != E; ++BB) {
-    for (BasicBlock::iterator Inst = BB->begin(), E = BB->end();
-         Inst != E; ++Inst) {
+    for (BasicBlock::iterator Iter = BB->begin(), E = BB->end();
+         Iter != E; ) {
+      Instruction *Inst = Iter++;
       if (StoreInst *Store = dyn_cast<StoreInst>(Inst)) {
         if (Store->getValueOperand()->getType()->isStructTy()) {
-          Stores.push_back(Store);
+          SplitUpStore(Store);
+          Changed = true;
         }
       } else if (LoadInst *Load = dyn_cast<LoadInst>(Inst)) {
         if (Load->getType()->isStructTy()) {
-          Loads.push_back(Load);
+          SplitUpLoad(Load);
+          Changed = true;
+        }
+      } else if (PHINode *Phi = dyn_cast<PHINode>(Inst)) {
+        if (Phi->getType()->isStructTy()) {
+          SplitUpPHINode(Phi);
+          Changed = true;
         }
       }
     }
   }
 
-  // Expand stores first.  This will introduce extractvalue
-  // instructions.
-  for (SmallVectorImpl<StoreInst *>::iterator Inst = Stores.begin(),
-         E = Stores.end(); Inst != E; ++Inst) {
-    ExpandStore(*Inst);
-    Changed = true;
+  // Expand out all the extractvalue instructions.  Also collect up
+  // the insertvalue instructions for later deletion so that we do not
+  // need to make extra passes across the whole function.
+  SmallVector<Instruction *, 10> ToErase;
+  for (Function::iterator BB = Func.begin(), E = Func.end();
+       BB != E; ++BB) {
+    for (BasicBlock::iterator Iter = BB->begin(), E = BB->end();
+         Iter != E; ) {
+      Instruction *Inst = Iter++;
+      if (ExtractValueInst *EV = dyn_cast<ExtractValueInst>(Inst)) {
+        ExpandExtractValue(EV);
+        Changed = true;
+      } else if (isa<InsertValueInst>(Inst)) {
+        ToErase.push_back(Inst);
+        Changed = true;
+      }
+    }
+  }
+  // Delete the insertvalue instructions.  These can reference each
+  // other, so we must do dropAllReferences() before doing
+  // eraseFromParent(), otherwise we will try to erase instructions
+  // that are still referenced.
+  for (SmallVectorImpl<Instruction *>::iterator I = ToErase.begin(),
+           E = ToErase.end();
+       I != E; ++I) {
+    (*I)->dropAllReferences();
   }
-  // Expanding loads will remove the extractvalue instructions we
-  // previously introduced.
-  for (SmallVectorImpl<LoadInst *>::iterator Inst = Loads.begin(),
-         E = Loads.end(); Inst != E; ++Inst) {
-    ExpandLoad(*Inst);
-    Changed = true;
+  for (SmallVectorImpl<Instruction *>::iterator I = ToErase.begin(),
+           E = ToErase.end();
+       I != E; ++I) {
+    (*I)->eraseFromParent();
   }
   return Changed;
 }
diff --git a/lib/Transforms/NaCl/ExpandUtils.cpp b/lib/Transforms/NaCl/ExpandUtils.cpp
index 75f5dad114..6914b6dd84 100644
--- a/lib/Transforms/NaCl/ExpandUtils.cpp
+++ b/lib/Transforms/NaCl/ExpandUtils.cpp
@@ -59,27 +59,3 @@ Function *llvm::RecreateFunction(Function *Func, FunctionType *NewType) {
                                Func->getFunctionType()->getPointerTo()));
   return NewFunc;
 }
-
-void llvm::ReplaceUsesOfStructWithFields(
-    Value *StructVal, const SmallVectorImpl<Value *> &Fields) {
-  while (!StructVal->use_empty()) {
-    User *U = StructVal->use_back();
-    ExtractValueInst *Field = dyn_cast<ExtractValueInst>(U);
-    if (!Field) {
-      errs() << "Use: " << *U << "\n";
-      report_fatal_error("ReplaceUsesOfStructWithFields: "
-                         "Struct use site is not an extractvalue");
-    }
-    if (Field->getNumIndices() != 1) {
-      // If we wanted to handle this case, we could split the
-      // extractvalue into two extractvalues and run ExpandLoad()
-      // multiple times.
-      errs() << "Use: " << *U << "\n";
-      report_fatal_error("ReplaceUsesOfStructWithFields: Unexpected indices");
-    }
-    unsigned Index = Field->getIndices()[0];
-    assert(Index < Fields.size());
-    Field->replaceAllUsesWith(Fields[Index]);
-    Field->eraseFromParent();
-  }
-}
diff --git a/lib/Transforms/NaCl/PNaClABISimplify.cpp b/lib/Transforms/NaCl/PNaClABISimplify.cpp
index db9f9e6ebf..e50112c741 100644
--- a/lib/Transforms/NaCl/PNaClABISimplify.cpp
+++ b/lib/Transforms/NaCl/PNaClABISimplify.cpp
@@ -30,6 +30,9 @@ void llvm::PNaClABISimplifyAddPreOptPasses(PassManager &PM) {
 
   // Expand out some uses of struct types.
   PM.add(createExpandArithWithOverflowPass());
+  // ExpandStructRegs must be run after ExpandArithWithOverflow to
+  // expand out the insertvalue instructions that
+  // ExpandArithWithOverflow introduces.
   PM.add(createExpandStructRegsPass());
 
   PM.add(createExpandVarArgsPass());
diff --git a/test/Transforms/NaCl/expand-arith-with-overflow.ll b/test/Transforms/NaCl/expand-arith-with-overflow.ll
index d3bb8f0392..bff6388e11 100644
--- a/test/Transforms/NaCl/expand-arith-with-overflow.ll
+++ b/test/Transforms/NaCl/expand-arith-with-overflow.ll
@@ -1,10 +1,14 @@
-; RUN: opt %s -expand-arith-with-overflow -S | FileCheck %s
+; RUN: opt %s -expand-arith-with-overflow -expand-struct-regs -S | FileCheck %s
+; RUN: opt %s -expand-arith-with-overflow -expand-struct-regs -S | \
+; RUN:     FileCheck %s -check-prefix=CLEANUP
 
 declare {i32, i1} @llvm.umul.with.overflow.i32(i32, i32)
 declare {i64, i1} @llvm.umul.with.overflow.i64(i64, i64)
 declare {i16, i1} @llvm.uadd.with.overflow.i16(i16, i16)
 
-; CHECK-NOT: with.overflow
+; CLEANUP-NOT: with.overflow
+; CLEANUP-NOT: extractvalue
+; CLEANUP-NOT: insertvalue
 
 
 define void @umul32_by_const(i32 %x, i32* %result_val, i1* %result_overflow) {
diff --git a/test/Transforms/NaCl/expand-struct-regs.ll b/test/Transforms/NaCl/expand-struct-regs.ll
index a81176b5e6..8291ec5069 100644
--- a/test/Transforms/NaCl/expand-struct-regs.ll
+++ b/test/Transforms/NaCl/expand-struct-regs.ll
@@ -1,4 +1,9 @@
 ; RUN: opt %s -expand-struct-regs -S | FileCheck %s
+; RUN: opt %s -expand-struct-regs -S | FileCheck %s -check-prefix=CLEANUP
+
+; These two instructions should not appear in the output:
+; CLEANUP-NOT: extractvalue
+; CLEANUP-NOT: insertvalue
 
 %struct = type { i8, i32 }
 
@@ -59,3 +64,52 @@ define void @const_struct_store(%struct* %ptr) {
 ; CHECK: define void @const_struct_store
 ; CHECK: store i8 99
 ; CHECK: store i32 1234
+
+
+define void @struct_phi_node(%struct* %ptr) {
+entry:
+  %val = load %struct* %ptr
+  br label %bb
+bb:
+  %phi = phi %struct [ %val, %entry ]
+  ret void
+}
+; CHECK: bb:
+; CHECK-NEXT: %phi.index{{.*}} = phi i8 [ %val.field{{.*}}, %entry ]
+; CHECK-NEXT: %phi.index{{.*}} = phi i32 [ %val.field{{.*}}, %entry ]
+
+
+define void @struct_phi_node_multiple_entry(i1 %arg, %struct* %ptr) {
+entry:
+  %val = load %struct* %ptr
+  br i1 %arg, label %bb, label %bb
+bb:
+  %phi = phi %struct [ %val, %entry ], [ %val, %entry ]
+  ret void
+}
+; CHECK: bb:
+; CHECK-NEXT: %phi.index{{.*}} = phi i8 [ %val.field{{.*}}, %entry ], [ %val.field{{.*}}, %entry ]
+; CHECK-NEXT: %phi.index{{.*}} = phi i32 [ %val.field{{.*}}, %entry ], [ %val.field{{.*}}, %entry ]
+
+
+define void @insert_and_extract(i8* %out0, i32* %out1) {
+  %temp = insertvalue %struct undef, i8 100, 0
+  %sval = insertvalue %struct %temp, i32 200, 1
+  %field0 = extractvalue %struct %sval, 0
+  %field1 = extractvalue %struct %sval, 1
+  store i8 %field0, i8* %out0
+  store i32 %field1, i32* %out1
+  ret void
+}
+; CHECK: define void @insert_and_extract(i8* %out0, i32* %out1) {
+; CHECK-NEXT: store i8 100, i8* %out0
+; CHECK-NEXT: store i32 200, i32* %out1
+; CHECK-NEXT: ret void
+
+
+define i32 @extract_from_constant() {
+  %ev = extractvalue %struct { i8 99, i32 888 }, 1
+  ret i32 %ev
+}
+; CHECK: define i32 @extract_from_constant() {
+; CHECK-NEXT: ret i32 888