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

lib/Transforms/NaCl/ExpandStructRegs.cpp

 //===- ExpandStructRegs.cpp - Expand out variables with struct type--------===//
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===----------------------------------------------------------------------===//
 //
 // This pass expands out some uses of LLVM variables
 // (a.k.a. registers) of struct type.  It replaces loads and stores of
 // 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 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, and the SROA pass sometimes converts this code to using
 // insertvalue instructions too.
 //
 // 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
 //    struct-passing generally uses "byval"/"sret", this does not
 //    matter.
 //
 // Other limitations:
 //
 //  * ExpandStructRegs does not attempt to use memcpy() where that
 //    might be more appropriate than copying fields individually.
@@ skipping 21 matching lines @@
     }
 
     virtual bool runOnFunction(Function &F);
   };
 }
 
 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) {
       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();
+
+  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;
     EVIndexes.push_back(Index);
 
     Value *TrueVal = CopyDebug(
         ExtractValueInst::Create(Select->getTrueValue(), EVIndexes,
                                  Select->getName() + ".extract", Select),
         Select);
     Value *FalseVal = CopyDebug(
         ExtractValueInst::Create(Select->getFalseValue(), EVIndexes,
                                  Select->getName() + ".extract", Select),
         Select);
     Value *NewSelect = CopyDebug(
         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,
                                 Select->getName() + ".insert", Select),
         Select);
   }
   Select->replaceAllUsesWith(NewStruct);
   Select->eraseFromParent();
+
+  return NeedsAnotherPass;
 }
 
 template <class InstType>
 static void ProcessLoadOrStoreAttrs(InstType *Dest, InstType *Src) {
   CopyDebug(Dest, Src);
   Dest->setVolatile(Src->isVolatile());
   if (Src->isAtomic()) {
     errs() << "Use: " << *Src << "\n";
     report_fatal_error("Atomic struct loads/stores not supported");
   }
   // Make a pessimistic assumption about alignment.  Preserving
   // alignment information here is tricky and is not really desirable
   // for PNaCl because mistakes here could lead to non-portable
   // behaviour.
   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;
     Indexes.push_back(ConstantInt::get(Store->getContext(), APInt(32, 0)));
     Indexes.push_back(ConstantInt::get(Store->getContext(), APInt(32, Index)));
     Value *GEP = CopyDebug(GetElementPtrInst::Create(
                                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(),
                                             EVIndexes, "", Store);
     StoreInst *NewStore = new StoreInst(Field, GEP, 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;
     Indexes.push_back(ConstantInt::get(Load->getContext(), APInt(32, 0)));
     Indexes.push_back(ConstantInt::get(Load->getContext(), APInt(32, Index)));
     Value *GEP = CopyDebug(
         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.
     SmallVector<unsigned, 1> EVIndexes;
     EVIndexes.push_back(Index);
     NewStruct = CopyDebug(
         InsertValueInst::Create(NewStruct, NewLoad, EVIndexes,
                                 Load->getName() + ".insert", Load), 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
   // instructions which we will expand out later.
   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 (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() {
   return new ExpandStructRegs();
 }