[MoveOnlyAddressChecker] Fix representation for used fields.

include/swift/SIL/FieldSensitivePrunedLiveness.h

@@ -345,6 +345,16 @@ struct TypeTreeLeafTypeRange {
     return TypeTreeLeafTypeRange(start, end);
   }
 
+  /// Whether \p bits contains any of the in-range bits.
+  bool intersects(SmallBitVector const &bits) const {
+    for (auto element : getRange()) {
+      if (bits.test(element)) {
+        return true;
+      }
+    }
+    return false;
+  }
+
   /// Is the given leaf type specified by \p singleLeafElementNumber apart of
   /// our \p range of leaf type values in the our larger type.
   bool contains(SubElementOffset singleLeafElementNumber) const {
@@ -359,7 +369,7 @@ struct TypeTreeLeafTypeRange {
   }
 
   /// Sets each bit in \p bits corresponding to an element of this range.
-  void setBits(SmallBitVector &bits) {
+  void setBits(SmallBitVector &bits) const {
     for (auto element : getRange()) {
       bits.set(element);
     }
@@ -696,6 +706,14 @@ class FieldSensitivePrunedLiveness {
       }
     }
 
+    /// Record that the instruction uses the bits in \p bits.
+    void addUses(SmallBitVector const &bits, bool lifetimeEnding) {
+      liveBits |= bits;
+      if (lifetimeEnding) {
+        consumingBits |= bits;
+      }
+    }
+
     /// Populates the provided vector with contiguous ranges of bits which are
     /// users of the same sort.
     void getContiguousRanges(
@@ -838,6 +856,9 @@ class FieldSensitivePrunedLiveness {
   void updateForUse(SILInstruction *user, TypeTreeLeafTypeRange span,
                     bool lifetimeEnding);
 
+  void updateForUse(SILInstruction *user, SmallBitVector const &bits,
+                    bool lifetimeEnding);
+
   void getBlockLiveness(SILBasicBlock *bb, TypeTreeLeafTypeRange span,
                         SmallVectorImpl<FieldSensitivePrunedLiveBlocks::IsLive>
                             &resultingFoundLiveness) const {
@@ -862,6 +883,14 @@ class FieldSensitivePrunedLiveness {
                         SmallBitVector &liveOutBits,
                         SmallBitVector &deadBits) const;
 
+  InterestingUser &getOrCreateInterestingUser(SILInstruction *user) {
+    auto iter = users.find(user);
+    if (iter == users.end()) {
+      iter = users.insert({user, InterestingUser(getNumSubElements())}).first;
+    }
+    return *&iter->second;
+  }
+
   /// If \p user has had uses recored, return a pointer to the InterestingUser
   /// where they've been recorded.
   InterestingUser const *getInterestingUser(SILInstruction *user) const {
@@ -885,11 +914,12 @@ class FieldSensitivePrunedLiveness {
   bool isInterestingUserOfKind(SILInstruction *user, IsInterestingUser kind,
                                TypeTreeLeafTypeRange range) const {
     auto *record = getInterestingUser(user);
-    if (!record)
+    if (!record) {
       return kind == IsInterestingUser::NonUser;
+    }
 
     for (auto element : range.getRange()) {
-      if (isInterestingUser(user, element) != kind)
+      if (record->isInterestingUser(element) != kind)
         return false;
     }
     return true;
@@ -918,11 +948,12 @@ class FieldSensitivePrunedLiveness {
   /// argument must be copied.
   void addInterestingUser(SILInstruction *user, TypeTreeLeafTypeRange range,
                           bool lifetimeEnding) {
-    auto iter = users.find(user);
-    if (iter == users.end()) {
-      iter = users.insert({user, InterestingUser(getNumSubElements())}).first;
-    }
-    iter->second.addUses(range, lifetimeEnding);
+    getOrCreateInterestingUser(user).addUses(range, lifetimeEnding);
+  }
+
+  void addInterestingUser(SILInstruction *user, SmallBitVector const &bits,
+                          bool lifetimeEnding) {
+    getOrCreateInterestingUser(user).addUses(bits, lifetimeEnding);
   }
 };
 
@@ -1036,6 +1067,11 @@ class FieldSensitivePrunedLiveRange : public FieldSensitivePrunedLiveness {
   void updateForUse(SILInstruction *user, TypeTreeLeafTypeRange span,
                     bool lifetimeEnding);
 
+  /// Customize updateForUse for FieldSensitivePrunedLiveness such that we check
+  /// that we consider defs as stopping liveness from being propagated up.
+  void updateForUse(SILInstruction *user, SmallBitVector const &bits,
+                    bool lifetimeEnding);
+
   /// Compute the boundary from the blocks discovered during liveness analysis.
   ///
   /// Precondition: \p liveness.getDiscoveredBlocks() is a valid list of all
@@ -1107,6 +1143,14 @@ class FieldSensitiveSSAPrunedLiveRange
     return inst == defInst.first && defInst.second->contains(bit);
   }
 
+  bool isDef(SILInstruction *inst, SmallBitVector const &bits) const {
+    if (inst != defInst.first)
+      return false;
+    SmallBitVector defBits(bits.size());
+    defInst.second->setBits(defBits);
+    return (defBits & bits) == bits;
+  }
+
   bool isDef(SILInstruction *inst, TypeTreeLeafTypeRange span) const {
     return inst == defInst.first &&
            defInst.second->setIntersection(span).has_value();
@@ -1217,6 +1261,30 @@ class FieldSensitiveMultiDefPrunedLiveRange
         *iter, [&](TypeTreeLeafTypeRange span) { return span.contains(bit); });
   }
 
+  bool isDef(SILValue value, SmallBitVector const &bits) const {
+    assert(isInitialized());
+    auto iter = defs.find(cast<SILNode>(value));
+    if (!iter)
+      return false;
+    SmallBitVector allBits(bits.size());
+    for (auto range : *iter) {
+      range.setBits(allBits);
+    }
+    return (bits & allBits) == bits;
+  }
+
+  bool isDef(SILInstruction *inst, SmallBitVector const &bits) const {
+    assert(isInitialized());
+    auto iter = defs.find(cast<SILNode>(inst));
+    if (!iter)
+      return false;
+    SmallBitVector allBits(bits.size());
+    for (auto range : *iter) {
+      range.setBits(allBits);
+    }
+    return (bits & allBits) == bits;
+  }
+
   bool isDef(SILInstruction *inst, TypeTreeLeafTypeRange span) const {
     assert(isInitialized());
     auto iter = defs.find(cast<SILNode>(inst));

lib/SIL/Utils/FieldSensitivePrunedLiveness.cpp

@@ -609,6 +609,16 @@ void FieldSensitivePrunedLiveness::updateForUse(SILInstruction *user,
   addInterestingUser(user, range, lifetimeEnding);
 }
 
+void FieldSensitivePrunedLiveness::updateForUse(SILInstruction *user,
+                                                SmallBitVector const &bits,
+                                                bool lifetimeEnding) {
+  for (auto bit : bits.set_bits()) {
+    liveBlocks.updateForUse(user, bit);
+  }
+
+  addInterestingUser(user, bits, lifetimeEnding);
+}
+
 //===----------------------------------------------------------------------===//
 //                    MARK: FieldSensitivePrunedLiveRange
 //===----------------------------------------------------------------------===//
@@ -822,6 +832,42 @@ void FieldSensitivePrunedLiveRange<LivenessWithDefs>::updateForUse(
   FieldSensitivePrunedLiveness::updateForUse(user, range, lifetimeEnding);
 }
 
+template <typename LivenessWithDefs>
+void FieldSensitivePrunedLiveRange<LivenessWithDefs>::updateForUse(
+    SILInstruction *user, SmallBitVector const &bits, bool lifetimeEnding) {
+  PRUNED_LIVENESS_LOG(
+      llvm::dbgs()
+      << "Begin FieldSensitivePrunedLiveRange<LivenessWithDefs>::updateForUse "
+         "for: "
+      << *user);
+  PRUNED_LIVENESS_LOG(llvm::dbgs()
+                      << "Looking for def instruction earlier in the block!\n");
+
+  auto *parentBlock = user->getParent();
+  for (auto ii = std::next(user->getReverseIterator()),
+            ie = parentBlock->rend();
+       ii != ie; ++ii) {
+    // If we find the def, just mark this instruction as being an interesting
+    // instruction.
+    if (asImpl().isDef(&*ii, bits)) {
+      PRUNED_LIVENESS_LOG(llvm::dbgs() << "    Found def: " << *ii);
+      PRUNED_LIVENESS_LOG(
+          llvm::dbgs()
+          << "    Marking inst as interesting user and returning!\n");
+      addInterestingUser(user, bits, lifetimeEnding);
+      return;
+    }
+  }
+
+  // Otherwise, just delegate to our parent class's update for use. This will
+  // update liveness for our predecessor blocks and add this instruction as an
+  // interesting user.
+  PRUNED_LIVENESS_LOG(llvm::dbgs()
+                      << "No defs found! Delegating to "
+                         "FieldSensitivePrunedLiveness::updateForUse.\n");
+  FieldSensitivePrunedLiveness::updateForUse(user, bits, lifetimeEnding);
+}
+
 //===----------------------------------------------------------------------===//
 //                    MARK: Boundary Computation Utilities
 //===----------------------------------------------------------------------===//