Fix operand ownership of mark_dependence [nonescaping] of address values

Author: meg-gupta

include/swift/SIL/OwnershipUtils.h

@@ -823,6 +823,7 @@ class InteriorPointerOperandKind {
     RefTailAddr,
     OpenExistentialBox,
     ProjectBox,
+    MarkDependenceNonEscaping,
   };
 
 private:
@@ -851,6 +852,12 @@ class InteriorPointerOperandKind {
       return Kind::OpenExistentialBox;
     case SILInstructionKind::ProjectBoxInst:
       return Kind::ProjectBox;
+    case SILInstructionKind::MarkDependenceInst: {
+      auto *mdi = cast<MarkDependenceInst>(use->getUser());
+      return mdi->isNonEscaping() && mdi->getType().isAddress()
+                 ? Kind::MarkDependenceNonEscaping
+                 : Kind::Invalid;
+    }
     }
   }
 
@@ -869,6 +876,12 @@ class InteriorPointerOperandKind {
       return Kind::OpenExistentialBox;
     case ValueKind::ProjectBoxInst:
       return Kind::ProjectBox;
+    case ValueKind::MarkDependenceInst: {
+      auto *mdi = cast<MarkDependenceInst>(value->getDefiningInstruction());
+      return mdi->isNonEscaping() && mdi->getType().isAddress()
+                 ? Kind::MarkDependenceNonEscaping
+                 : Kind::Invalid;
+    }
     }
   }
 
@@ -931,6 +944,13 @@ struct InteriorPointerOperand {
           &cast<SingleValueInstruction>(resultValue)->getAllOperands()[0];
       return InteriorPointerOperand(op, kind);
     }
+    case InteriorPointerOperandKind::MarkDependenceNonEscaping: {
+      auto *mdi =
+          cast<MarkDependenceInst>(resultValue->getDefiningInstruction());
+      assert(mdi->isNonEscaping() && mdi->getType().isAddress());
+      return InteriorPointerOperand(
+          &mdi->getAllOperands()[MarkDependenceInst::Base], kind);
+    }
     }
     llvm_unreachable("covered switch");
   }
@@ -968,6 +988,8 @@ struct InteriorPointerOperand {
       return cast<OpenExistentialBoxInst>(operand->getUser());
     case InteriorPointerOperandKind::ProjectBox:
       return cast<ProjectBoxInst>(operand->getUser());
+    case InteriorPointerOperandKind::MarkDependenceNonEscaping:
+      return cast<MarkDependenceInst>(operand->getUser());
     }
     llvm_unreachable("Covered switch isn't covered?!");
   }

lib/SIL/IR/OperandOwnership.cpp

@@ -672,9 +672,12 @@ OperandOwnershipClassifier::visitMarkDependenceInst(MarkDependenceInst *mdi) {
       /*allowUnowned*/true);
   }
   if (mdi->isNonEscaping()) {
-    // This creates a "dependent value", just like on-stack partial_apply, which
-    // we treat like a borrow.
-    return OperandOwnership::Borrow;
+    if (!mdi->getType().isAddress()) {
+      // This creates a "dependent value", just like on-stack partial_apply,
+      // which we treat like a borrow.
+      return OperandOwnership::Borrow;
+    }
+    return OperandOwnership::InteriorPointer;
   }
   if (mdi->hasUnresolvedEscape()) {
     // This creates a dependent value that may extend beyond the parent's

lib/SILOptimizer/SILCombiner/SILCombinerBuiltinVisitors.cpp

@@ -163,6 +163,7 @@ SILCombiner::optimizeBuiltinCOWBufferForReadingOSSA(BuiltinInst *bi) {
             return;
           case InteriorPointerOperandKind::OpenExistentialBox:
           case InteriorPointerOperandKind::ProjectBox:
+          case InteriorPointerOperandKind::MarkDependenceNonEscaping:
             // Can not mark this immutable.
             return;
           }

test/SILOptimizer/ossa_lifetime_completion.sil

@@ -956,3 +956,25 @@ exit:
 die:
   unreachable
 }
+
+class Klass {}
+
+struct Wrapper {
+  var c: Klass
+}
+
+// CHECK-LABEL: begin running test {{.*}} on testInteriorMarkDepNonEscAddressValue: ossa_lifetime_completion
+// CHECK-LABEL: sil [ossa] @testInteriorMarkDepNonEscAddressValue : {{.*}} {
+// CHECK: mark_dependence
+// CHECK: end_borrow
+// CHECK-LABEL: } // end sil function 'testInteriorMarkDepNonEscAddressValue'
+// CHECK-LABEL: end running test {{.*}} on testInteriorMarkDepNonEscAddressValue: ossa_lifetime_completion
+sil [ossa] @testInteriorMarkDepNonEscAddressValue : $@convention(thin) (@owned Wrapper, @inout Klass) -> () {
+bb0(%0 : @owned $Wrapper, %1 : $*Klass):
+  specify_test "ossa_lifetime_completion %2 liveness"
+  %2 = begin_borrow %0 : $Wrapper
+  %3 = struct_extract %2 : $Wrapper, #Wrapper.c
+  %4 = mark_dependence [nonescaping] %1 : $*Klass on %3 : $Klass
+  end_borrow %2 : $Wrapper
+  unreachable
+}