FieldSensitivePrunedLiveness: inject_enum_addr only initializes the tag of enums.

Fixes a crash when optional-chaining a move-only function that produces an
address-only value; FSPL would previously see the apply that initializes the
payload and the `inject_enum_tag` as overlapping defs, leading to mishandling
when trying to maximize the lifetime of the enum.

Author: jckarter

include/swift/SIL/FieldSensitivePrunedLiveness.h

@@ -315,12 +315,27 @@ struct TypeTreeLeafTypeRange {
       return std::nullopt;
 
     // A drop_deinit only consumes the deinit bit of its operand.
-    auto *ddi = dyn_cast<DropDeinitInst>(op->getUser());
-    if (ddi) {
+    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.
@@ -984,7 +999,7 @@ class FieldSensitivePrunedLiveness {
     return *&iter->second;
   }
 
-  /// If \p user has had uses recored, return a pointer to the InterestingUser
+  /// If \p user has had uses recorded, return a pointer to the InterestingUser
   /// where they've been recorded.
   InterestingUser const *getInterestingUser(SILInstruction *user) const {
     auto iter = users.find(user);

test/SILGen/moveonly_optional_operations.swift

@@ -23,6 +23,9 @@ struct NC: ~Copyable {
     borrowing func b() {}
     mutating func m() {}
     consuming func c() {}
+
+    consuming func c2() -> NC { c2() } // expected-warning{{}}
+    consuming func c3() -> NCAO { c3() } // expected-warning{{}}
 }
 
 struct NCAO: ~Copyable {
@@ -31,6 +34,9 @@ struct NCAO: ~Copyable {
     borrowing func b() {}
     mutating func m() {}
     consuming func c() {}
+
+    consuming func c2() -> NC { c2() } // expected-warning{{}}
+    consuming func c3() -> NCAO { c3() } // expected-warning{{}}
 }
 
 func borrowingChains(nc: borrowing NC?, // expected-error{{'nc' is borrowed and cannot be consumed}}
@@ -140,3 +146,27 @@ func consumingChains2(nc: consuming NC?,
     ncao?.m()
     ncao?.c()
 }
+
+func consumingSwitchSubject(nc: consuming NC?,
+                            ncao: consuming NCAO?) {
+    switch nc?.c2() {
+    default:
+      break
+    }
+    switch ncao?.c2() {
+    default:
+      break
+    }
+}
+
+func consumingSwitchSubject2(nc: consuming NC?,
+                             ncao: consuming NCAO?) {
+    switch nc?.c3() {
+    default:
+      break
+    }
+    switch ncao?.c3() {
+    default:
+      break
+    }
+}