[mandatory-inlining] Fix lifetime extension error.

This bug is caused by a quirk in the API of the linear lifetime
checker. Specifically, even though valueHasLinearLifetime is passed a SILValue
(the value whose lifetime one is checking), really it doesnt care about that
value (except for error diagnostics). Really it just cares about the parent
block of the value since it assumes that the value is guaranteed to dominate all
uses.

This creates a footgun when if one is writing code using "generic ossa/non-ossa"
routines on SILBuilder (the emit*Operation methods), if one in non-ossa code
calls that function, it returns the input value of the strong_retain. This
causes the linear lifetime error, to use the parent block of the argument of the
retain, instead of the parent block of the retain itself. This then causes it to
find the wrong leaking blocks and thus insert destroys in the wrong places.

I fix this problem in this commit by noting that the partial apply is our
original insertion point for the copy, so of course it is going to be in the
same block. So I changed the linear lifetime checker to check for leaks with
respect to the partial applies result.

In a subsequent commit, I am going to add a new API on top of this that is based
around the use of the value by the partial apply (maybe
extendLifetimeFromUseToInsertionPoint?). By using the use, it will express in
code more clearly what is happening here and will insert the copy for you.

rdar://54234011

Author: gottesmm

lib/SILOptimizer/Mandatory/MandatoryInlining.cpp

@@ -137,9 +137,21 @@ static void fixupReferenceCounts(
       }
 
       visitedBlocks.clear();
+
       // If we need to insert compensating destroys, do so.
+      //
+      // NOTE: We use pai here since in non-ossa code emitCopyValueOperation
+      // returns the operand of the strong_retain which may have a ValueBase
+      // that is not in the same block. An example of where this is important is
+      // if we are performing emitCopyValueOperation in non-ossa code on an
+      // argument when the partial_apply is not in the entrance block. In truth,
+      // the linear lifetime checker does not /actually/ care what the value is
+      // (ignoring diagnostic error msgs that we do not care about here), it
+      // just cares about the block the value is in. In a forthcoming commit, I
+      // am going to change this to use a different API on the linear lifetime
+      // checker that makes this clearer.
       auto error =
