[pruned-liveness] Implement FieldSensitiveAddressPrunedLiveness.

include/swift/SIL/PrunedLiveness.h

@@ -94,7 +94,9 @@
 #ifndef SWIFT_SILOPTIMIZER_UTILS_PRUNEDLIVENESS_H
 #define SWIFT_SILOPTIMIZER_UTILS_PRUNEDLIVENESS_H
 
+#include "swift/AST/TypeExpansionContext.h"
 #include "swift/SIL/SILBasicBlock.h"
+#include "swift/SIL/SILFunction.h"
 #include "llvm/ADT/MapVector.h"
 #include "llvm/ADT/PointerIntPair.h"
 #include "llvm/ADT/SmallVector.h"
@@ -236,6 +238,8 @@ class PrunedLiveBlocks {
     assert(!discoveredBlocks || discoveredBlocks->empty());
   }
 
+  unsigned getNumBitsToTrack() const { return numBitsToTrack; }
+
   bool empty() const { return liveBlocks.empty(); }
 
   void clear() {
@@ -290,7 +294,6 @@ class PrunedLiveBlocks {
       return;
     }
 
-    assert(liveBlockIter->second.size() == 1);
     liveBlockIter->second.getLiveness(startBitNo, endBitNo, foundLivenessInfo);
   }
 
@@ -553,6 +556,289 @@ struct PrunedLivenessBoundary {
                ArrayRef<SILBasicBlock *> postDomBlocks);
 };
 
