[MoveOnlyAddressChecker] Fix enum repr.

Change FieldSensitive's enum representation to allow distinguishing
among the elements with associated value.  Consider
`unchecked_take_enum_data_addr` to consume all other fields than that
taken.

rdar://125113258

Author: nate-chandler

include/swift/SIL/FieldSensitivePrunedLiveness.h

@@ -228,9 +228,19 @@ struct SubElementOffset {
   computeForValue(SILValue projectionFromRoot, SILValue rootValue);
 };
 
-/// Given a type T, this is the number of leaf field types in T's type tree. A
-/// leaf field type is a descendent field of T that does not have any
-/// descendent's itself.
+/// Counts the leaf fields aggregated together into a particular type.
+///
+/// Defined in such a way as to enable walking up the tree of aggregations
+/// node-by-node, visiting each type along the way.
+///
+/// The definition is given recursively as follows:
+/// a an atom  => count(a) := 1
+/// t a tuple  => count(t) := sum(t.elements, { elt in count(type(elt)) })
+/// 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
 struct TypeSubElementCount {
   unsigned number;
 
@@ -267,6 +277,11 @@ struct TypeSubElementCount {
 
 class FieldSensitivePrunedLiveness;
 
+enum NeedsDestroy_t {
+  DoesNotNeedDestroy = false,
+  NeedsDestroy = true,
+};
+
 /// A span of leaf elements in the sub-element break down of the linearization
 /// of the type tree of a type T.
 struct TypeTreeLeafTypeRange {
@@ -308,10 +323,10 @@ struct TypeTreeLeafTypeRange {
                   SmallVectorImpl<TypeTreeLeafTypeRange> &ranges);
 
   static void constructProjectionsForNeededElements(
-      SILValue rootValue, SILInstruction *insertPt,
+      SILValue rootValue, SILInstruction *insertPt, DominanceInfo *domTree,
       SmallBitVector &neededElements,
-      SmallVectorImpl<std::pair<SILValue, TypeTreeLeafTypeRange>>
-          &resultingProjections);
+      SmallVectorImpl<std::tuple<SILValue, TypeTreeLeafTypeRange,
+                                 NeedsDestroy_t>> &resultingProjections);
 
   static void visitContiguousRanges(
       SmallBitVector const &bits,
@@ -386,7 +401,9 @@ struct TypeTreeLeafTypeRange {
   /// common with filterBitVector.
   void constructFilteredProjections(
       SILValue value, SILInstruction *insertPt, SmallBitVector &filterBitVector,
-      llvm::function_ref<bool(SILValue, TypeTreeLeafTypeRange)> callback);
+      DominanceInfo *domTree,
+      llvm::function_ref<bool(SILValue, TypeTreeLeafTypeRange, NeedsDestroy_t)>
+          callback);
 
   void print(llvm::raw_ostream &os) const {
     os << "TypeTreeLeafTypeRange: (start: " << startEltOffset

lib/SIL/Utils/FieldSensitivePrunedLiveness.cpp

@@ -83,20 +83,20 @@ TypeSubElementCount::TypeSubElementCount(SILType type, SILModule &mod,
     return;
   }
 
-  // If we have an enum, we add one for tracking if the base enum is set and use
-  // the remaining bits for the max sized payload. This ensures that if we have
-  // a smaller sized payload, we still get all of the bits set, allowing for a
-  // homogeneous representation.
   if (auto *enumDecl = type.getEnumOrBoundGenericEnum()) {
     unsigned numElements = 0;
     for (auto *eltDecl : enumDecl->getAllElements()) {
       if (!eltDecl->hasAssociatedValues())
         continue;
       auto elt = type.getEnumElementType(eltDecl, mod, context);
-      numElements = std::max(numElements,
-                             unsigned(TypeSubElementCount(elt, mod, context)));
+      numElements += unsigned(TypeSubElementCount(elt, mod, context));
     }
     number = numElements + 1;
+    if (type.isValueTypeWithDeinit()) {
+      // 'self' has its own liveness represented as an additional field at the
+      // end of the structure.
+      ++number;
+    }
     return;
   }
 
@@ -190,19 +190,20 @@ SubElementOffset::computeForAddress(SILValue projectionDerivedFromRoot,
       continue;
     }
 
-    // In the case of enums, we note that our representation is:
-    //
-    //                   ---------|Enum| ---
-    //                  /                   \
-    //                 /                     \
-    //                v                       v
-    //  |Bits for Max Sized Payload|    |Discrim Bit|
-    //
-    // So our payload is always going to start at the current field number since
-    // we are the left most child of our parent enum. So we just need to look
-    // through to our parent enum.
     if (auto *enumData = dyn_cast<UncheckedTakeEnumDataAddrInst>(
             projectionDerivedFromRoot)) {
+      auto ty = enumData->getOperand()->getType();
+      auto *enumDecl = enumData->getEnumDecl();
+      for (auto *element : enumDecl->getAllElements()) {
+        if (!element->hasAssociatedValues())
+          continue;
+        if (element == enumData->getElement())
+          break;
+        auto context = TypeExpansionContext(*rootAddress->getFunction());
+        auto elementTy = ty.getEnumElementType(element, mod, context);
+        finalSubElementOffset +=
+            unsigned(TypeSubElementCount(elementTy, mod, context));
+      }
       projectionDerivedFromRoot = enumData->getOperand();
       continue;
     }
@@ -376,7 +377,9 @@ SubElementOffset::computeForValue(SILValue projectionDerivedFromRoot,
 
 void TypeTreeLeafTypeRange::constructFilteredProjections(
     SILValue value, SILInstruction *insertPt, SmallBitVector &filterBitVector,
-    llvm::function_ref<bool(SILValue, TypeTreeLeafTypeRange)> callback) {
+    DominanceInfo *domTree,
+    llvm::function_ref<bool(SILValue, TypeTreeLeafTypeRange, NeedsDestroy_t)>
+        callback) {
   auto *fn = insertPt->getFunction();
   SILType type = value->getType();
 
@@ -408,7 +411,7 @@ void TypeTreeLeafTypeRange::constructFilteredProjections(
 
       auto newValue =
           builder.createStructElementAddr(insertPt->getLoc(), value, varDecl);
-      callback(newValue, TypeTreeLeafTypeRange(start, next));
+      callback(newValue, TypeTreeLeafTypeRange(start, next), NeedsDestroy);
       start = next;
     }
     if (type.isValueTypeWithDeinit()) {
@@ -419,30 +422,63 @@ void TypeTreeLeafTypeRange::constructFilteredProjections(
     return;
   }
 
-  // We only allow for enums that can be completely destroyed. If there is code
-  // where an enum should be partially destroyed, we need to treat the
-  // unchecked_take_enum_data_addr as a separate value whose liveness we are
-  // tracking.
   if (auto *enumDecl = type.getEnumOrBoundGenericEnum()) {
+    struct ElementRecord {
+      EnumElementDecl *element;
+      unsigned start;
+      unsigned next;
+    };
+    SmallVector<ElementRecord, 2> projectedElements;
     unsigned start = startEltOffset;
-
-    unsigned maxSubEltCount = 0;
     for (auto *eltDecl : enumDecl->getAllElements()) {
       if (!eltDecl->hasAssociatedValues())
         continue;
+
       auto nextType = type.getEnumElementType(eltDecl, fn);
-      maxSubEltCount =
-          std::max(maxSubEltCount, unsigned(TypeSubElementCount(nextType, fn)));
+      unsigned next = start + TypeSubElementCount(nextType, fn);
+      if (noneSet(filterBitVector, start, next)) {
+        start = next;
+        continue;
+      }
+
+      projectedElements.push_back({eltDecl, start, next});
+      start = next;
     }
 
-    // Add a bit for the case bit.
-    unsigned next = maxSubEltCount + 1;
+    // Add a bit for the discriminator.
+    unsigned next = start + 1;
+
+    if (!allSet(filterBitVector, start, next)) {
+      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;
+          }
+          if (!domTree->dominates(utedai, insertPt)) {
+            continue;
+          }
 
-    // Make sure we are all set.
-    assert(allSet(filterBitVector, start, next));
+          callback(utedai, TypeTreeLeafTypeRange(record.start, record.next),
+                   DoesNotNeedDestroy);
+          foundProjection = true;
+        }
+        assert(foundProjection ||
+               llvm::count_if(enumDecl->getAllElements(), [](auto *elt) {
+                 return elt->hasAssociatedValues();
+               }) == 1);
+      }
+      return;
+    }
 
     // Then just pass back our enum base value as the pointer.
-    callback(value, TypeTreeLeafTypeRange(start, next));
+    callback(value, TypeTreeLeafTypeRange(start, next), NeedsDestroy);
 
     // Then set start to next and assert we covered the entire end elt offset.
     start = next;
@@ -463,7 +499,7 @@ void TypeTreeLeafTypeRange::constructFilteredProjections(
 
       auto newValue =
           builder.createTupleElementAddr(insertPt->getLoc(), value, index);
-      callback(newValue, TypeTreeLeafTypeRange(start, next));
+      callback(newValue, TypeTreeLeafTypeRange(start, next), NeedsDestroy);
       start = next;
     }
     assert(start == endEltOffset);
@@ -491,17 +527,34 @@ 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);
-    unsigned payloadUpperBound = 0;
-    if (inject->getElement()->hasAssociatedValues()) {
-      auto payloadTy = projectedValue->getType().getEnumElementType(
-          inject->getElement(), op->getFunction());
+    // TODO: account for deinit component if enum has deinit.
+    assert(!projectedValue->getType().isValueTypeWithDeinit());
+    ranges.push_back({upperBound - 1, upperBound});
+    return;
+  }
 
-      payloadUpperBound =
-          *startEltOffset + TypeSubElementCount(payloadTy, op->getFunction());
+  if (auto *utedai = dyn_cast<UncheckedTakeEnumDataAddrInst>(op->getUser())) {
+    auto *selected = utedai->getElement();
+    auto *enumDecl = utedai->getEnumDecl();
+    unsigned numAtoms = 0;
+    for (auto *element : enumDecl->getAllElements()) {
+      if (!element->hasAssociatedValues()) {
+        continue;
+      }
+      auto elementTy = projectedValue->getType().getEnumElementType(
+          element, op->getFunction());
+      auto elementAtoms =
+          unsigned(TypeSubElementCount(elementTy, op->getFunction()));
+      if (element != selected) {
+        ranges.push_back({*startEltOffset + numAtoms,
+                          *startEltOffset + numAtoms + elementAtoms});
+      }
+      numAtoms += elementAtoms;
     }
     // TODO: account for deinit component if enum has deinit.
     assert(!projectedValue->getType().isValueTypeWithDeinit());
-    ranges.push_back({payloadUpperBound, upperBound});
+    ranges.push_back(
+        {*startEltOffset + numAtoms, *startEltOffset + numAtoms + 1});
     return;
   }
 
@@ -521,17 +574,17 @@ void TypeTreeLeafTypeRange::get(
 }
 
 void TypeTreeLeafTypeRange::constructProjectionsForNeededElements(
-    SILValue rootValue, SILInstruction *insertPt,
+    SILValue rootValue, SILInstruction *insertPt, DominanceInfo *domTree,
     SmallBitVector &neededElements,
-    SmallVectorImpl<std::pair<SILValue, TypeTreeLeafTypeRange>>
+    SmallVectorImpl<std::tuple<SILValue, TypeTreeLeafTypeRange, NeedsDestroy_t>>
         &resultingProjections) {
   TypeTreeLeafTypeRange rootRange(rootValue);
   (void)rootRange;
   assert(rootRange.size() == neededElements.size());
 
-  StackList<std::pair<SILValue, TypeTreeLeafTypeRange>> worklist(
-      insertPt->getFunction());
-  worklist.push_back({rootValue, rootRange});
+  StackList<std::tuple<SILValue, TypeTreeLeafTypeRange, NeedsDestroy_t>>
+      worklist(insertPt->getFunction());
+  worklist.push_back({rootValue, rootRange, NeedsDestroy});
 
   // Temporary vector we use for our computation.
   SmallBitVector tmp(neededElements.size());
@@ -543,8 +596,10 @@ void TypeTreeLeafTypeRange::constructProjectionsForNeededElements(
 
   while (!worklist.empty()) {
     auto pair = worklist.pop_back_val();
-    auto value = pair.first;
-    auto range = pair.second;
+    SILValue value;
+    TypeTreeLeafTypeRange range;
+    NeedsDestroy_t needsDestroy;
+    std::tie(value, range, needsDestroy) = pair;
 
     tmp.reset();
     tmp.set(range.startEltOffset, range.endEltOffset);
@@ -561,17 +616,18 @@ void TypeTreeLeafTypeRange::constructProjectionsForNeededElements(
     // everything set in the range. In that case, we just add this range to the
     // result and continue.
     if (allInRange(tmp, range)) {
-      resultingProjections.emplace_back(value, range);
+      resultingProjections.emplace_back(value, range, needsDestroy);
       continue;
     }
 
     // Otherwise, we have a partial range. We need to split our range and then
     // recursively process those ranges looking for subranges that have
     // completely set bits.
     range.constructFilteredProjections(
-        value, insertPt, neededElements,
-        [&](SILValue subType, TypeTreeLeafTypeRange range) -> bool {
-          worklist.push_back({subType, range});
+        value, insertPt, neededElements, domTree,
+        [&](SILValue subType, TypeTreeLeafTypeRange range,
+            NeedsDestroy_t needsDestroy) -> bool {
+          worklist.push_back({subType, range, needsDestroy});
           return true;
         });
   }