Refactor the collection of type-dependent operands

Should be NFC in impact, but some of the existing patterns can produce
redundant dependencies in probably-obscure cases, so it's not purely a
refactor.

Author: rjmccall

lib/SIL/IR/SILInstructions.cpp

@@ -43,70 +43,105 @@ static void *allocateTrailingInst(SILFunction &F, CountTypes... counts) {
              alignof(Inst));
 }
 
+namespace {
+class TypeDependentOperandCollector {
+  SmallVector<CanOpenedArchetypeType, 4> rootOpenedArchetypes;
+  bool hasDynamicSelf = false;
+public:
+  void collect(CanType type);
+  void collect(SubstitutionMap subs);
+  void collect(SILType type) {
+    collect(type.getASTType());
+  }
+  template <class T>
+  void collect(ArrayRef<T> array) {
+    for (auto &elt: array)
+      collect(elt);
+  }
+
+  void collectAll() {}
+  template <class T, class... Ts>
+  void collectAll(T &&first, Ts &&...rest) {
+    collect(first);
+    collectAll(std::forward<Ts>(rest)...);
+  }
+
+  void addTo(SmallVectorImpl<SILValue> &typeDependentOperands,
+             SILFunction &f);
+};
+
+}
+
 /// Collect root open archetypes from a given type into \p RootOpenedArchetypes.
 /// \p RootOpenedArchetypes is being used as a set. We don't use a real set type
 /// here for performance reasons.