+/// Given a type T and a descendent field F in T's type tree, then the
+/// sub-element number of F is the first leaf element of the type tree in its
+/// linearized representation.
+struct SubElementNumber {
+  unsigned number;
+
+  SubElementNumber(unsigned number) : number(number) {}
+
+  /// Given an arbitrary projection \p projectionFromRoot from the \p
+  /// rootAddress, compute the sub element number for that \p SILValue. The sub
+  /// element number of a type T is always the index of its first leaf node
+  /// descendent in the type tree.
+  ///
+  /// DISCUSSION: This works for non-leaf types in the type tree as well as
+  /// normal leaf elements. It is the index of the first leaf element that is a
+  /// sub element of the root SILType that this projection will effect. The rest
+  /// of the elements effected can be found by computing the number of leaf sub
+  /// elements of \p projectionFromRoot's type and adding this to the result of
+  /// this function.
+  ///
+  /// \returns None if we didn't know how to compute sub-element for this
+  /// projection.
+  static Optional<SubElementNumber> compute(SILValue projectionFromRoot,
+                                            SILValue root);
+
+  operator unsigned() const { return number; }
+};
+
+/// 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.
+struct TypeSubElementCount {
+  unsigned number;
+
+  TypeSubElementCount(unsigned number) : number(number) {}
+
+  /// Given a type \p type, compute the total number of leaf sub-elements of \p
+  /// type in the type tree.
+  ///
+  /// Some interesting properties of this computation:
+  ///
+  /// 1. When applied to the root type, this equals the total number of bits of
+  /// liveness that we track.
+  ///
+  /// 2. When applied to a field type F of the type tree for a type T,
+  /// computeNumLeafSubElements(F) when added to F's start sub element number
+  /// will go to the next sibling node in the type tree, walking up the tree and
+  /// attempting to find siblings if no further siblings exist.
+  TypeSubElementCount(SILType type, SILModule &mod,
+                      TypeExpansionContext context);
+
+  TypeSubElementCount(SILValue value)
+      : TypeSubElementCount(value->getType(), *value->getModule(),
+                            TypeExpansionContext(*value->getFunction())) {}
+
+  operator unsigned() const { return number; }
+};
+
+/// A span of leaf elements in the sub-element break down of the linearization
+/// of the type tree of a type T.
+struct TypeTreeLeafTypeRange {
+  SubElementNumber startEltOffset;
+  SubElementNumber endEltOffset;
+
+  /// The leaf type range for the entire type tree.
+  TypeTreeLeafTypeRange(SILValue rootAddress)
+      : startEltOffset(0), endEltOffset(TypeSubElementCount(rootAddress)) {}
+
+  /// 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.
+  TypeTreeLeafTypeRange(SILValue projectedAddress, SILValue rootAddress)
+      : startEltOffset(
+            *SubElementNumber::compute(projectedAddress, rootAddress)),
+        endEltOffset(startEltOffset + TypeSubElementCount(projectedAddress)) {}
+
+  /// 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(SubElementNumber singleLeafElementNumber) const {
+    return startEltOffset <= singleLeafElementNumber &&
+           singleLeafElementNumber < endEltOffset;
+  }
+
+  /// Returns true if either of this overlaps at all with the given range.
+  bool contains(TypeTreeLeafTypeRange range) const {
+    if (startEltOffset <= range.startEltOffset &&
+        range.startEltOffset < endEltOffset)
+      return true;
+
+    // If our start and end offsets, our extent is only 1 and we know that our
+    // value
+    unsigned rangeLastElt = range.endEltOffset - 1;
+    if (range.startEltOffset == rangeLastElt)
+      return false;
+
+    // Othrwise, see if endEltOffset - 1 is within the range.
+    return startEltOffset <= rangeLastElt && rangeLastElt < endEltOffset;
+  }
+};
+
+/// This is exactly like pruned liveness except that instead of tracking a
+/// single bit of liveness, it tracks multiple bits of liveness for leaf type
+/// tree nodes of an allocation one is calculating pruned liveness for.
+///
+/// DISCUSSION: One can view a type T as a tree with recursively each field F of
+/// the type T being a child of T in the tree. We say recursively since the tree
+/// unfolds for F and its children as well.
+class FieldSensitiveAddressPrunedLiveness {
+  PrunedLiveBlocks liveBlocks;
+
+  struct InterestingUser {
+    TypeTreeLeafTypeRange subEltSpan;
+    bool isConsuming;
+
+    InterestingUser(TypeTreeLeafTypeRange subEltSpan, bool isConsuming)
+        : subEltSpan(subEltSpan), isConsuming(isConsuming) {}
+
+    InterestingUser &operator&=(bool otherValue) {
+      isConsuming &= otherValue;
+      return *this;
+    }
+  };
+
+  /// Map all "interesting" user instructions in this def's live range to a pair
+  /// consisting of the SILValue that it uses and a flag indicating whether they
+  /// must end the lifetime.
+  ///
+  /// Lifetime-ending users are always on the boundary so are always
+  /// interesting.
+  ///
+  /// Non-lifetime-ending uses within a LiveWithin block are interesting because
+  /// they may be the last use in the block.
+  ///
+  /// Non-lifetime-ending within a LiveOut block are uninteresting.
+  llvm::SmallMapVector<SILInstruction *, InterestingUser, 8> users;
+
+  /// The root address of our type tree.
+  SILValue rootAddress;
+
+public:
+  FieldSensitiveAddressPrunedLiveness(
+      SILFunction *fn, SILValue rootValue,
+      SmallVectorImpl<SILBasicBlock *> *discoveredBlocks = nullptr)
+      : liveBlocks(TypeSubElementCount(rootValue), discoveredBlocks),
+        rootAddress(rootValue) {}
+
+  bool empty() const {
+    assert(!liveBlocks.empty() || users.empty());
+    return liveBlocks.empty();
+  }
+
+  void clear() {
+    liveBlocks.clear();
+    users.clear();
+  }
+
+  SILValue getRootAddress() const { return rootAddress; }
+
+  unsigned numLiveBlocks() const { return liveBlocks.numLiveBlocks(); }
+
+  /// If the constructor was provided with a vector to populate, then this
+  /// returns the list of all live blocks with no duplicates.
+  ArrayRef<SILBasicBlock *> getDiscoveredBlocks() const {
+    return liveBlocks.getDiscoveredBlocks();
+  }
+
+  using UserRange =
+      iterator_range<const std::pair<SILInstruction *, InterestingUser> *>;
+  UserRange getAllUsers() const {
+    return llvm::make_range(users.begin(), users.end());
+  }
+
+  using UserBlockRange = TransformRange<
+      UserRange, function_ref<SILBasicBlock *(
+                     const std::pair<SILInstruction *, InterestingUser> &)>>;
+  UserBlockRange getAllUserBlocks() const {
+    function_ref<SILBasicBlock *(
+        const std::pair<SILInstruction *, InterestingUser> &)>
+        op;
+    op = [](const std::pair<SILInstruction *, InterestingUser> &pair)
+        -> SILBasicBlock * { return pair.first->getParent(); };
+    return UserBlockRange(getAllUsers(), op);
+  }
+
+  void initializeDefBlock(SILBasicBlock *defBB, TypeTreeLeafTypeRange span) {
+    liveBlocks.initializeDefBlock(defBB, span.startEltOffset,
+                                  span.endEltOffset);
+  }
+
+  /// For flexibility, \p lifetimeEnding is provided by the
+  /// caller. PrunedLiveness makes no assumptions about the def-use
+  /// relationships that generate liveness. For example, use->isLifetimeEnding()
+  /// cannot distinguish the end of the borrow scope that defines this extended
+  /// live range vs. a nested borrow scope within the extended live range.
+  ///
+  /// Also for flexibility, \p affectedAddress must be a derived projection from
+  /// the base that \p user is affecting.
+  void updateForUse(SILInstruction *user, TypeTreeLeafTypeRange span,
+                    bool lifetimeEnding);
+
+  void getBlockLiveness(
+      SILBasicBlock *bb, TypeTreeLeafTypeRange span,
+      SmallVectorImpl<PrunedLiveBlocks::IsLive> &resultingFoundLiveness) const {
+    liveBlocks.getBlockLiveness(bb, span.startEltOffset, span.endEltOffset,
+                                resultingFoundLiveness);
+  }
+
+  /// Return the liveness for this specific sub-element of our root value.
+  PrunedLiveBlocks::IsLive getBlockLiveness(SILBasicBlock *bb,
+                                            unsigned subElementNumber) const {
+    SmallVector<PrunedLiveBlocks::IsLive, 1> isLive;
+    liveBlocks.getBlockLiveness(bb, subElementNumber, subElementNumber + 1,
+                                isLive);
+    return isLive[0];
+  }
+
+  void getBlockLiveness(
+      SILBasicBlock *bb,
+      SmallVectorImpl<PrunedLiveBlocks::IsLive> &foundLiveness) const {
+    liveBlocks.getBlockLiveness(bb, 0, liveBlocks.getNumBitsToTrack(),
+                                foundLiveness);
+  }
+
+  enum IsInterestingUser { NonUser, NonLifetimeEndingUse, LifetimeEndingUse };
+
+  /// Return a result indicating whether the given user was identified as an
+  /// interesting use of the current def and whether it ends the lifetime.
+  std::pair<IsInterestingUser, Optional<TypeTreeLeafTypeRange>>
+  isInterestingUser(SILInstruction *user) const {
+    auto useIter = users.find(user);
+    if (useIter == users.end())
+      return {NonUser, None};
+    auto isInteresting =
+        useIter->second.isConsuming ? LifetimeEndingUse : NonLifetimeEndingUse;
+    return {isInteresting, useIter->second.subEltSpan};
+  }
+
+  unsigned getNumSubElements() const { return liveBlocks.getNumBitsToTrack(); }
+
+  /// Return true if \p inst occurs before the liveness boundary. Used when the
+  /// client already knows that inst occurs after the start of liveness.
+  void isWithinBoundary(SILInstruction *inst, SmallBitVector &outVector) const;
+};
+
+/// Record the last use points and CFG edges that form the boundary of
+/// FieldSensitiveAddressPrunedLiveness. It does this on a per type tree leaf
+/// node basis.
+struct FieldSensitiveAddressPrunedLivenessBoundary {
+  /// The list of last users and an associated SILValue that is the address that
+  /// is being used. The address can be used to determine the start sub element
+  /// number of the user in the type tree and the end sub element number.
+  ///
+  /// TODO (MG): If we don't eventually need to store the SILValue here (I am
+  /// not sure yet...), just store a tuple with the start/end sub element
+  /// number.
+  SmallVector<std::tuple<SILInstruction *, TypeTreeLeafTypeRange>, 8> lastUsers;
+
+  /// Blocks where the value was live out but had a successor that was dead.
+  SmallVector<SILBasicBlock *, 8> boundaryEdges;
+
+  void clear() {
+    lastUsers.clear();
+    boundaryEdges.clear();
+  }
+
+  /// Compute the boundary from the blocks discovered during liveness analysis.
+  ///
+  /// Precondition: \p liveness.getDiscoveredBlocks() is a valid list of all
+  /// live blocks with no duplicates.
+  ///
+  /// The computed boundary will completely post-dominate, including dead end
+  /// paths. The client should query DeadEndBlocks to ignore those dead end
+  /// paths.
+  void compute(const FieldSensitiveAddressPrunedLiveness &liveness);
+
+private:
+  void
+  findLastUserInBlock(SILBasicBlock *bb,
+                      FieldSensitiveAddressPrunedLivenessBoundary &boundary,
+                      const FieldSensitiveAddressPrunedLiveness &liveness,
+                      unsigned subElementNumber);
+};
+
 } // namespace swift
 
 #endif