Merge pull request #15554 from jrose-apple/SILGenPattern-cleanup

Author: jrose-apple

lib/SILGen/SILGenPattern.cpp

@@ -1690,32 +1690,63 @@ void PatternMatchEmission::emitIsDispatch(ArrayRef<RowToSpecialize> rows,
 
 namespace {
   struct CaseInfo {
-    EnumElementDecl *FormalElement;
+    SmallVector<SpecializedRow, 2> SpecializedRows;
     Pattern *FirstMatcher;
     bool Irrefutable = false;
-    SmallVector<SpecializedRow, 2> SpecializedRows;
+  };
+
+  class CaseBlocks {
+    // These vectors are completely parallel, but the switch instructions want
+    // only the first two, so we split them up.
+    SmallVector<std::pair<EnumElementDecl *, SILBasicBlock *>, 4> CaseBBs;
+    SmallVector<ProfileCounter, 4> CaseCounts;
+    SmallVector<CaseInfo, 4> CaseInfos;
+    SILBasicBlock *DefaultBB = nullptr;
+
+  public:
+    /// Create destination blocks for switching over the cases in an enum
+    /// defined by \p rows.
+    CaseBlocks(SILGenFunction &SGF,
+               ArrayRef<RowToSpecialize> rows,
+               CanType sourceType,
+               SILBasicBlock *curBB);
+
+    ArrayRef<std::pair<EnumElementDecl *, SILBasicBlock *>>
+    getCaseBlocks() const {
+      return CaseBBs;
+    }
+
+    ArrayRef<ProfileCounter> getCounts() const { return CaseCounts; }
+
+    SILBasicBlock *getDefaultBlock() const { return DefaultBB; }
+
+    void forEachCase(llvm::function_ref<void(EnumElementDecl *,
+                                             SILBasicBlock *,
+                                             const CaseInfo &)> op) const {
+      for_each(CaseBBs, CaseInfos,
+               [op](std::pair<EnumElementDecl *, SILBasicBlock *> casePair,
+                    const CaseInfo &info) {
+        op(casePair.first, casePair.second, info);
+      });
+    }
+
+    bool hasAnyRefutableCase() const {
+      return llvm::any_of(CaseInfos, [](const CaseInfo &info) {
+        return !info.Irrefutable;
+      });
+    }
   };
 } // end anonymous namespace
 
-/// Create destination blocks for switching over the cases in an enum defined
-/// by \p rows.
-static void generateEnumCaseBlocks(
+CaseBlocks::CaseBlocks(
     SILGenFunction &SGF,
     ArrayRef<RowToSpecialize> rows,
     CanType sourceType,
-    SILBasicBlock *curBB,
-    SmallVectorImpl<std::pair<EnumElementDecl *, SILBasicBlock *>> &caseBBs,
-    SmallVectorImpl<ProfileCounter> &caseCounts,
-    SmallVectorImpl<CaseInfo> &caseInfos,
-    SILBasicBlock *&defaultBB) {
+    SILBasicBlock *curBB) {
 
-  assert(caseBBs.empty());
-  assert(caseCounts.empty());
-  assert(caseInfos.empty());
-  assert(defaultBB == nullptr);
-
-  caseBBs.reserve(rows.size());
-  caseInfos.reserve(rows.size());
+  CaseBBs.reserve(rows.size());
+  CaseInfos.reserve(rows.size());
+  CaseCounts.reserve(rows.size());
 
   auto enumDecl = sourceType.getEnumOrBoundGenericEnum();
 
@@ -1733,22 +1764,21 @@ static void generateEnumCaseBlocks(
     }
     assert(formalElt->getParentEnum() == enumDecl);
 
-    unsigned index = caseInfos.size();
+    unsigned index = CaseInfos.size();
     auto insertionResult = caseToIndex.insert({formalElt, index});
     if (!insertionResult.second) {
       index = insertionResult.first->second;
     } else {
       curBB = SGF.createBasicBlock(curBB);
-      caseBBs.push_back({formalElt, curBB});
-      caseInfos.push_back(CaseInfo());
-      caseInfos.back().FormalElement = formalElt;
-      caseInfos.back().FirstMatcher = row.Pattern;
-      caseCounts.push_back(row.Count);
+      CaseBBs.push_back({formalElt, curBB});
+      CaseInfos.push_back(CaseInfo());
+      CaseInfos.back().FirstMatcher = row.Pattern;
+      CaseCounts.push_back(row.Count);
     }
     assert(caseToIndex[formalElt] == index);
-    assert(caseBBs[index].first == formalElt);
+    assert(CaseBBs[index].first == formalElt);
 
-    auto &info = caseInfos[index];
+    auto &info = CaseInfos[index];
     info.Irrefutable = (info.Irrefutable || row.Irrefutable);
     info.SpecializedRows.push_back(SpecializedRow());
     auto &specRow = info.SpecializedRows.back();
@@ -1765,7 +1795,7 @@ static void generateEnumCaseBlocks(
     }
   }
 
-  assert(caseBBs.size() == caseInfos.size());
+  assert(CaseBBs.size() == CaseInfos.size());
 
   // Check to see if the enum may have values beyond the cases we can see
   // at compile-time. This includes future cases (for resilient enums) and
@@ -1789,7 +1819,7 @@ static void generateEnumCaseBlocks(
   }
 
   if (!canAssumeExhaustive)
-    defaultBB = SGF.createBasicBlock(curBB);
+    DefaultBB = SGF.createBasicBlock(curBB);
 }
 
 /// Perform specialized dispatch for a sequence of EnumElementPattern or an
@@ -1804,37 +1834,23 @@ void PatternMatchEmission::emitEnumElementDispatchWithOwnership(
   CanType sourceType = rows[0].Pattern->getType()->getCanonicalType();
 
   // Collect the cases and specialized rows.
-  //
-  // These vectors are completely parallel, but the switch
-  // instructions want only the first information, so we split them up.
-  SmallVector<std::pair<EnumElementDecl *, SILBasicBlock *>, 4> caseBBs;
-  SmallVector<ProfileCounter, 4> caseCounts;
-  SmallVector<CaseInfo, 4> caseInfos;
-  SILBasicBlock *defaultBB = nullptr;
-
-  generateEnumCaseBlocks(SGF, rows, sourceType, SGF.B.getInsertionBB(),
-                         caseBBs, caseCounts, caseInfos, defaultBB);
+  CaseBlocks blocks{SGF, rows, sourceType, SGF.B.getInsertionBB()};
 
   SILLocation loc = PatternMatchStmt;
   loc.setDebugLoc(rows[0].Pattern);
   // SEMANTIC SIL TODO: Once we have the representation of a switch_enum that
   // can take a +0 value, this extra copy should be a borrow.
   SILValue srcValue = src.getFinalManagedValue().copy(SGF, loc).forward(SGF);
-  // FIXME: Pass caseCounts in here as well, as it is in
-  // emitEnumElementDispatch.
-  ArrayRef<ProfileCounter> caseCountsArrayRef = caseCounts;
-  SGF.B.createSwitchEnum(loc, srcValue, defaultBB, caseBBs, caseCountsArrayRef,
+  SGF.B.createSwitchEnum(loc, srcValue, blocks.getDefaultBlock(),
+                         blocks.getCaseBlocks(), blocks.getCounts(),
                          defaultCastCount);
 
   // Okay, now emit all the cases.
-  for (unsigned i = 0, e = caseInfos.size(); i != e; ++i) {
-    auto &caseInfo = caseInfos[i];
+  blocks.forEachCase([&](EnumElementDecl *elt, SILBasicBlock *caseBB,
+                         const CaseInfo &caseInfo) {
     SILLocation loc = caseInfo.FirstMatcher;
     auto &specializedRows = caseInfo.SpecializedRows;
 
-    EnumElementDecl *elt = caseBBs[i].first;
-    EnumElementDecl *formalElt = caseInfos[i].FormalElement;
-    SILBasicBlock *caseBB = caseBBs[i].second;
     SGF.B.setInsertionPoint(caseBB);
 
     // We're in conditionally-executed code; enter a scope.
@@ -1908,8 +1924,8 @@ void PatternMatchEmission::emitEnumElementDispatchWithOwnership(
 
       CanType substEltTy =
           sourceType
-              ->getTypeOfMember(SGF.SGM.M.getSwiftModule(), formalElt,
-                                formalElt->getArgumentInterfaceType())
+              ->getTypeOfMember(SGF.SGM.M.getSwiftModule(), elt,
+                                elt->getArgumentInterfaceType())
               ->getCanonicalType();
 
       AbstractionPattern origEltTy =
@@ -1923,10 +1939,10 @@ void PatternMatchEmission::emitEnumElementDispatchWithOwnership(
 
     handleCase(eltCMV, specializedRows, outerFailure);
     assert(!SGF.B.hasValidInsertionPoint() && "did not end block");
-  }
+  });
 
   // Emit the default block if we needed one.
-  if (defaultBB) {
+  if (SILBasicBlock *defaultBB = blocks.getDefaultBlock()) {
     SGF.B.setInsertionPoint(defaultBB);
     SGF.B.createOwnedPHIArgument(src.getType());
     outerFailure(rows.back().Pattern);
@@ -1949,16 +1965,7 @@ void PatternMatchEmission::emitEnumElementDispatch(
   CanType sourceType = rows[0].Pattern->getType()->getCanonicalType();
 
   // Collect the cases and specialized rows.
-  //
-  // These vectors are completely parallel, but the switch
-  // instructions want only the first information, so we split them up.
-  SmallVector<std::pair<EnumElementDecl*, SILBasicBlock*>, 4> caseBBs;
-  SmallVector<CaseInfo, 4> caseInfos;
-  SmallVector<ProfileCounter, 4> caseCounts;
-  SILBasicBlock *defaultBB = nullptr;
-
-  generateEnumCaseBlocks(SGF, rows, sourceType, SGF.B.getInsertionBB(),
-                         caseBBs, caseCounts, caseInfos, defaultBB);
+  CaseBlocks blocks{SGF, rows, sourceType, SGF.B.getInsertionBB()};
 
   // Emit the switch_enum{_addr} instruction.
   bool addressOnlyEnum = src.getType().isAddress();
@@ -1977,12 +1984,9 @@ void PatternMatchEmission::emitEnumElementDispatch(
       assert(!SGF.F.getModule().getOptions().EnableSILOwnership &&
              "TakeOnSuccess is not supported when compiling with ownership");
       // If any of the specialization cases is refutable, we must copy.
-      for (auto caseInfo : caseInfos)
-        if (!caseInfo.Irrefutable)
-          goto refutable;
-      break;
+      if (!blocks.hasAnyRefutableCase())
+        break;
 
-    refutable:
       src = ConsumableManagedValue(ManagedValue::forUnmanaged(src.getValue()),
                                    CastConsumptionKind::CopyOnSuccess);
       break;
@@ -1992,24 +1996,22 @@ void PatternMatchEmission::emitEnumElementDispatch(
   SILValue srcValue = src.getFinalManagedValue().forward(SGF);
   SILLocation loc = PatternMatchStmt;
   loc.setDebugLoc(rows[0].Pattern);
-  ArrayRef<ProfileCounter> caseCountsArrayRef = caseCounts;
   if (addressOnlyEnum) {
-    SGF.B.createSwitchEnumAddr(loc, srcValue, defaultBB, caseBBs,
-                               caseCountsArrayRef, defaultCaseCount);
+    SGF.B.createSwitchEnumAddr(loc, srcValue, blocks.getDefaultBlock(),
+                               blocks.getCaseBlocks(), blocks.getCounts(),
+                               defaultCaseCount);
   } else {
-    SGF.B.createSwitchEnum(loc, srcValue, defaultBB, caseBBs,
-                           caseCountsArrayRef, defaultCaseCount);
+    SGF.B.createSwitchEnum(loc, srcValue, blocks.getDefaultBlock(),
+                           blocks.getCaseBlocks(), blocks.getCounts(),
+                           defaultCaseCount);
   }
 
   // Okay, now emit all the cases.
-  for (unsigned i = 0, e = caseInfos.size(); i != e; ++i) {
-    auto &caseInfo = caseInfos[i];
+  blocks.forEachCase([&](EnumElementDecl *elt, SILBasicBlock *caseBB,
+                         const CaseInfo &caseInfo) {
     SILLocation loc = caseInfo.FirstMatcher;
     auto &specializedRows = caseInfo.SpecializedRows;
 
-    EnumElementDecl *elt = caseBBs[i].first;
-    EnumElementDecl *formalElt = caseInfos[i].FormalElement;
-    SILBasicBlock *caseBB = caseBBs[i].second;
     SGF.B.setInsertionPoint(caseBB);
 
     // We're in conditionally-executed code; enter a scope.
@@ -2131,8 +2133,8 @@ void PatternMatchEmission::emitEnumElementDispatch(
       // Reabstract to the substituted type, if needed.
 
       CanType substEltTy =
-        sourceType->getTypeOfMember(SGF.SGM.M.getSwiftModule(), formalElt,
-                                    formalElt->getArgumentInterfaceType())
+        sourceType->getTypeOfMember(SGF.SGM.M.getSwiftModule(), elt,
+                                    elt->getArgumentInterfaceType())
                   ->getCanonicalType();
 
       AbstractionPattern origEltTy =
@@ -2155,10 +2157,10 @@ void PatternMatchEmission::emitEnumElementDispatch(
 
     handleCase(eltCMV, specializedRows, *innerFailure);
     assert(!SGF.B.hasValidInsertionPoint() && "did not end block");
-  }
+  });
 
   // Emit the default block if we needed one.
-  if (defaultBB) {
+  if (SILBasicBlock *defaultBB = blocks.getDefaultBlock()) {
     SGF.B.setInsertionPoint(defaultBB);
     outerFailure(rows.back().Pattern);
   }