-static void collectDependentTypeInfo(
-    CanType Ty, SmallVectorImpl<CanOpenedArchetypeType> &RootOpenedArchetypes,
-    bool &hasDynamicSelf) {
-  if (!Ty)
+void TypeDependentOperandCollector::collect(CanType type) {
+  if (!type)
     return;
-  if (Ty->hasDynamicSelfType())
+  if (type->hasDynamicSelfType())
     hasDynamicSelf = true;
-  if (!Ty->hasOpenedExistential())
+  if (!type->hasOpenedExistential())
     return;
-  Ty.visit([&](CanType t) {
+  type.visit([&](CanType t) {
     if (const auto opened = dyn_cast<OpenedArchetypeType>(t)) {
       const auto root = opened.getRoot();
 
       // Add this root opened archetype if it was not seen yet.
       // We don't use a set here, because the number of open archetypes
       // is usually very small and using a real set may introduce too
       // much overhead.
-      if (std::find(RootOpenedArchetypes.begin(), RootOpenedArchetypes.end(),
-                    root) == RootOpenedArchetypes.end())
-        RootOpenedArchetypes.push_back(root);
+      if (std::find(rootOpenedArchetypes.begin(), rootOpenedArchetypes.end(),
+                    root) == rootOpenedArchetypes.end())
+        rootOpenedArchetypes.push_back(root);
     }
   });
 }
 
-/// Takes a set of root opened archetypes as input and produces a set of
-/// references to their definitions.
-static void buildTypeDependentOperands(
-    SmallVectorImpl<CanOpenedArchetypeType> &RootOpenedArchetypes,
-    bool hasDynamicSelf, SmallVectorImpl<SILValue> &TypeDependentOperands,
-    SILFunction &F) {
+/// Collect type dependencies from the replacement types of a
+/// substitution map.
+void TypeDependentOperandCollector::collect(SubstitutionMap subs) {
+  for (Type replacement : subs.getReplacementTypes()) {
+    // Substitutions in SIL should really be canonical.
+    auto ReplTy = replacement->getCanonicalType();
+    collect(ReplTy);
+  }
+}
 
-  for (const auto &archetype : RootOpenedArchetypes) {
+/// Given that we've collected a set of type dependencies, add operands
+/// for those dependencies to the given vector.
+void TypeDependentOperandCollector::addTo(SmallVectorImpl<SILValue> &operands,
+                                          SILFunction &F) {
+  size_t firstArchetypeOperand = operands.size();
+  for (CanOpenedArchetypeType archetype : rootOpenedArchetypes) {
     SILValue def = F.getModule().getRootOpenedArchetypeDef(archetype, &F);
     assert(def->getFunction() == &F &&
            "def of root opened archetype is in wrong function");
-    TypeDependentOperands.push_back(def);
+
+    // The archetypes in rootOpenedArchetypes have already been uniqued,
+    // but a single instruction can open multiple archetypes (e.g.
+    // open_pack_element), so we also unique the actual operand values.
+    // As above, we assume there are very few values in practice and so
+    // a linear scan is better than maintaining a set.
+    if (std::find(operands.begin() + firstArchetypeOperand, operands.end(),
+                  def) == operands.end())
+      operands.push_back(def);
   }
   if (hasDynamicSelf)
-    TypeDependentOperands.push_back(F.getDynamicSelfMetadata());
+    operands.push_back(F.getDynamicSelfMetadata());
 }
 
 /// Collects all root opened archetypes from a type and a substitution list, and
 /// forms a corresponding list of operands.
 /// We need to know the number of root opened archetypes to estimate the number
 /// of corresponding operands for the instruction being formed, because we need
 /// to reserve enough memory for these operands.
+template <class... Sources>
 static void collectTypeDependentOperands(
-                      SmallVectorImpl<SILValue> &TypeDependentOperands,
-                      SILFunction &F,
-                      CanType Ty,
-                      SubstitutionMap subs = { }) {
-  SmallVector<CanOpenedArchetypeType, 4> RootOpenedArchetypes;
-  bool hasDynamicSelf = false;
-  collectDependentTypeInfo(Ty, RootOpenedArchetypes, hasDynamicSelf);
-  for (Type replacement : subs.getReplacementTypes()) {
-    // Substitutions in SIL should really be canonical.
-    auto ReplTy = replacement->getCanonicalType();
-    collectDependentTypeInfo(ReplTy, RootOpenedArchetypes, hasDynamicSelf);
-  }
-  buildTypeDependentOperands(RootOpenedArchetypes, hasDynamicSelf,
-                             TypeDependentOperands, F);
+                      SmallVectorImpl<SILValue> &typeDependentOperands,
+                      SILFunction &F, Sources &&... sources) {
+  TypeDependentOperandCollector collector;
+  collector.collectAll(std::forward<Sources>(sources)...);
+  collector.addTo(typeDependentOperands, F);
 }
 
 //===----------------------------------------------------------------------===//
@@ -284,10 +319,7 @@ AllocRefInst *AllocRefInst::create(SILDebugLocation Loc, SILFunction &F,
   assert(!objc || ElementTypes.empty());
   SmallVector<SILValue, 8> AllOperands(ElementCountOperands.begin(),
                                        ElementCountOperands.end());
-  for (SILType ElemType : ElementTypes) {
-    collectTypeDependentOperands(AllOperands, F, ElemType.getASTType());
-  }
-  collectTypeDependentOperands(AllOperands, F, ObjectType.getASTType());
+  collectTypeDependentOperands(AllOperands, F, ElementTypes, ObjectType);
   auto Size = totalSizeToAlloc<swift::Operand, SILType>(AllOperands.size(),
                                                         ElementTypes.size());
   auto Buffer = F.getModule().allocateInst(Size, alignof(AllocRefInst));
@@ -304,10 +336,7 @@ AllocRefDynamicInst::create(SILDebugLocation DebugLoc, SILFunction &F,
   SmallVector<SILValue, 8> AllOperands(ElementCountOperands.begin(),
                                        ElementCountOperands.end());
   AllOperands.push_back(metatypeOperand);
-  collectTypeDependentOperands(AllOperands, F, ty.getASTType());
-  for (SILType ElemType : ElementTypes) {
-    collectTypeDependentOperands(AllOperands, F, ElemType.getASTType());
-  }
+  collectTypeDependentOperands(AllOperands, F, ty, ElementTypes);
   auto Size = totalSizeToAlloc<swift::Operand, SILType>(AllOperands.size(),
                                                         ElementTypes.size());
   auto Buffer = F.getModule().allocateInst(Size, alignof(AllocRefDynamicInst));
@@ -1249,14 +1278,13 @@ UncheckedRefCastAddrInst *
 UncheckedRefCastAddrInst::create(SILDebugLocation Loc, SILValue src,
         CanType srcType, SILValue dest, CanType targetType, SILFunction &F) {
   SILModule &Mod = F.getModule();
-  SmallVector<SILValue, 8> TypeDependentOperands;
-  collectTypeDependentOperands(TypeDependentOperands, F, srcType);
-  collectTypeDependentOperands(TypeDependentOperands, F, targetType);
+  SmallVector<SILValue, 4> allOperands;
+  collectTypeDependentOperands(allOperands, F, srcType, targetType);
   unsigned size =
-      totalSizeToAlloc<swift::Operand>(2 + TypeDependentOperands.size());
+      totalSizeToAlloc<swift::Operand>(2 + allOperands.size());
   void *Buffer = Mod.allocateInst(size, alignof(UncheckedRefCastAddrInst));
   return ::new (Buffer) UncheckedRefCastAddrInst(Loc, src, srcType,
-    dest, targetType, TypeDependentOperands);
+    dest, targetType, allOperands);
 }
 
 UnconditionalCheckedCastAddrInst::UnconditionalCheckedCastAddrInst(
@@ -1269,14 +1297,13 @@ UnconditionalCheckedCastAddrInst *
 UnconditionalCheckedCastAddrInst::create(SILDebugLocation Loc, SILValue src,
         CanType srcType, SILValue dest, CanType targetType, SILFunction &F) {
   SILModule &Mod = F.getModule();
-  SmallVector<SILValue, 8> TypeDependentOperands;
-  collectTypeDependentOperands(TypeDependentOperands, F, srcType);
-  collectTypeDependentOperands(TypeDependentOperands, F, targetType);
+  SmallVector<SILValue, 4> allOperands;
+  collectTypeDependentOperands(allOperands, F, srcType, targetType);
   unsigned size =
-      totalSizeToAlloc<swift::Operand>(2 + TypeDependentOperands.size());
+      totalSizeToAlloc<swift::Operand>(2 + allOperands.size());
   void *Buffer = Mod.allocateInst(size, alignof(UnconditionalCheckedCastAddrInst));
   return ::new (Buffer) UnconditionalCheckedCastAddrInst(Loc, src, srcType,
-    dest, targetType, TypeDependentOperands);
+    dest, targetType, allOperands);
 }
 
 CheckedCastAddrBranchInst::CheckedCastAddrBranchInst(
@@ -1300,14 +1327,13 @@ CheckedCastAddrBranchInst::create(SILDebugLocation DebugLoc,
          ProfileCounter Target1Count, ProfileCounter Target2Count,
          SILFunction &F) {
   SILModule &Mod = F.getModule();
-  SmallVector<SILValue, 8> TypeDependentOperands;
-  collectTypeDependentOperands(TypeDependentOperands, F, srcType);
-  collectTypeDependentOperands(TypeDependentOperands, F, targetType);
+  SmallVector<SILValue, 4> allOperands;
+  collectTypeDependentOperands(allOperands, F, srcType, targetType);
   unsigned size =
-      totalSizeToAlloc<swift::Operand>(2 + TypeDependentOperands.size());
+      totalSizeToAlloc<swift::Operand>(2 + allOperands.size());
   void *Buffer = Mod.allocateInst(size, alignof(CheckedCastAddrBranchInst));
   return ::new (Buffer) CheckedCastAddrBranchInst(DebugLoc, consumptionKind,
-    src, srcType, dest, targetType, TypeDependentOperands,
+    src, srcType, dest, targetType, allOperands,
     successBB, failureBB, Target1Count, Target2Count);
 }