[MoveOnlyAddressChecker] Handle enum deinits.

Previously, the enum representation was fixed to represent the different
cases payloads separately with the unchecked_take_enum_data_addr
instruction consuming all fields but that whose address is obtained.
In a few places, handling for enum deinits was left undone with an
assertion that the enum not have one.

Here, the deinit bit of the enum is shifted to the end.  And the
assertions are replaced with handling.  Finally, the logic for inserting
destroys after switch_enum_addr instructions is fixed.

Author: nate-chandler

include/swift/SIL/FieldSensitivePrunedLiveness.h

@@ -239,8 +239,11 @@ struct SubElementOffset {
 /// s a struct => count(s) := sum(s.fields, { f in count(type(f)) })
 ///                             + s.hasDeinit
 /// e an enum  => count(e) := sum(e.elements, { elt in count(type(elt)) })
-///                             + e.hasDeinit
 ///                             + 1 // discriminator
+///                             + e.hasDeinit
+///
+/// The deinit bit is at the end to make drop_deinit produce a value whose
+/// leaves are contiguous.
 struct TypeSubElementCount {
   unsigned number;
 

lib/SIL/Utils/FieldSensitivePrunedLiveness.cpp

@@ -375,6 +375,39 @@ SubElementOffset::computeForValue(SILValue projectionDerivedFromRoot,
 //                        MARK: TypeTreeLeafTypeRange
 //===----------------------------------------------------------------------===//
 
+/// Whether \p targetInst is dominated by one of the provided switch_enum_addr's
+/// destination blocks whose corresponding enum element has no associated values
+/// which need to be destroyed (i.e. either it has no associated values or they
+/// are trivial).
+static bool isDominatedByPayloadlessSwitchEnumAddrDests(
+    SILInstruction *targetInst, ArrayRef<SwitchEnumAddrInst *> seais,
+    DominanceInfo *domTree) {
+  if (seais.empty())
+    return false;
+  auto *target = targetInst->getParent();
+  for (auto *seai : seais) {
+    if (!domTree->dominates(seai, targetInst)) {
+      continue;
+    }
+    auto size = seai->getNumCases();
+    auto ty = seai->getOperand()->getType();
+    for (unsigned index = 0; index < size; ++index) {
+      auto pair = seai->getCase(index);
+      auto *eltDecl = pair.first;
+      if (eltDecl->hasAssociatedValues()) {
+        auto eltTy = ty.getEnumElementType(eltDecl, seai->getFunction());
+        if (!eltTy.isTrivial(*seai->getFunction())) {
+          continue;
+        }
+      }
+      auto *block = pair.second;
+      if (domTree->dominates(block, target))
+        return true;
+    }
+  }
+  return false;
+}
+
 void TypeTreeLeafTypeRange::constructFilteredProjections(
     SILValue value, SILInstruction *insertPt, SmallBitVector &filterBitVector,
     DominanceInfo *domTree,
@@ -429,60 +462,72 @@ void TypeTreeLeafTypeRange::constructFilteredProjections(
       unsigned next;
     };
     SmallVector<ElementRecord, 2> projectedElements;
-    unsigned start = startEltOffset;
+    unsigned runningStart = startEltOffset;
     for (auto *eltDecl : enumDecl->getAllElements()) {
       if (!eltDecl->hasAssociatedValues())
         continue;
 
-      auto nextType = type.getEnumElementType(eltDecl, fn);
-      unsigned next = start + TypeSubElementCount(nextType, fn);
-      if (noneSet(filterBitVector, start, next)) {
-        start = next;
+      auto eltTy = type.getEnumElementType(eltDecl, fn);
+      unsigned next = runningStart + TypeSubElementCount(eltTy, fn);
+      if (noneSet(filterBitVector, runningStart, next)) {
+        runningStart = next;
         continue;
       }
 
-      projectedElements.push_back({eltDecl, start, next});
-      start = next;
+      projectedElements.push_back({eltDecl, runningStart, next});
+      runningStart = next;
     }
+    assert((runningStart + 1 + (type.isValueTypeWithDeinit() ? 1 : 0)) ==
+           endEltOffset);
 
-    // Add a bit for the discriminator.
-    unsigned next = start + 1;
-
-    if (!allSet(filterBitVector, start, next)) {
+    if (!allSet(filterBitVector, startEltOffset, endEltOffset)) {
+      TinyPtrVector<SwitchEnumAddrInst *> seais;
       for (auto record : projectedElements) {
         // Find a preexisting unchecked_take_enum_data_addr that dominates
         // insertPt.
         bool foundProjection = false;
-        for (auto *user : value->getUsers()) {
-          auto *utedai = dyn_cast<UncheckedTakeEnumDataAddrInst>(user);
-          if (!utedai) {
-            continue;
-          }
-          if (utedai->getElement() == record.element) {
-            continue;
+        StackList<SILValue> worklist(value->getFunction());
+        worklist.push_back(value);
+        while (!worklist.empty()) {
+          auto v = worklist.pop_back_val();
+          for (auto *user : v->getUsers()) {
+            if (auto *ddi = dyn_cast<DropDeinitInst>(user)) {
+              worklist.push_back(ddi);
+              continue;
+            }
+            if (auto *seai = dyn_cast<SwitchEnumAddrInst>(user)) {
+              seais.push_back(seai);
+            }
+            auto *utedai = dyn_cast<UncheckedTakeEnumDataAddrInst>(user);
+            if (!utedai) {
+              continue;
+            }
+            if (utedai->getElement() != record.element) {
+              continue;
+            }
+            if (!domTree->dominates(utedai, insertPt)) {
+              continue;
+            }
+
+            callback(utedai, TypeTreeLeafTypeRange(record.start, record.next),
+                     NeedsDestroy);
+            foundProjection = true;
           }
-          if (!domTree->dominates(utedai, insertPt)) {
-            continue;
-          }
-
-          callback(utedai, TypeTreeLeafTypeRange(record.start, record.next),
-                   DoesNotNeedDestroy);
-          foundProjection = true;
         }
+        (void)foundProjection;
         assert(foundProjection ||
-               llvm::count_if(enumDecl->getAllElements(), [](auto *elt) {
-                 return elt->hasAssociatedValues();
-               }) == 1);
+               llvm::count_if(
+                   enumDecl->getAllElements(),
+                   [](auto *elt) { return elt->hasAssociatedValues(); }) == 1 ||
+               isDominatedByPayloadlessSwitchEnumAddrDests(insertPt, seais,
+                                                           domTree));
       }
       return;
     }
 
     // Then just pass back our enum base value as the pointer.
-    callback(value, TypeTreeLeafTypeRange(start, next), NeedsDestroy);
-
-    // Then set start to next and assert we covered the entire end elt offset.
-    start = next;
-    assert(start == endEltOffset);
+    callback(value, TypeTreeLeafTypeRange(startEltOffset, endEltOffset),
+             NeedsDestroy);
     return;
   }
 
@@ -526,9 +571,12 @@ void TypeTreeLeafTypeRange::get(
 
   // An `inject_enum_addr` only initializes the enum tag.
   if (auto inject = dyn_cast<InjectEnumAddrInst>(op->getUser())) {
-    auto upperBound = *startEltOffset + TypeSubElementCount(projectedValue);
-    // TODO: account for deinit component if enum has deinit.
-    assert(!projectedValue->getType().isValueTypeWithDeinit());
+    // Subtract the deinit bit, if any: the discriminator bit is before it:
+    //
+    // [ case1 bits ..., case2 bits, ..., discriminator bit, deinit bit ]
+    auto deinitBits = projectedValue->getType().isValueTypeWithDeinit() ? 1 : 0;
+    auto upperBound =
+        *startEltOffset + TypeSubElementCount(projectedValue) - deinitBits;
     ranges.push_back({upperBound - 1, upperBound});
     return;
   }
@@ -551,10 +599,11 @@ void TypeTreeLeafTypeRange::get(
       }
       numAtoms += elementAtoms;
     }
-    // TODO: account for deinit component if enum has deinit.
-    assert(!projectedValue->getType().isValueTypeWithDeinit());
+    // The discriminator bit is consumed.
     ranges.push_back(
         {*startEltOffset + numAtoms, *startEltOffset + numAtoms + 1});
+    // The deinit bit is _not_ consumed.  A drop_deinit is required to
+    // consumingly switch an enum with a deinit.
     return;
   }