PNaCl: Handle nested structure types in -expand-struct-regs.

lib/Transforms/NaCl/ExpandStructRegs.cpp

Index: lib/Transforms/NaCl/ExpandStructRegs.cpp
diff --git a/lib/Transforms/NaCl/ExpandStructRegs.cpp b/lib/Transforms/NaCl/ExpandStructRegs.cpp
index 71f7f286f82232a7b27459909b9ba373b1a42a05..000e08c5672d06c9ec56caaba565b77832b9c4a9 100644
--- a/lib/Transforms/NaCl/ExpandStructRegs.cpp
+++ b/lib/Transforms/NaCl/ExpandStructRegs.cpp
@@ -21,7 +21,6 @@
 //
 // ExpandStructRegs does not handle:
 //
-//  * Nested struct types.
 //  * Array types.
 //  * Function types containing arguments or return values of struct
 //    type without the "byval" or "sret" attributes.  Since by-value
@@ -63,20 +62,27 @@ char ExpandStructRegs::ID = 0;
 INITIALIZE_PASS(ExpandStructRegs, "expand-struct-regs",
                 "Expand out variables with struct types", false, false)
 
-static void SplitUpPHINode(PHINode *Phi) {
+static bool DoAnotherPass(Type *Ty) { return isa<StructType>(Ty); }
+static bool DoAnotherPass(Value *V) { return DoAnotherPass(V->getType()); }
+
+static bool SplitUpPHINode(PHINode *Phi) {
   StructType *STy = cast<StructType>(Phi->getType());
 
   Value *NewStruct = UndefValue::get(STy);
   Instruction *NewStructInsertPt = Phi->getParent()->getFirstInsertionPt();
 
+  bool NeedsAnotherPass = false;
+
   // 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);
+    Type *ElemTy = STy->getElementType(Index);
+    NeedsAnotherPass = NeedsAnotherPass || DoAnotherPass(ElemTy);
+
+    PHINode *NewPhi = PHINode::Create(ElemTy, Phi->getNumIncomingValues(),
+                                      Phi->getName() + ".index", Phi);
     CopyDebug(NewPhi, Phi);
     for (unsigned PhiIndex = 0; PhiIndex < Phi->getNumIncomingValues();
          ++PhiIndex) {
@@ -96,12 +102,15 @@ static void SplitUpPHINode(PHINode *Phi) {
   }
   Phi->replaceAllUsesWith(NewStruct);
   Phi->eraseFromParent();
+
+  return NeedsAnotherPass;
 }
 
-static void SplitUpSelect(SelectInst *Select) {
+static bool SplitUpSelect(SelectInst *Select) {
   StructType *STy = cast<StructType>(Select->getType());
   Value *NewStruct = UndefValue::get(STy);
 
+  bool NeedsAnotherPass = false;
   // Create a separate SelectInst for each struct field.
   for (unsigned Index = 0; Index < STy->getNumElements(); ++Index) {
     SmallVector<unsigned, 1> EVIndexes;
@@ -119,6 +128,8 @@ static void SplitUpSelect(SelectInst *Select) {
         SelectInst::Create(Select->getCondition(), TrueVal, FalseVal,
                            Select->getName() + ".index", Select), Select);
 
+    NeedsAnotherPass = NeedsAnotherPass || DoAnotherPass(NewSelect);
+
     // Reconstruct the original struct value.
     NewStruct = CopyDebug(
         InsertValueInst::Create(NewStruct, NewSelect, EVIndexes,
@@ -127,6 +138,8 @@ static void SplitUpSelect(SelectInst *Select) {
   }
   Select->replaceAllUsesWith(NewStruct);
   Select->eraseFromParent();
+
+  return NeedsAnotherPass;
 }
 
 template <class InstType>
@@ -144,8 +157,10 @@ static void ProcessLoadOrStoreAttrs(InstType *Dest, InstType *Src) {
   Dest->setAlignment(1);
 }
 
-static void SplitUpStore(StoreInst *Store) {
+static bool SplitUpStore(StoreInst *Store) {
   StructType *STy = cast<StructType>(Store->getValueOperand()->getType());
+
+  bool NeedsAnotherPass = false;
   // Create a separate store instruction for each struct field.
   for (unsigned Index = 0; Index < STy->getNumElements(); ++Index) {
     SmallVector<Value *, 2> Indexes;
@@ -155,6 +170,9 @@ static void SplitUpStore(StoreInst *Store) {
                                Store->getPointerOperand(), Indexes,
                                Store->getPointerOperand()->getName() + ".index",
                                Store), Store);
+    NeedsAnotherPass =
+        NeedsAnotherPass || DoAnotherPass(GEP->getType()->getContainedType(0));
+
     SmallVector<unsigned, 1> EVIndexes;
     EVIndexes.push_back(Index);
     Value *Field = ExtractValueInst::Create(Store->getValueOperand(),
@@ -163,12 +181,15 @@ static void SplitUpStore(StoreInst *Store) {
     ProcessLoadOrStoreAttrs(NewStore, Store);
   }
   Store->eraseFromParent();
+
+  return NeedsAnotherPass;
 }
 
-static void SplitUpLoad(LoadInst *Load) {
+static bool SplitUpLoad(LoadInst *Load) {
   StructType *STy = cast<StructType>(Load->getType());
   Value *NewStruct = UndefValue::get(STy);
 
+  bool NeedsAnotherPass = false;
   // Create a separate load instruction for each struct field.
   for (unsigned Index = 0; Index < STy->getNumElements(); ++Index) {
     SmallVector<Value *, 2> Indexes;
@@ -178,6 +199,8 @@ static void SplitUpLoad(LoadInst *Load) {
         GetElementPtrInst::Create(Load->getPointerOperand(), Indexes,
                                   Load->getName() + ".index", Load), Load);
     LoadInst *NewLoad = new LoadInst(GEP, Load->getName() + ".field", Load);
+
+    NeedsAnotherPass = NeedsAnotherPass || DoAnotherPass(NewLoad);
     ProcessLoadOrStoreAttrs(NewLoad, Load);
 
     // Reconstruct the struct value.
@@ -189,44 +212,118 @@ static void SplitUpLoad(LoadInst *Load) {
   }
   Load->replaceAllUsesWith(NewStruct);
   Load->eraseFromParent();
+
+  return NeedsAnotherPass;
 }
 
-static bool ExpandExtractValue(ExtractValueInst *EV) {
+static bool ExpandExtractValue(ExtractValueInst *EV, bool &NeedsAnotherPass) {
   // Search for the insertvalue instruction that inserts the struct field
   // referenced by this extractvalue instruction, excluding CmpXchg which
   // returns a struct and is handled by RewriteAtomics.
   Value *StructVal = EV->getAggregateOperand();
-  Value *ResultField;
+  Value *ResultField = nullptr;
+  size_t I = 0;

[Mark Seaborn, 2014/12/09 20:39:46]

"I" is too cryptic for an index used across multiple iterations in complex ways.
 Something like "EVIndex" (and "IVIndex" for "J") would be better.

What invariants apply to "I" across the loop?

[Richard Diamond, 2014/12/10 00:58:06]

Done.
I've added a comment explicating.

+
   if (isa<AtomicCmpXchgInst>(StructVal))
     return false;
+
   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");
+      size_t J = 0;
+      for (; I < EV->getIndices().size() && J < IV->getIndices().size();
+           ++J, ++I) {

[Mark Seaborn, 2014/12/09 20:39:46]

This doesn't look right.  You increment "I", but "I" never gets reset.  This
code has to search multiple InsertValueInsts until it finds a matching one. 
What happens if an InsertValueInst only partially matches?  Will "I" have been
advanced when we start checking a later InsertValueInst?

I find this new version of ExpandExtractValue() hard to understand.

[Richard Diamond, 2014/12/10 00:58:05]

Yes, see line 242 ((EVIndex + 1 == EV->getIndices().size())'s 'else' block)
I apologize; I went ahead and added many comments w/ examples explaining all the
cases handled in this section.

+        const bool Equal = (EV->getIndices()[I] == IV->getIndices()[J]);
+        if (J + 1 == IV->getIndices().size() && Equal) {
+          if (I + 1 == EV->getIndices().size()) {
+            // Match
+            ResultField = IV->getInsertedValueOperand();
+          } else {
+            StructVal = IV->getInsertedValueOperand();
+            ++I;
+          }
+          break;
+        } else if (!Equal) {
+          // No match.  Try the next struct value in the chain.
+          StructVal = IV->getAggregateOperand();
+          break;
+        }
       }
-      if (EV->getIndices()[0] == IV->getIndices()[0]) {
-        ResultField = IV->getInsertedValueOperand();
+      if (ResultField) {
+        break;
+      } else if (I == EV->getIndices().size()) {
+        // We've found an insertvalue that inserts at one or more levels deeper
+        // than this extractvalue.
+        NeedsAnotherPass = true;
+        SmallVector<unsigned, 4> Indices(IV->getIndices().begin() + J,
+                                         IV->getIndices().end());
+
+        InsertValueInst *Insert = InsertValueInst::Create(
+            UndefValue::get(EV->getType()), IV->getInsertedValueOperand(),
+            Indices, "", EV);
+        ResultField = CopyDebug(Insert, EV);
         break;
       }

[Mark Seaborn, 2014/12/09 20:39:46]

Can we reach after this "}" with StructVal not having been modified?  If so,
this would go into an infinite loop.  It's not really clear to me whether
there's an input that could cause this.

[Richard Diamond, 2014/12/10 00:58:05]

No, because all such cases would be caught by the branch at line 253 (EVIndex ==
EV->getIndices().size()) (the inner loop only breaks if it sets ResultValue,
StructVal, or EVIndex < EV->getIndices().size(). In the later case, the
aforementioned branch would be hit).

-      // 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());
+      SmallVector<unsigned, 4> Indices(EV->getIndices().begin() + I,
+                                       EV->getIndices().end());
+      ResultField = ConstantExpr::getExtractValue(C, Indices);
+      break;
+    } else if (isa<LoadInst>(StructVal)) {
+      ResultField = StructVal;
       break;
     } else {
       errs() << "Value: " << *StructVal << "\n";
       report_fatal_error("Unrecognized struct value");
     }
   }
+
+  assert(ResultField);

[Mark Seaborn, 2014/12/09 20:39:46]

Are there inputs for which this assertion would fail?

[Richard Diamond, 2014/12/10 00:58:06]

There shouldn't be, but that doesn't prevent a dev (such as myself!) from making
such a mistake in the future.

   EV->replaceAllUsesWith(ResultField);
   EV->eraseFromParent();
   return true;
 }
 
+static bool ExpandExtractValues(Function &Func) {
+  bool Changed = false;
+  bool NeedsAnotherPass = false;
+
+  // 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 (auto &BB : Func) {
+    for (BasicBlock::iterator Iter = BB.begin(), E = BB.end();
+         Iter != E; ) {
+      Instruction* Inst = Iter++;

[Mark Seaborn, 2014/12/09 20:39:46]

Nit: LLVM style is " *"

[Richard Diamond, 2014/12/10 00:58:06]

Whoops. Done.

+      if (ExtractValueInst *EV = dyn_cast<ExtractValueInst>(Inst)) {
+        Changed |= ExpandExtractValue(EV, NeedsAnotherPass);
+      } else if (isa<InsertValueInst>(Inst)) {
+        ToErase.push_back(Inst);

[Mark Seaborn, 2014/12/09 20:39:46]

ExpandExtractValue() can create InsertValueInsts, but it looks like they won't
get added to ToErase, which runs the risk of causing a failure later.

I guess this explains why you do "if (!NeedsAnotherPass)" immediately below. 
This isn't very obvious, though.

[Richard Diamond, 2014/12/10 00:58:05]

Actually, extra passes in ExpandExtractValues are impossible because
ExpandExtractValue never inserts ExtractValueInsts. I removed the loop.

+        Changed = true;
+      }
+    }
+  }
+
+  if (!NeedsAnotherPass) {
+    // 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 (Instruction *I : ToErase) {
+      I->dropAllReferences();
+    }
+
+    for (Instruction *I : ToErase) {
+      I->eraseFromParent();
+    }
+  }
+
+  return (NeedsAnotherPass && ExpandExtractValues(Func)) || Changed;
+}
+
 bool ExpandStructRegs::runOnFunction(Function &Func) {
   bool Changed = false;
+  bool NeedsAnotherPass = false;
 
   // Split up aggregate loads, stores and phi nodes into operations on

[Mark Seaborn, 2014/12/09 20:39:46]

I think there's a completeness problem here: This no longer splits up
loads/stores/phis into operations on scalar types.  Since this loop is not
iterated, doesn't it do only one level of splitting?

Shouldn't this loop be iterated so that the comment is accurate?

Or maybe each call to SplitUpStore(), SplitUpLoad() etc. should be iterated?

[Richard Diamond, 2014/12/10 00:58:05]

I was using recursion, however this patchset is old; JF had me switch to a
looped impl.

   // scalar types.  This inserts extractvalue and insertvalue
@@ -238,60 +335,33 @@ bool ExpandStructRegs::runOnFunction(Function &Func) {
       Instruction *Inst = Iter++;
       if (StoreInst *Store = dyn_cast<StoreInst>(Inst)) {
         if (Store->getValueOperand()->getType()->isStructTy()) {
-          SplitUpStore(Store);
+          NeedsAnotherPass = SplitUpStore(Store) || NeedsAnotherPass;

[JF, 2014/12/09 19:20:22]

I think this is cleaner:
  NeedsAnotherPass |= foo();

[Richard Diamond, 2014/12/09 19:39:23]

Done.

           Changed = true;
         }
       } else if (LoadInst *Load = dyn_cast<LoadInst>(Inst)) {
         if (Load->getType()->isStructTy()) {
-          SplitUpLoad(Load);
+          NeedsAnotherPass = SplitUpLoad(Load) || NeedsAnotherPass;
           Changed = true;
         }
       } else if (PHINode *Phi = dyn_cast<PHINode>(Inst)) {
         if (Phi->getType()->isStructTy()) {
-          SplitUpPHINode(Phi);
+          NeedsAnotherPass = SplitUpPHINode(Phi) || NeedsAnotherPass;
           Changed = true;
         }
       } else if (SelectInst *Select = dyn_cast<SelectInst>(Inst)) {
         if (Select->getType()->isStructTy()) {
-          SplitUpSelect(Select);
+          NeedsAnotherPass = SplitUpSelect(Select) || NeedsAnotherPass;
           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)) {
-        Changed |= ExpandExtractValue(EV);
-      } else if (isa<InsertValueInst>(Inst)) {
-        ToErase.push_back(Inst);
-        Changed = true;
-      }
-    }
+  if (!NeedsAnotherPass) {
+    Changed = ExpandExtractValues(Func) || Changed;
   }
-  // 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();
-  }
-  for (SmallVectorImpl<Instruction *>::iterator I = ToErase.begin(),
-           E = ToErase.end();
-       I != E; ++I) {
-    (*I)->eraseFromParent();
-  }
-  return Changed;
+
+  return (NeedsAnotherPass && runOnFunction(Func)) || Changed;

[JF, 2014/12/09 19:20:22]

I'm not a fan of the tail recursion, could you just make it a loop, and then the
ExpandExtractValues call can be after the loop. You could just return Changed
here.

[Richard Diamond, 2014/12/09 19:39:23]

Done.

 }
 
 FunctionPass *llvm::createExpandStructRegsPass() {