[region-isolation] Fix the dataflow and add support for project_block_storage

docs/SIL.rst

@@ -3844,6 +3844,8 @@ Notionally, ``alloc_pack_metadata`` is a stack allocation instruction.  See the
 section above on stack discipline.  The corresponding stack deallocation
 instruction is ``dealloc_pack_metadata``.
 
+Only valid in Lowered SIL.
+
 alloc_ref
 `````````
 ::
@@ -4124,6 +4126,8 @@ instruction after its operand) must be cleaned up here.
 on Stack Discipline above.  The operand must be an ``alloc_pack_metadata``
 instruction.
 
+Only valid in Lowered SIL.
+
 dealloc_box
 ```````````
 ::

include/swift/SIL/MemAccessUtils.h

@@ -1591,6 +1591,9 @@ inline bool isAccessStorageTypeCast(SingleValueInstruction *svi) {
   switch (svi->getKind()) {
   default:
     return false;
+  // This extracts out the block storage from an alloc_stack. We do not want
+  // to treat it as any more than a cast of the underlying value.
+  case SILInstructionKind::ProjectBlockStorageInst:
   // Simply pass-thru the incoming address.  But change its type!
   case SILInstructionKind::MoveOnlyWrapperToCopyableAddrInst:
   case SILInstructionKind::CopyableToMoveOnlyWrapperAddrInst:

include/swift/SIL/Projection.h

@@ -94,8 +94,9 @@ enum class ProjectionKind : unsigned {
   RefCast = 1,
   BitwiseCast = 2,
   TailElems = 3,
+  BlockStorageCast = 4,
   FirstPointerKind = Upcast,
-  LastPointerKind = TailElems,
+  LastPointerKind = BlockStorageCast,
 
   // Index Projection Kinds
   FirstIndexKind = 7,
@@ -120,6 +121,7 @@ static inline bool isCastProjectionKind(ProjectionKind Kind) {
   case ProjectionKind::Upcast:
   case ProjectionKind::RefCast:
   case ProjectionKind::BitwiseCast:
+  case ProjectionKind::BlockStorageCast:
     return true;
   case ProjectionKind::Struct:
   case ProjectionKind::Tuple:
@@ -318,15 +320,19 @@ class Projection {
     auto *Ty = getPointer();
     assert(Ty->isCanonical());
     switch (getKind()) {
-      case ProjectionKind::Upcast:
-      case ProjectionKind::RefCast:
-      case ProjectionKind::BitwiseCast:
-        return SILType::getPrimitiveType(Ty->getCanonicalType(),
-                                         BaseType.getCategory());
-      case ProjectionKind::TailElems:
-        return SILType::getPrimitiveAddressType(Ty->getCanonicalType());
-      default:
-        llvm_unreachable("unknown cast projection type");
+    case ProjectionKind::Upcast:
+    case ProjectionKind::RefCast:
+    case ProjectionKind::BitwiseCast:
+      return SILType::getPrimitiveType(Ty->getCanonicalType(),
+                                       BaseType.getCategory());
+    case ProjectionKind::TailElems:
+      return SILType::getPrimitiveAddressType(Ty->getCanonicalType());
+    case ProjectionKind::BlockStorageCast: {
+      auto blockStorageTy = Ty->getCanonicalType()->castTo<SILBlockStorageType>();
+      return blockStorageTy->getCaptureAddressType();
+    }
+    default:
+      llvm_unreachable("unknown cast projection type");
     }
   }
 
@@ -421,6 +427,7 @@ class Projection {
     case ProjectionKind::BitwiseCast:
       return true;
     case ProjectionKind::Upcast:
+    case ProjectionKind::BlockStorageCast:
     case ProjectionKind::Struct:
     case ProjectionKind::Tuple:
     case ProjectionKind::Index:
@@ -445,6 +452,7 @@ class Projection {
     case ProjectionKind::RefCast:
     case ProjectionKind::Tuple:
     case ProjectionKind::Upcast:
+    case ProjectionKind::BlockStorageCast:
     case ProjectionKind::Box:
     case ProjectionKind::TailElems:
       return false;

include/swift/SILOptimizer/Utils/PartitionUtils.h

@@ -577,7 +577,9 @@ class Partition {
             fstReduced.elementToRegionMap.find(Element(sndRegionNumber));
         if (iter != fstReduced.elementToRegionMap.end()) {
           fstReduced.elementToRegionMap.insert({sndEltNumber, iter->second});
-          if (fstReduced.fresh_label < Region(sndEltNumber))
+          // We want fresh_label to always be one element larger than our
+          // maximum element.
+          if (fstReduced.fresh_label <= Region(sndEltNumber))
             fstReduced.fresh_label = Region(sndEltNumber + 1);
           continue;
         }
@@ -597,15 +599,10 @@ class Partition {
           fstIter.first->getSecond() =
               fstIter.first->second->merge(sndIter->second);
       }
-      if (fstReduced.fresh_label < sndRegionNumber)
+      if (fstReduced.fresh_label <= sndRegionNumber)
         fstReduced.fresh_label = Region(sndEltNumber + 1);
     }
 
-    LLVM_DEBUG(llvm::dbgs() << "JOIN PEFORMED: \nFST: ";
-               fst.print(llvm::dbgs()); llvm::dbgs() << "SND: ";
-               snd.print(llvm::dbgs()); llvm::dbgs() << "RESULT: ";
-               fstReduced.print(llvm::dbgs()););
-
     assert(fstReduced.is_canonical_correct());
 
     // fst_reduced is now the join

lib/SIL/Utils/Projection.cpp

@@ -154,6 +154,13 @@ Projection::Projection(SingleValueInstruction *I) : Value() {
     }
     break;
   }
+  case SILInstructionKind::ProjectBlockStorageInst: {
+    auto *Ty = I->getType().getASTType().getPointer();
+    assert(Ty->isCanonical());
+    Value = ValueTy(ProjectionKind::BlockStorageCast, Ty);
+    assert(getKind() == ProjectionKind::BlockStorageCast);
+    break;
+  }
   case SILInstructionKind::UpcastInst: {
     auto *Ty = I->getType().getASTType().getPointer();
     assert(Ty->isCanonical());
@@ -199,6 +206,7 @@ SILType Projection::getType(SILType BaseType, SILModule &M,
   case ProjectionKind::Tuple:
     return BaseType.getTupleElementType(getIndex());
   case ProjectionKind::Upcast:
+  case ProjectionKind::BlockStorageCast:
   case ProjectionKind::RefCast:
   case ProjectionKind::BitwiseCast:
   case ProjectionKind::TailElems:
@@ -238,6 +246,8 @@ Projection::createObjectProjection(SILBuilder &B, SILLocation Loc,
     return nullptr;
   case ProjectionKind::Box:
     return nullptr;
+  case ProjectionKind::BlockStorageCast:
+    return nullptr;
   case ProjectionKind::Upcast:
     return B.createUpcast(Loc, Base, getCastType(BaseTy));
   case ProjectionKind::RefCast:
@@ -285,6 +295,8 @@ Projection::createAddressProjection(SILBuilder &B, SILLocation Loc,
     return B.createRefTailAddr(Loc, Base, getCastType(BaseTy));
   case ProjectionKind::Box:
     return B.createProjectBox(Loc, Base, getIndex());
+  case ProjectionKind::BlockStorageCast:
+    return B.createProjectBlockStorage(Loc, Base);
   case ProjectionKind::Upcast:
     return B.createUpcast(Loc, Base, getCastType(BaseTy));
   case ProjectionKind::RefCast:
@@ -600,6 +612,10 @@ void Projection::print(raw_ostream &os, SILType baseType) const {
     os << " Box over";
     break;
   }
+  case ProjectionKind::BlockStorageCast: {
+    os << "BlockStorageCast";
+    break;
+  }
   case ProjectionKind::Upcast: {
     os << "UpCast";
     break;
@@ -872,6 +888,7 @@ SILValue Projection::getOperandForAggregate(SILInstruction *I) const {
     case ProjectionKind::Class:
     case ProjectionKind::TailElems:
     case ProjectionKind::Box:
+    case ProjectionKind::BlockStorageCast:
     case ProjectionKind::Upcast:
     case ProjectionKind::RefCast:
     case ProjectionKind::BitwiseCast:
@@ -931,6 +948,7 @@ static bool isSupportedProjection(const Projection &p) {
   case ProjectionKind::Enum:
   case ProjectionKind::Box:
   case ProjectionKind::Upcast:
+  case ProjectionKind::BlockStorageCast:
   case ProjectionKind::RefCast:
   case ProjectionKind::BitwiseCast:
   case ProjectionKind::TailElems:

lib/SIL/Verifier/SILVerifier.cpp

@@ -6273,13 +6273,17 @@ class SILVerifier : public SILVerifierBase<SILVerifier> {
     require(apmi->getIntroducer()->mayRequirePackMetadata(),
             "Introduces instruction of kind which cannot emit on-stack pack "
             "metadata");
+    require(F.getModule().getStage() == SILStage::Lowered,
+            "Only supported in lowered SIL");
   }
 
   void checkDeallocPackMetadataInst(DeallocPackMetadataInst *dpmi) {
     auto *apmi = dpmi->getOperand()->getDefiningInstruction();
     require(apmi, "Must have instruction operand.");
     require(isa<AllocPackMetadataInst>(apmi),
             "Must have alloc_pack_metadata operand");
+    require(F.getModule().getStage() == SILStage::Lowered,
+            "Only supported in lowered SIL");
   }
 
   void checkMoveOnlyWrapperToCopyableAddrInst(