[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;
 
@@ -308,7 +318,7 @@ 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);
@@ -386,6 +396,7 @@ struct TypeTreeLeafTypeRange {
   /// common with filterBitVector.
   void constructFilteredProjections(
       SILValue value, SILInstruction *insertPt, SmallBitVector &filterBitVector,
+      DominanceInfo *domTree,
       llvm::function_ref<bool(SILValue, TypeTreeLeafTypeRange)> callback);
 
   void print(llvm::raw_ostream &os) const {

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,6 +377,7 @@ SubElementOffset::computeForValue(SILValue projectionDerivedFromRoot,
 
 void TypeTreeLeafTypeRange::constructFilteredProjections(
     SILValue value, SILInstruction *insertPt, SmallBitVector &filterBitVector,
+    DominanceInfo *domTree,
     llvm::function_ref<bool(SILValue, TypeTreeLeafTypeRange)> callback) {
   auto *fn = insertPt->getFunction();
   SILType type = value->getType();
@@ -419,27 +421,56 @@ 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));
+          foundProjection = true;
+        }
+        // assert(foundProjection);
+      }
+      return;
+    }
 
     // Then just pass back our enum base value as the pointer.
     callback(value, TypeTreeLeafTypeRange(start, next));
@@ -491,17 +522,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,7 +569,7 @@ void TypeTreeLeafTypeRange::get(
 }
 
 void TypeTreeLeafTypeRange::constructProjectionsForNeededElements(
-    SILValue rootValue, SILInstruction *insertPt,
+    SILValue rootValue, SILInstruction *insertPt, DominanceInfo *domTree,
     SmallBitVector &neededElements,
     SmallVectorImpl<std::pair<SILValue, TypeTreeLeafTypeRange>>
         &resultingProjections) {
@@ -569,7 +617,7 @@ void TypeTreeLeafTypeRange::constructProjectionsForNeededElements(
     // recursively process those ranges looking for subranges that have
     // completely set bits.
     range.constructFilteredProjections(
-        value, insertPt, neededElements,
+        value, insertPt, neededElements, domTree,
         [&](SILValue subType, TypeTreeLeafTypeRange range) -> bool {
           worklist.push_back({subType, range});
           return true;

lib/SILOptimizer/Mandatory/MoveOnlyAddressCheckerUtils.cpp

@@ -505,6 +505,8 @@ namespace {
 struct UseState {
   MarkUnresolvedNonCopyableValueInst *address;
 
+  DominanceInfo *const domTree;
+
   // FIXME: [partial_consume_of_deiniting_aggregate_with_drop_deinit] Keep track
   //        of the projections out of which a use emerges and use that to tell
   //        whether a particular partial consume is allowed.
@@ -605,6 +607,8 @@ struct UseState {
   /// scopes.
   DebugValueInst *debugValue = nullptr;
 
+  UseState(DominanceInfo *domTree) : domTree(domTree) {}
+
   SILFunction *getFunction() const { return address->getFunction(); }
 
   /// The number of fields in the exploded type.
@@ -1426,8 +1430,6 @@ struct MoveOnlyAddressCheckerPImpl {
 
   SILFunction *fn;
 
-  DominanceInfo *domTree;
-
   /// A set of mark_unresolved_non_copyable_value that we are actually going to
   /// process.
   llvm::SmallSetVector<MarkUnresolvedNonCopyableValueInst *, 32>
@@ -1459,9 +1461,9 @@ struct MoveOnlyAddressCheckerPImpl {
       SILFunction *fn, DiagnosticEmitter &diagnosticEmitter,
       DominanceInfo *domTree, PostOrderAnalysis *poa,
       borrowtodestructure::IntervalMapAllocator &allocator)
-      : fn(fn), domTree(domTree), deleter(),
-        canonicalizer(fn, domTree, deleter),
-        diagnosticEmitter(diagnosticEmitter), poa(poa), allocator(allocator) {
+      : fn(fn), deleter(), canonicalizer(fn, domTree, deleter),
+        addressUseState(domTree), diagnosticEmitter(diagnosticEmitter),
+        poa(poa), allocator(allocator) {
     deleter.setCallbacks(std::move(
         InstModCallbacks().onDelete([&](SILInstruction *instToDelete) {
           if (auto *mvi =
@@ -1597,6 +1599,11 @@ shouldVisitAsEndPointUse(Operand *op) {
   if (isa<DropDeinitInst>(op->getUser())) {
     return TransitiveAddressWalkerTransitiveUseVisitation::BothUserAndUses;
   }
+  // An unchecked_take_enum_data_addr consumes all bits except the remaining
+  // element's.
+  if (isa<UncheckedTakeEnumDataAddrInst>(op->getUser())) {
+    return TransitiveAddressWalkerTransitiveUseVisitation::BothUserAndUses;
+  }
   return TransitiveAddressWalkerTransitiveUseVisitation::OnlyUses;
 }
 
@@ -1762,6 +1769,7 @@ static std::optional<PartialMutationError>
 shouldEmitPartialMutationError(UseState &useState, PartialMutation::Kind kind,
                                SILInstruction *user, SILType useType,
                                TypeTreeLeafTypeRange usedBits) {
+  // FIXME: FIXME_NOW: This needs to be in terms of a use operand, right???
   SILFunction *fn = useState.getFunction();
 
   // We walk down from our ancestor to our projection, emitting an error if
@@ -2662,6 +2670,20 @@ bool GatherUsesVisitor::visitUse(Operand *op) {
     }
   }
 
+  if (auto *seai = dyn_cast<SwitchEnumAddrInst>(user)) {
+    SmallVector<TypeTreeLeafTypeRange, 2> leafRanges;
+    TypeTreeLeafTypeRange::get(op, getRootAddress(), leafRanges);
+    if (!leafRanges.size()) {
+      LLVM_DEBUG(llvm::dbgs() << "Failed to form leaf type range!\n");
+      return false;
+    }
+    for (auto leafRange : leafRanges) {
+      useState.recordLivenessUse(user, leafRange);
+    }
+
+    return true;
+  }
+
   // If we don't fit into any of those categories, just track as a liveness
   // use. We assume all such uses must only be reads to the memory. So we assert
   // to be careful.
@@ -3019,7 +3041,8 @@ static void insertDestroyBeforeInstruction(UseState &addressUseState,
       RegularLocation::getAutoGeneratedLocation(nextInstruction->getLoc());
   SmallVector<std::pair<SILValue, TypeTreeLeafTypeRange>> valuesToDestroy;
   TypeTreeLeafTypeRange::constructProjectionsForNeededElements(
-      baseAddress, nextInstruction, bv, valuesToDestroy);
+      baseAddress, nextInstruction, addressUseState.domTree, bv,
+      valuesToDestroy);
   while (!valuesToDestroy.empty()) {
     auto pair = valuesToDestroy.pop_back_val();
     if (pair.first->getType().isTrivial(*nextInstruction->getFunction()))

lib/SILOptimizer/Mandatory/MoveOnlyTypeUtils.cpp

@@ -105,7 +105,25 @@ TypeOffsetSizePair::walkOneLevelTowardsChild(
       // The only possible child of Optional is the wrapped type.
       return {{ancestorOffsetSize, ancestorType.getOptionalObjectType()}};
     }
-    llvm_unreachable("Cannot find child type of enum!\n");
+    unsigned elementOffset = ancestorOffsetSize.startOffset;
+    for (auto *element : enumDecl->getAllElements()) {
+      if (!element->hasAssociatedValues()) {
+        continue;
+      }
+      SILType elementTy = ancestorType.getEnumElementType(element, fn);
+      unsigned elementSize = TypeSubElementCount(elementTy, fn);
+      // iterOffset + size(tupleChild) is the offset of the next tuple
+      // element. If our target offset is less than that, then we know that
+      // the target type must be a child of this tuple element type.
+      if (elementOffset + elementSize > startOffset) {
+        return {{{elementOffset, elementSize}, elementTy}};
+      }
+
+      // Otherwise, add the new size of this field to iterOffset so we visit
+      // our sibling type next.
+      elementOffset += elementSize;
+    }
+    llvm_unreachable("Not a child of this enum?!");
   }
 
   llvm_unreachable("Hit a leaf type?! Should have handled it earlier");

lib/SILOptimizer/Mandatory/MoveOnlyUtils.cpp

@@ -318,6 +318,9 @@ bool noncopyable::memInstMustConsume(Operand *memOper) {
     auto convention = applySite.getArgumentConvention(*memOper);
     return !convention.isInoutConvention();
   }
+  case SILInstructionKind::UncheckedTakeEnumDataAddrInst: {
+    return true;
+  }
   }
 }
 

test/SILOptimizer/moveonly_generics_complex.swift

@@ -1,8 +1,9 @@
-// RUN: %target-swift-frontend                               \
-// RUN:     -emit-sil -verify                                \
-// RUN:     %s                                               \
-// RUN:     -enable-experimental-feature BuiltinModule       \
-// RUN:     -enable-experimental-feature NoncopyableGenerics \
+// RUN: %target-swift-frontend                                      \
+// RUN:     -emit-sil -verify                                       \
+// RUN:     %s                                                      \
+// RUN:     -enable-experimental-feature BuiltinModule              \
+// RUN:     -enable-experimental-feature NoncopyableGenerics        \
+// RUN:     -enable-experimental-feature MoveOnlyPartialConsumption \
 // RUN:     -sil-verify-all
 
 // REQUIRES: asserts
@@ -62,8 +63,7 @@ func _withUnprotectedUnsafeTemporaryAllocation<T: ~Copyable, R: ~Copyable>(
     Builtin.stackDealloc(stackAddress)
   }
   switch consume result {
-  // FIXME: This shouldn't be diagnosed.  But it's better than miscompiling.
-  case .success(let success): return success // expected-error{{cannot partially consume 'unknown'}}
+  case .success(let success): return success
   case .failure(let error): return try _rethrowsViaClosure { throw error }
   }
 }