diff --git a/llvm/lib/Bitcode/Reader/ValueList.cpp b/llvm/lib/Bitcode/Reader/ValueList.cpp --- a/llvm/lib/Bitcode/Reader/ValueList.cpp +++ b/llvm/lib/Bitcode/Reader/ValueList.cpp @@ -8,6 +8,7 @@ #include "ValueList.h" #include "llvm/ADT/SmallVector.h" +#include "llvm/ADT/TinyPtrVector.h" #include "llvm/IR/Argument.h" #include "llvm/IR/Constant.h" #include "llvm/IR/Constants.h" @@ -148,75 +149,122 @@ /// constant. Instead of doing this, we look at all the uses and rewrite all /// the place holders at once for any constant that uses a placeholder. void BitcodeReaderValueList::resolveConstantForwardRefs() { - // Sort the values by-pointer so that they are efficient to look up with a - // binary search. - llvm::sort(ResolveConstants); - + DenseMap ConstantMapping; + ConstantMapping.reserve(ResolveConstants.size()); + + // For each constant, compute how many of its operands needs to be rewritten. + // We use these counts to drive the topological sort. + DenseMap NumRewrittenOperands; + SmallVector RewriteWorklist; + auto CountOperandsOfUsers = [&](Constant *C) { + for (Use &ConstantUse : C->uses()) { + auto *ConstantUser = dyn_cast(ConstantUse.getUser()); + if (ConstantUser && !isa(ConstantUser)) { + unsigned NumOperandsFound = ++NumRewrittenOperands[ConstantUser]; + if (NumOperandsFound == 1) + RewriteWorklist.push_back(ConstantUser); + } + } + }; + for (auto ResolveConstant : ResolveConstants) + CountOperandsOfUsers(ResolveConstant.first); + while (!RewriteWorklist.empty()) + CountOperandsOfUsers(RewriteWorklist.pop_back_val()); + + // List of constants in sorted order. + SmallVector RewrittenConstants; + // Temporary vector for operands. SmallVector NewOps; - - while (!ResolveConstants.empty()) { - Value *RealVal = operator[](ResolveConstants.back().second); - Constant *Placeholder = ResolveConstants.back().first; - ResolveConstants.pop_back(); - - // Loop over all users of the placeholder, updating them to reference the - // new value. If they reference more than one placeholder, update them all - // at once. - while (!Placeholder->use_empty()) { - auto UI = Placeholder->user_begin(); - User *U = *UI; - - // If the using object isn't uniqued, just update the operands. This - // handles instructions and initializers for global variables. - if (!isa(U) || isa(U)) { - UI.getUse().set(RealVal); - continue; + // Mapping from constants which will be rewritten to placeholders for + // those constants. + DenseMap> ConstToPlaceholders; + + // Add a constant to the sorted order, and add its uses to the worklist. + auto AddConstantToOrder = [&](Constant *C) { + RewrittenConstants.push_back(C); + for (auto &ConstantUse : C->uses()) { + auto *ConstantUser = dyn_cast(ConstantUse.getUser()); + if (ConstantUser && !isa(ConstantUser) && + --NumRewrittenOperands[ConstantUser] == 0) { + RewriteWorklist.push_back(ConstantUser); } + } + }; + + // Initialize the worklist with the roots for the topological sort: + // placeholders which don't depend on another placeholder. + for (auto ResolveConstant : ResolveConstants) { + Constant *MappedC = cast(operator[](ResolveConstant.second)); + if (NumRewrittenOperands.lookup(MappedC)) { + // We have a placeholder that resolves to a constant that needs to + // be rewritten; note it so we can process it later. + ConstToPlaceholders[MappedC].push_back(ResolveConstant.first); + } else { + AddConstantToOrder(ResolveConstant.first); + ConstantMapping[ResolveConstant.first] = MappedC; + } + } + ResolveConstants.clear(); - // Otherwise, we have a constant that uses the placeholder. Replace that - // constant with a new constant that has *all* placeholder uses updated. - Constant *UserC = cast(U); - for (User::op_iterator I = UserC->op_begin(), E = UserC->op_end(); I != E; - ++I) { - Value *NewOp; - if (!isa(*I)) { - // Not a placeholder reference. - NewOp = *I; - } else if (*I == Placeholder) { - // Common case is that it just references this one placeholder. - NewOp = RealVal; - } else { - // Otherwise, look up the placeholder in ResolveConstants. - ResolveConstantsTy::iterator It = llvm::lower_bound( - ResolveConstants, - std::pair(cast(*I), 0)); - assert(It != ResolveConstants.end() && It->first == *I); - NewOp = operator[](It->second); - } - - NewOps.push_back(cast(NewOp)); - } + // Perform the topological sort. This is Kahn's algorithm. + while (!RewriteWorklist.empty()) { + Constant *C = RewriteWorklist.pop_back_val(); - // Make the new constant. - Constant *NewC; - if (ConstantArray *UserCA = dyn_cast(UserC)) { - NewC = ConstantArray::get(UserCA->getType(), NewOps); - } else if (ConstantStruct *UserCS = dyn_cast(UserC)) { - NewC = ConstantStruct::get(UserCS->getType(), NewOps); - } else if (isa(UserC)) { - NewC = ConstantVector::get(NewOps); - } else { - assert(isa(UserC) && "Must be a ConstantExpr."); - NewC = cast(UserC)->getWithOperands(NewOps); - } + // Add this constant, and placeholders that refer to it, to the sorted + // order. Add its uses to the worklist. + AddConstantToOrder(C); + for (Constant *Placeholder : ConstToPlaceholders.lookup(C)) + AddConstantToOrder(Placeholder); + } + + // Compute the rest of the ConstantMapping map. The topological sort ensures + // we can look up the rewritten operands in ConstantMapping. + for (Constant *C : RewrittenConstants) { + if (isa(C)) { + assert(ConstantMapping.lookup(C) && "Should already be mapped"); + continue; + } - UserC->replaceAllUsesWith(NewC); - UserC->destroyConstant(); - NewOps.clear(); + // Remap the constant's operands. + for (auto &Operand : C->operands()) { + Constant *OperandVal = cast(Operand); + Constant *NewOp = ConstantMapping.lookup(OperandVal); + NewOps.push_back(NewOp ? NewOp : OperandVal); } - // Update all ValueHandles, they should be the only users at this point. - Placeholder->replaceAllUsesWith(RealVal); - delete cast(Placeholder); + // Make the new constant. + Constant *NewC; + if (ConstantArray *CA = dyn_cast(C)) { + NewC = ConstantArray::get(CA->getType(), NewOps); + } else if (ConstantStruct *CS = dyn_cast(C)) { + NewC = ConstantStruct::get(CS->getType(), NewOps); + } else if (isa(C)) { + NewC = ConstantVector::get(NewOps); + } else { + assert(isa(C) && "Must be a ConstantExpr."); + NewC = cast(C)->getWithOperands(NewOps); + } + NewOps.clear(); + + // Map this constant, and any placeholders that correspond to this + // constant, to the new value. + ConstantMapping[C] = NewC; + for (Constant *Placeholder : ConstToPlaceholders.lookup(C)) + ConstantMapping[Placeholder] = NewC; + } + + // Rewrite the uses of each constant. The topological sort ensures we never + // RAUW a constant that's used by other constants. + for (Constant *C : reverse(RewrittenConstants)) { + assert(all_of(C->users(), + [](User *U) -> bool { + return !isa(U) || isa(U); + }) && + "Rewriting constants in wrong order"); + C->replaceAllUsesWith(ConstantMapping.lookup(C)); + if (auto *Placeholder = dyn_cast(C)) + delete Placeholder; + else + C->destroyConstant(); } }