[MoveOnlyAddressChecker] Fix enum representation.

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 {
@@ -291,83 +306,27 @@ struct TypeTreeLeafTypeRange {
   /// The leaf type sub-range of the type tree of \p rootAddress, consisting of
   /// \p projectedAddress and all of \p projectedAddress's descendent fields in
   /// the type tree.
-  ///
-  /// \returns None if we are unable to understand the path in between \p
-  /// projectedAddress and \p rootAddress.
-  static std::optional<TypeTreeLeafTypeRange> get(SILValue projectedValue,
-                                                  SILValue rootValue) {
+  static void get(SILValue projectedValue, SILValue rootValue,
+                  SmallVectorImpl<TypeTreeLeafTypeRange> &ranges) {
     auto startEltOffset = SubElementOffset::compute(projectedValue, rootValue);
     if (!startEltOffset)
-      return std::nullopt;
-    return {{*startEltOffset,
-             *startEltOffset + TypeSubElementCount(projectedValue)}};
+      return;
+    ranges.push_back({*startEltOffset,
+                      *startEltOffset + TypeSubElementCount(projectedValue)});
   }
 
   /// Which bits of \p rootValue are involved in \p op.
   ///
   /// This is a subset of (usually equal to) the bits of op->getType() in \p
   /// rootValue.
-  static std::optional<TypeTreeLeafTypeRange> get(Operand *op,
-                                                  SILValue rootValue) {
-    auto projectedValue = op->get();
-    auto startEltOffset = SubElementOffset::compute(projectedValue, rootValue);
-    if (!startEltOffset)
-      return std::nullopt;
-
-    // A drop_deinit only consumes the deinit bit of its operand.
-    if (isa<DropDeinitInst>(op->getUser())) {
-      auto upperBound = *startEltOffset + TypeSubElementCount(projectedValue);
-      return {{upperBound - 1, upperBound}};
-    }
-
-    // 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());
-
-        payloadUpperBound = *startEltOffset
-          + TypeSubElementCount(payloadTy, op->getFunction());
-      }
-      // TODO: account for deinit component if enum has deinit.
-      assert(!projectedValue->getType().isValueTypeWithDeinit());
-      return {{payloadUpperBound, upperBound}};
-    }
-
-    // Uses that borrow a value do not involve the deinit bit.
-    //
-    // FIXME: This shouldn't be limited to applies.
-    unsigned deinitBitOffset = 0;
-    if (op->get()->getType().isValueTypeWithDeinit() &&
-        op->getOperandOwnership() == OperandOwnership::Borrow &&
-        ApplySite::isa(op->getUser())) {
-      deinitBitOffset = 1;
-    }
-
-    return {{*startEltOffset, *startEltOffset +
-                                  TypeSubElementCount(projectedValue) -
-                                  deinitBitOffset}};
-  }
-
-  /// Given a type \p rootType and a set of needed elements specified by the bit
-  /// vector \p neededElements, place into \p foundContiguousTypeRanges a set of
-  /// TypeTreeLeafTypeRanges that are associated with the bit vectors
-  /// elements. As a constraint, we ensure that if \p neededElements has bits
-  /// set that are part of subsequent fields of a type that is only partially
-  /// needed, the two fields are represented as separate ranges. This ensures
-  /// that it is easy to use this API to correspond to independent operations
-  /// for the fields.
-  static void convertNeededElementsToContiguousTypeRanges(
-      SILFunction *fn, SILType rootType, SmallBitVector &neededElements,
-      SmallVectorImpl<TypeTreeLeafTypeRange> &foundContiguousTypeRanges);
+  static void get(Operand *op, SILValue rootValue,
+                  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,
@@ -442,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
@@ -533,10 +494,10 @@ class FieldSensitivePrunedLiveBlocks {
 
     /// Returns the liveness in \p resultingFoundLiveness. We only return the
     /// bits for endBitNo - startBitNo.
-    void getLiveness(unsigned startBitNo, unsigned endBitNo,
+    void getLiveness(SmallBitVector const &bitsOfInterest,
                      SmallVectorImpl<IsLive> &resultingFoundLiveness) const {
-      for (unsigned i = startBitNo, e = endBitNo; i != e; ++i) {
-        resultingFoundLiveness.push_back(getLiveness(i));
+      for (auto bit : bitsOfInterest.set_bits()) {
+        resultingFoundLiveness.push_back(getLiveness(bit));
       }
     }
 
@@ -654,17 +615,26 @@ class FieldSensitivePrunedLiveBlocks {
   void getBlockLiveness(SILBasicBlock *bb, unsigned startBitNo,
                         unsigned endBitNo,
                         SmallVectorImpl<IsLive> &foundLivenessInfo) const {
+    SmallBitVector bits(*numBitsToTrack);
+    for (auto index = startBitNo; index < endBitNo; ++index) {
+      bits.set(index);
+    }
+    getBlockLiveness(bb, bits, foundLivenessInfo);
+  }
+
+  void getBlockLiveness(SILBasicBlock *bb, SmallBitVector const &bits,
+                        SmallVectorImpl<IsLive> &foundLivenessInfo) const {
     assert(isInitialized());
     auto liveBlockIter = liveBlocks.find(bb);
     if (liveBlockIter == liveBlocks.end()) {
-      for (unsigned i : range(endBitNo - startBitNo)) {
-        (void)i;
+      for (auto bit : bits.set_bits()) {
+        (void)bit;
         foundLivenessInfo.push_back(Dead);
       }
       return;
     }
 
-    liveBlockIter->second.getLiveness(startBitNo, endBitNo, foundLivenessInfo);
+    liveBlockIter->second.getLiveness(bits, foundLivenessInfo);
   }
 
   llvm::StringRef getStringRef(IsLive isLive) const;
@@ -972,6 +942,12 @@ class FieldSensitivePrunedLiveness {
                                 resultingFoundLiveness);
   }
 
+  void getBlockLiveness(SILBasicBlock *bb, SmallBitVector const &bits,
+                        SmallVectorImpl<FieldSensitivePrunedLiveBlocks::IsLive>
+                            &foundLivenessInfo) const {
+    liveBlocks.getBlockLiveness(bb, bits, foundLivenessInfo);
+  }
+
   /// Return the liveness for this specific sub-element of our root value.
   FieldSensitivePrunedLiveBlocks::IsLive
   getBlockLiveness(SILBasicBlock *bb, unsigned subElementNumber) const {
@@ -1016,16 +992,16 @@ class FieldSensitivePrunedLiveness {
     return record->isInterestingUser(element);
   }
 
-  /// Whether \p user uses the fields in \p range as indicated by \p kind.
+  /// Whether \p user uses the fields in \p bits as indicated by \p kind.
   bool isInterestingUserOfKind(SILInstruction *user, IsInterestingUser kind,
-                               TypeTreeLeafTypeRange range) const {
+                               SmallBitVector const &bits) const {
     auto *record = getInterestingUser(user);
     if (!record) {
       return kind == IsInterestingUser::NonUser;
     }
 
-    for (auto element : range.getRange()) {
-      if (record->isInterestingUser(element) != kind)
+    for (auto bit : bits.set_bits()) {
+      if (record->isInterestingUser(bit) != kind)
         return false;
     }
     return true;
@@ -1172,10 +1148,10 @@ class FieldSensitivePrunedLiveRange : public FieldSensitivePrunedLiveness {
       : FieldSensitivePrunedLiveness(fn, discoveredBlocks) {}
 
   /// Check if \p inst occurs in between the definition of a def and the
-  /// liveness boundary for bits in \p span.
+  /// liveness boundary for \p bits.
   ///
-  /// NOTE: It is assumed that \p inst is correctly described by span.
-  bool isWithinBoundary(SILInstruction *inst, TypeTreeLeafTypeRange span) const;
+  /// NOTE: It is assumed that \p inst is correctly described by \p bits.
+  bool isWithinBoundary(SILInstruction *inst, SmallBitVector const &bits) const;
 
   /// Customize updateForUse for FieldSensitivePrunedLiveness such that we check
   /// that we consider defs as stopping liveness from being propagated up.