[field-sensitive-pruned-liveness] Some small changes.

gottesmm · gottesmm · commit 4ad2bc93bad7 · 2023-01-24T15:28:12.000-08:00
Specifically:

1. I added to the documentation at the top of the file that our representation
allows for partial init/reinit of structs/tuples from parts.

2. I renamed SubElementNumber to SubElementOffset. This I think fits the actual
use case better and makes it clearer what one is working with (the offset inside
a type of a subelement of the type).

3. I added some small helpers to TypeSubElementCount and SubElementOffset for
adding/subtracting from them.

4. I added the ability to iterate over just consuming/nonconsuming users in
FieldSensitivePrunedLiveness. Just a useful little helper.
diff --git a/include/swift/SIL/FieldSensitivePrunedLiveness.h b/include/swift/SIL/FieldSensitivePrunedLiveness.h
@@ -16,7 +16,9 @@
 #include "swift/AST/TypeExpansionContext.h"
 #include "swift/Basic/Debug.h"
 #include "swift/Basic/FrozenMultiMap.h"
+#include "swift/Basic/STLExtras.h"
 #include "swift/SIL/PrunedLiveness.h"
+#include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallBitVector.h"
 #include "llvm/Support/raw_ostream.h"
 
@@ -59,6 +61,13 @@ namespace swift {
 /// [0, 1, 1, 0, 0]
 /// ```
 ///
+/// NOTE: Our representation allows for partial initialization/reinitialization
+/// since we do not include a bit for each level of struct/tuple. The effect of
+/// this is that we cannot distinguish a single field type from its child
+/// field. This makes this type not suited for projection operations. This is a
+/// trade-off that was made to make it easy to support partial
+/// initialization/reinitialization of liveness.
+///
 /// Linearized Representation of Enums
 /// ----------------------------------
 ///
@@ -141,28 +150,36 @@ namespace swift {
 /// we would only use a single bit in our linearized representation, just for
 /// the discriminator value.
 ///
-/// Enums and Partial Initialization
-/// --------------------------------
+/// Enums and Partial Initialization/Deinitialization
+/// -------------------------------------------------
 ///
 /// One property of our representation of structs and tuples is that a code
-/// generator can reinitialize a struct/tuple completely just by re-initializing
-/// each of its sub-types individually. This is not possible for enums in our
-/// representation since if one just took the leaf nodes for the payload, one
-/// would not update the bit for the enum case itself and any additional spare
-/// bits. Luckily for us, this is actually impossible to do in SIL since it is
-/// impossible to dynamically change the payload of an enum without destroying
-/// the original enum and its payload since that would be a verifier caught
-/// leak.
-struct SubElementNumber {
+/// generator can init/reinit/deinit a struct/tuple completely just by
+/// performing the relevant operation on each of its sub-types
+/// individually. This is not possible for enums in our representation since if
+/// one just took the leaf nodes for the payload, one would not update the bit
+/// for the enum case itself and any additional spare bits. Luckily for us, this
+/// is safe to assume since we only use this in misc address contexts and move
+/// only object contexts in Raw SIL which have the following invariants:
+///
+/// 1. SIL addresses of enum type cannot dynamically change the payload of an
+///    enum without destroying the original enum completely. So such an address
+///    can never be reinitialized by storing into the payload of an enum.
+///
+/// 2. It is illegal in SIL to unchecked_enum_data a move only type in Raw
+///    SIL. One must instead use a switch_enum which creates a new value for the
+///    destructured enum. We when writing such verifiers consider the switch to
+///    produce an entire new value rather than a derived forwarding value.
+struct SubElementOffset {
   /// Our internal sub element representation number. We force 32 bits so that
   /// our type tree span is always pointer width. This is convenient for storing
   /// it in other data structures.
   uint32_t number;
 
-  SubElementNumber(unsigned number) : number(number) {}
+  SubElementOffset(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
+  /// rootValue, 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.
   ///
@@ -175,7 +192,7 @@ struct SubElementNumber {
   ///
   /// \returns None if we didn't know how to compute sub-element for this
   /// projection.
-  static Optional<SubElementNumber> compute(SILValue projectionFromRoot,
+  static Optional<SubElementOffset> compute(SILValue projectionFromRoot,
                                             SILValue root) {
     assert(projectionFromRoot->getType().getCategory() ==
                root->getType().getCategory() &&
@@ -187,10 +204,15 @@ struct SubElementNumber {
 
   operator unsigned() const { return number; }
 
+  SubElementOffset &operator+=(unsigned other) {
+    number += other;
+    return *this;
+  }
+
 private:
-  static Optional<SubElementNumber>
+  static Optional<SubElementOffset>
   computeForAddress(SILValue projectionFromRoot, SILValue rootAddress);
-  static Optional<SubElementNumber> computeForValue(SILValue projectionFromRoot,
+  static Optional<SubElementOffset> computeForValue(SILValue projectionFromRoot,
                                                     SILValue rootValue);
 };
 
@@ -226,6 +248,11 @@ struct TypeSubElementCount {
                             TypeExpansionContext(*value->getFunction())) {}
 
   operator unsigned() const { return number; }
+
+  TypeSubElementCount operator-=(unsigned other) {
+    *this = TypeSubElementCount(unsigned(*this) - other);
+    return *this;
+  }
 };
 
 class FieldSensitivePrunedLiveness;
@@ -234,18 +261,18 @@ class FieldSensitivePrunedLiveness;
 /// of the type tree of a type T.
 struct TypeTreeLeafTypeRange {
   friend FieldSensitivePrunedLiveness;
-  SubElementNumber startEltOffset;
-  SubElementNumber endEltOffset;
 
-public:
+  SubElementOffset startEltOffset;
+  SubElementOffset endEltOffset;
+
   TypeTreeLeafTypeRange() : startEltOffset(0), endEltOffset(0) {}
 
-  TypeTreeLeafTypeRange(SubElementNumber start, SubElementNumber end)
+  TypeTreeLeafTypeRange(SubElementOffset start, SubElementOffset end)
       : startEltOffset(start), endEltOffset(end) {}
 
   /// The leaf type range for the entire type tree.
-  TypeTreeLeafTypeRange(SILValue rootAddress)
-      : startEltOffset(0), endEltOffset(TypeSubElementCount(rootAddress)) {}
+  TypeTreeLeafTypeRange(SILValue rootValue)
+      : startEltOffset(0), endEltOffset(TypeSubElementCount(rootValue)) {}
 
   /// The leaf type range for the entire type tree.
   TypeTreeLeafTypeRange(SILType rootType, SILFunction *fn)
@@ -257,14 +284,13 @@ struct TypeTreeLeafTypeRange {
   ///
   /// \returns None if we are unable to understand the path in between \p
   /// projectedAddress and \p rootAddress.
-  static Optional<TypeTreeLeafTypeRange> get(SILValue projectedAddress,
-                                             SILValue rootAddress) {
-    auto startEltOffset =
-        SubElementNumber::compute(projectedAddress, rootAddress);
+  static Optional<TypeTreeLeafTypeRange> get(SILValue projectedValue,
+                                             SILValue rootValue) {
+    auto startEltOffset = SubElementOffset::compute(projectedValue, rootValue);
     if (!startEltOffset)
       return None;
     return {{*startEltOffset,
-             *startEltOffset + TypeSubElementCount(projectedAddress)}};
+             *startEltOffset + TypeSubElementCount(projectedValue)}};
   }
 
   /// Given a type \p rootType and a set of needed elements specified by the bit
@@ -311,7 +337,7 @@ struct TypeTreeLeafTypeRange {
 
   /// 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 {
+  bool contains(SubElementOffset singleLeafElementNumber) const {
     return startEltOffset <= singleLeafElementNumber &&
            singleLeafElementNumber < endEltOffset;
   }
@@ -343,6 +369,12 @@ struct TypeTreeLeafTypeRange {
   void dump() const { print(llvm::dbgs()); }
 };
 
+inline llvm::raw_ostream &operator<<(llvm::raw_ostream &os,
+                                     const TypeTreeLeafTypeRange &value) {
+  value.print(os);
+  return os;
+}
+
 /// 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.
@@ -353,10 +385,12 @@ struct TypeTreeLeafTypeRange {
 class FieldSensitivePrunedLiveness {
   PrunedLiveBlocks liveBlocks;
 
+public:
   struct InterestingUser {
     TypeTreeLeafTypeRange subEltSpan;
     bool isConsuming;
 
+    InterestingUser() : subEltSpan(), isConsuming(false) {}
     InterestingUser(TypeTreeLeafTypeRange subEltSpan, bool isConsuming)
         : subEltSpan(subEltSpan), isConsuming(isConsuming) {}
 
@@ -366,6 +400,7 @@ class FieldSensitivePrunedLiveness {
     }
   };
 
+private:
   /// 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.
@@ -420,6 +455,40 @@ class FieldSensitivePrunedLiveness {
     return llvm::make_range(users.begin(), users.end());
   }
 
+  using LifetimeEndingUserRange = OptionalTransformRange<
+      UserRange,
+      function_ref<Optional<std::pair<SILInstruction *, TypeTreeLeafTypeRange>>(
+          const std::pair<SILInstruction *, InterestingUser> &)>>;
+  LifetimeEndingUserRange getAllLifetimeEndingUses() const {
+    function_ref<Optional<std::pair<SILInstruction *, TypeTreeLeafTypeRange>>(
+        const std::pair<SILInstruction *, InterestingUser> &)>
+        op;
+    op = [](const std::pair<SILInstruction *, InterestingUser> &pair)
+        -> Optional<std::pair<SILInstruction *, TypeTreeLeafTypeRange>> {
+      if (pair.second.isConsuming)
+        return {{pair.first, pair.second.subEltSpan}};
+      return None;
+    };
+    return LifetimeEndingUserRange(getAllUsers(), op);
+  }
+
+  using NonLifetimeEndingUserRange = OptionalTransformRange<
+      UserRange,
+      function_ref<Optional<std::pair<SILInstruction *, TypeTreeLeafTypeRange>>(
+          const std::pair<SILInstruction *, InterestingUser> &)>>;
+  NonLifetimeEndingUserRange getAllNonLifetimeEndingUses() const {
+    function_ref<Optional<std::pair<SILInstruction *, TypeTreeLeafTypeRange>>(
+        const std::pair<SILInstruction *, InterestingUser> &)>
+        op;
+    op = [](const std::pair<SILInstruction *, InterestingUser> &pair)
+        -> Optional<std::pair<SILInstruction *, TypeTreeLeafTypeRange>> {
+      if (!pair.second.isConsuming)
+        return {{pair.first, pair.second.subEltSpan}};
+      return None;
+    };
+    return NonLifetimeEndingUserRange(getAllUsers(), op);
+  }
+
   using UserBlockRange = TransformRange<
       UserRange, function_ref<SILBasicBlock *(
                      const std::pair<SILInstruction *, InterestingUser> &)>>;
diff --git a/lib/SIL/Utils/FieldSensitivePrunedLiveness.cpp b/lib/SIL/Utils/FieldSensitivePrunedLiveness.cpp
@@ -88,16 +88,16 @@ TypeSubElementCount::TypeSubElementCount(SILType type, SILModule &mod,
 //                           MARK: SubElementNumber
 //===----------------------------------------------------------------------===//
 
-Optional<SubElementNumber>
-SubElementNumber::computeForAddress(SILValue projectionDerivedFromRoot,
+Optional<SubElementOffset>
+SubElementOffset::computeForAddress(SILValue projectionDerivedFromRoot,
                                     SILValue rootAddress) {
-  unsigned finalSubElementNumber = 0;
+  unsigned finalSubElementOffset = 0;
   SILModule &mod = *rootAddress->getModule();
 
   while (1) {
     // If we got to the root, we're done.
     if (rootAddress == projectionDerivedFromRoot)
-      return {SubElementNumber(finalSubElementNumber)};
+      return {SubElementOffset(finalSubElementOffset)};
 
     if (auto *pbi = dyn_cast<ProjectBoxInst>(projectionDerivedFromRoot)) {
       projectionDerivedFromRoot = pbi->getOperand();
@@ -115,7 +115,7 @@ SubElementNumber::computeForAddress(SILValue projectionDerivedFromRoot,
 
       // Keep track of what subelement is being referenced.
       for (unsigned i : range(teai->getFieldIndex())) {
-        finalSubElementNumber += TypeSubElementCount(
+        finalSubElementOffset += TypeSubElementCount(
             tupleType.getTupleElementType(i), mod,
             TypeExpansionContext(*rootAddress->getFunction()));
       }
@@ -133,7 +133,7 @@ SubElementNumber::computeForAddress(SILValue projectionDerivedFromRoot,
         if (fieldDecl == seai->getField())
           break;
         auto context = TypeExpansionContext(*rootAddress->getFunction());
-        finalSubElementNumber += TypeSubElementCount(
+        finalSubElementOffset += TypeSubElementCount(
             type.getFieldType(fieldDecl, mod, context), mod, context);
       }
 
@@ -175,38 +175,75 @@ SubElementNumber::computeForAddress(SILValue projectionDerivedFromRoot,
   }
 }
 
-Optional<SubElementNumber>
-SubElementNumber::computeForValue(SILValue projectionDerivedFromRoot,
+Optional<SubElementOffset>
+SubElementOffset::computeForValue(SILValue projectionDerivedFromRoot,
                                   SILValue rootAddress) {
-  unsigned finalSubElementNumber = 0;
+  unsigned finalSubElementOffset = 0;
   SILModule &mod = *rootAddress->getModule();
 
   while (1) {
     // If we got to the root, we're done.
     if (rootAddress == projectionDerivedFromRoot)
-      return {SubElementNumber(finalSubElementNumber)};
+      return {SubElementOffset(finalSubElementOffset)};
 
     // Look through these single operand instructions.
     if (isa<BeginBorrowInst>(projectionDerivedFromRoot) ||
         isa<CopyValueInst>(projectionDerivedFromRoot)) {
       projectionDerivedFromRoot =
           cast<SingleValueInstruction>(projectionDerivedFromRoot)
               ->getOperand(0);
+      continue;
     }
 
     if (auto *teai = dyn_cast<TupleExtractInst>(projectionDerivedFromRoot)) {
       SILType tupleType = teai->getOperand()->getType();
 
       // Keep track of what subelement is being referenced.
       for (unsigned i : range(teai->getFieldIndex())) {
-        finalSubElementNumber += TypeSubElementCount(
+        finalSubElementOffset += TypeSubElementCount(
             tupleType.getTupleElementType(i), mod,
             TypeExpansionContext(*rootAddress->getFunction()));
       }
       projectionDerivedFromRoot = teai->getOperand();
       continue;
     }
 
+    if (auto *mvir = dyn_cast<MultipleValueInstructionResult>(
+            projectionDerivedFromRoot)) {
+      if (auto *dsi = dyn_cast<DestructureStructInst>(mvir->getParent())) {
+        SILType type = dsi->getOperand()->getType();
+
+        // Keep track of what subelement is being referenced.
+        unsigned resultIndex = mvir->getIndex();
+        StructDecl *structDecl = dsi->getStructDecl();
+        for (auto pair : llvm::enumerate(structDecl->getStoredProperties())) {
+          if (pair.index() == resultIndex)
+            break;
+          auto context = TypeExpansionContext(*rootAddress->getFunction());
+          finalSubElementOffset += TypeSubElementCount(
+              type.getFieldType(pair.value(), mod, context), mod, context);
+        }
+
+        projectionDerivedFromRoot = dsi->getOperand();
+        continue;
+      }
+
+      if (auto *dti = dyn_cast<DestructureTupleInst>(mvir->getParent())) {
+        SILType type = dti->getOperand()->getType();
+
+        // Keep track of what subelement is being referenced.
+        unsigned resultIndex = mvir->getIndex();
+        for (unsigned i : range(resultIndex)) {
+          auto context = TypeExpansionContext(*rootAddress->getFunction());
+          finalSubElementOffset +=
+              TypeSubElementCount(type.getTupleElementType(i), mod, context);
+        }
+
+        projectionDerivedFromRoot = dti->getOperand();
+        continue;
+      }
+    }
+
     if (auto *seai = dyn_cast<StructExtractInst>(projectionDerivedFromRoot)) {
       SILType type = seai->getOperand()->getType();
 
@@ -216,7 +253,7 @@ SubElementNumber::computeForValue(SILValue projectionDerivedFromRoot,
         if (fieldDecl == seai->getField())
           break;
         auto context = TypeExpansionContext(*rootAddress->getFunction());
-        finalSubElementNumber += TypeSubElementCount(
+        finalSubElementOffset += TypeSubElementCount(
             type.getFieldType(fieldDecl, mod, context), mod, context);
       }
 