-          valueHasLinearLifetime(copy, {applySite}, {}, visitedBlocks,
+          valueHasLinearLifetime(pai, {applySite}, {}, visitedBlocks,
                                  deadEndBlocks, errorBehavior, &leakingBlocks);
       if (error.getFoundLeak()) {
         while (!leakingBlocks.empty()) {
@@ -175,11 +187,22 @@ static void fixupReferenceCounts(
     // TODO: Do we need to lifetime extend here?
     case ParameterConvention::Direct_Unowned: {
       v = SILBuilderWithScope(pai).emitCopyValueOperation(loc, v);
-
       visitedBlocks.clear();
+
       // If we need to insert compensating destroys, do so.
+      //
+      // NOTE: We use pai here since in non-ossa code emitCopyValueOperation
+      // returns the operand of the strong_retain which may have a ValueBase
+      // that is not in the same block. An example of where this is important is
+      // if we are performing emitCopyValueOperation in non-ossa code on an
+      // argument when the partial_apply is not in the entrance block. In truth,
+      // the linear lifetime checker does not /actually/ care what the value is
+      // (ignoring diagnostic error msgs that we do not care about here), it
+      // just cares about the block the value is in. In a forthcoming commit, I
+      // am going to change this to use a different API on the linear lifetime
+      // checker that makes this clearer.
       auto error =
-          valueHasLinearLifetime(v, {applySite}, {}, visitedBlocks,
+          valueHasLinearLifetime(pai, {applySite}, {}, visitedBlocks,
                                  deadEndBlocks, errorBehavior, &leakingBlocks);
       if (error.getFoundError()) {
         while (!leakingBlocks.empty()) {
@@ -203,11 +226,22 @@ static void fixupReferenceCounts(
     // apply has another use that would destroy our value first.
     case ParameterConvention::Direct_Owned: {
       v = SILBuilderWithScope(pai).emitCopyValueOperation(loc, v);
-
       visitedBlocks.clear();
+
       // If we need to insert compensating destroys, do so.
+      //
+      // NOTE: We use pai here since in non-ossa code emitCopyValueOperation
+      // returns the operand of the strong_retain which may have a ValueBase
+      // that is not in the same block. An example of where this is important is
+      // if we are performing emitCopyValueOperation in non-ossa code on an
+      // argument when the partial_apply is not in the entrance block. In truth,
+      // the linear lifetime checker does not /actually/ care what the value is
+      // (ignoring diagnostic error msgs that we do not care about here), it
+      // just cares about the block the value is in. In a forthcoming commit, I
+      // am going to change this to use a different API on the linear lifetime
+      // checker that makes this clearer.
       auto error =
-          valueHasLinearLifetime(v, {applySite}, {}, visitedBlocks,
+          valueHasLinearLifetime(pai, {applySite}, {}, visitedBlocks,
                                  deadEndBlocks, errorBehavior, &leakingBlocks);
       if (error.getFoundError()) {
         while (!leakingBlocks.empty()) {

test/Interpreter/mandatory_inlining.swift

@@ -2,20 +2,60 @@
 
 // REQUIRES: executable_test
 
-func test(reportError: ((String) -> (Void))? = nil) {
-  let reportError = reportError ?? { error in
-    print(error)
+// This file contains various scenarios that validate that various functional
+// behaviors of mandatory inlining work as expected. Please give each test a
+// descriptive name and a large comment explaining what you are testing.
+
+import StdlibUnittest
+
+var Tests = TestSuite("MandatoryInlining Functional Tests")
+defer { runAllTests() }
+
+func doNotInsertReleaseInLoopIfValueConsumedInLoop(
+  callBack: ((String) -> (Void))? = nil) {
+  let callBack = callBack ?? { msg in
+    print(msg)
   }
 
   for _ in 0..<1 {
-    reportError("foo bar baz")
+    callBack("foo bar baz")
   }
 }
 
-func main() {
-  test { error in
-      print(error)
+// When using the linear lifetime checker to insert compensating releases, if we
+// find a double consume due to a loop, do not insert an apply at that call site
+// that is in the loop. There is already a compensating retain in the loop that
+// we are ignoring. This test functionally makes sure we no longer do this.
+Tests.test("doNotInsertReleaseInLoopIfValueConsumedInLoop") {
+  doNotInsertReleaseInLoopIfValueConsumedInLoop { msg in
+      print(msg)
   }
 }
 
-main()
+
+@propertyWrapper
+public struct Published<Value> {
+    public var wrappedValue: Value
+    public init(initialValue value: Value) { wrappedValue = value }
+}
+
+private class Model {
+    var selectedContact: Int? {
+        get { _selectedContact }
+        set {
+            if let contact = newValue {
+                self._selectedContact = contact
+            }
+        }
+    }
+
+    @Published private var _selectedContact: Int? = nil
+}
+
+Tests.test("testNonOSSAPartialApplyLifetimeExtension") {
+  // If we mess this up, when we set model.selectedContact to nil, we release
+  // self incorrectly due to mandatory inlining getting the wrong begin lifetime
+  // block (which in this case is the entry block).
+  let model = Model()
+  model.selectedContact = nil
+}

test/SILOptimizer/mandatory_inlining.sil

@@ -1320,7 +1320,7 @@ bb5:
 // CHECK-NEXT: strong_release [[ARG]]
 // CHECK-NEXT: tuple
 // CHECK-NEXT: return
-// CHECK-NEXT: } // end sil function 'test_apply_pai_in_diamond'
+// CHECK: } // end sil function 'test_apply_pai_in_diamond'
 sil @test_apply_pai_in_diamond : $@convention(thin) (@guaranteed C) ->  () {
 bb0(%0: $C):
   strong_retain %0 : $C
@@ -1340,3 +1340,64 @@ bb3:
   %tuple = tuple()
   return %tuple : $()
 }
+
+// CHECK-LABEL: sil @test_apply_pai_in_diamond_non_post_dominate_non_post_dominate : $@convention(thin) (@guaranteed C) ->  () {
+// CHECK: bb0([[ARG:%.*]] : $C):
+// CHECK-NEXT:   cond_br undef, [[BB_TRAMPOLINE:bb[0-9]+]], [[BB_DIAMOND_BEGIN:bb[0-9]+]]
+//
+// CHECK: [[BB_TRAMPOLINE]]:
+// CHECK-NEXT: br [[BB_END:bb[0-9]+]]
+//
+// CHECK: [[BB_DIAMOND_BEGIN]]:
+// CHECK-NEXT:   strong_retain [[ARG]]
+// CHECK-NEXT:   strong_retain [[ARG]]
+// CHECK-NEXT:   cond_br undef, [[NO_CALL_BB:bb[0-9]+]], [[CALL_BB:bb[0-9]+]]
+//
+// CHECK: [[NO_CALL_BB]]:
+// CHECK-NEXT: strong_release [[ARG]]
+// CHECK-NEXT: br [[MERGE_BB:bb[0-9]+]]
+//
+// CHECK: [[CALL_BB]]:
+// CHECK-NEXT:  function_ref use_c
+// CHECK-NEXT:  [[FUNC:%.*]] = function_ref @use_c : $@convention(thin)
+// CHECK-NEXT:  apply [[FUNC]]
+// CHECK-NEXT:  tuple
+// CHECK-NEXT:  strong_release [[ARG]]
+// CHECK-NEXT:  br [[MERGE_BB]]
+//
+// CHECK: [[MERGE_BB]]:
+// CHECK-NEXT: strong_release [[ARG]]
+// CHECK-NEXT: br [[BB_END]]
+//
+// CHECK: [[BB_END]]:
+// CHECK-NEXT: tuple
+// CHECK-NEXT: return
+// CHECK: } // end sil function 'test_apply_pai_in_diamond_non_post_dominate_non_post_dominate'
+sil @test_apply_pai_in_diamond_non_post_dominate_non_post_dominate : $@convention(thin) (@guaranteed C) ->  () {
+bb0(%0: $C):
+  cond_br undef, bb1, bb2
+
+bb1:
+  br bb6
+
+bb2:
+  strong_retain %0 : $C
+  %closure_fun = function_ref @guaranteed_closure_func : $@convention(thin) (@guaranteed C) -> ()
+  %closure = partial_apply [callee_guaranteed] %closure_fun(%0) : $@convention(thin) (@guaranteed C) -> ()
+  cond_br undef, bb3, bb4
+
+bb3:
+  br bb5
+
+bb4:
+  apply %closure() : $@callee_guaranteed () -> ()
+  br bb5
+
+bb5:
+  strong_release %closure : $@callee_guaranteed () -> ()
+  br bb6
+
+bb6:
+  %tuple = tuple()
+  return %tuple : $()
+}