Improve funciton signature @owned return result to "not owned" conversion

More specifically, this handles a case of self-recursion.

With this change we convert 11 @owned return value to "not owned", while
we convert 179 @owned parameter to @guanrateed.

rdar://24022375

Author: trentxintong

include/swift/SILOptimizer/Analysis/ARCAnalysis.h

@@ -114,7 +114,13 @@ class ConsumedReturnValueToEpilogueRetainMatcher {
 public:
   enum class ExitKind { Return, Throw };
 
-  enum class FindRetainKind { None, Found, Blocked };
+  /// The state on how retains are found in a basic block.
+  enum class FindRetainKind { 
+    None,      ///< Did not find a retain.
+    Found,     ///< Found a retain.
+    Recursion, ///< Found a retain and its due to self-recursion.
+    Blocked    ///< Found a blocking instructions, i.e. MayDecrement.
+  };
 
   using RetainKindValue = std::pair<FindRetainKind, SILInstruction *>;
 

lib/SILOptimizer/Analysis/ARCAnalysis.cpp

@@ -26,6 +26,22 @@
 
 using namespace swift;
 
+using BasicBlockRetainValue = std::pair<SILBasicBlock *, SILValue>;
+
+
+//===----------------------------------------------------------------------===//
+//                            Utilities 
+//===----------------------------------------------------------------------===//
+
+static SILValue
+findIncomingValue(SILBasicBlock *BB, 
+     llvm::SmallVector<std::pair<SILBasicBlock *, SILValue>, 2> Values) {
+  for (auto &X : Values)
+    if (X.first == BB)
+      return X.second;
+  return SILValue();
+}
+
 //===----------------------------------------------------------------------===//
 //                             Decrement Analysis
 //===----------------------------------------------------------------------===//
@@ -481,7 +497,6 @@ void ConsumedReturnValueToEpilogueRetainMatcher::recompute() {
   findMatchingRetains(&*BB);
 }
 
-
 void
 ConsumedReturnValueToEpilogueRetainMatcher::
 findMatchingRetains(SILBasicBlock *BB) {
@@ -502,47 +517,69 @@ findMatchingRetains(SILBasicBlock *BB) {
   // OK. we've found the return value, now iterate on the CFG to find all the
   // post-dominating retains.
   constexpr unsigned WorkListMaxSize = 8;
-  RV = RCFI->getRCIdentityRoot(RV);
+
   llvm::DenseSet<SILBasicBlock *> RetainFrees;
-  llvm::SmallVector<SILBasicBlock *, 4> WorkList;
+  llvm::SmallVector<BasicBlockRetainValue, 4> WorkList;
   llvm::DenseSet<SILBasicBlock *> HandledBBs;
-  WorkList.push_back(BB);
+  WorkList.push_back(std::make_pair(BB, RV));
   HandledBBs.insert(BB);
   while (!WorkList.empty()) {
-    auto *CBB = WorkList.pop_back_val();
-    RetainKindValue Kind = findMatchingRetainsInner(CBB, RV);
-
     // Too many blocks ?.
     if (WorkList.size() > WorkListMaxSize) {
       EpilogueRetainInsts.clear();
       return;
     }
 
+    // Try to find a retain %value in this basic block.
+    auto BVP = WorkList.pop_back_val();
+    RetainKindValue Kind = findMatchingRetainsInner(BVP.first, BVP.second);
+
+    // We've found a retain on this path.
+    if (Kind.first == FindRetainKind::Found) { 
+      EpilogueRetainInsts.push_back(Kind.second);
+      continue;
+    }
+
     // There is a MayDecrement instruction.
     if (Kind.first == FindRetainKind::Blocked)
       return;
+
+    // There is a self-recursion. Use the apply instruction as the retain.
+    if (Kind.first == FindRetainKind::Recursion) {
+      EpilogueRetainInsts.push_back(Kind.second);
+      continue;
+    }
 
+    // Did not find a retain in this block, try to go to its predecessors.
     if (Kind.first == FindRetainKind::None) {
-      RetainFrees.insert(CBB);
-
       // We can not find a retain in a block with no predecessors.
-      if (CBB->getPreds().begin() == CBB->getPreds().end()) {
+      if (BVP.first->getPreds().begin() == BVP.first->getPreds().end()) {
         EpilogueRetainInsts.clear();
         return;
       }
 
+      // This block does not have a retain.
+      RetainFrees.insert(BVP.first);
+
       // Check the predecessors.
-      for (auto X : CBB->getPreds()){
+      bool FoundArgValues = false;
+      llvm::SmallVector<std::pair<SILBasicBlock *, SILValue>, 2> IncomingValues;
+      if (auto *SA = dyn_cast<SILArgument>(BVP.second))
+        FoundArgValues = SA->getIncomingValues(IncomingValues);
+
+      for (auto X : BVP.first->getPreds()){
         if (HandledBBs.find(X) != HandledBBs.end())
           continue;
-        WorkList.push_back(X);
+
+        // Try to use the predecessor edge-value.
+        if (FoundArgValues)
+          WorkList.push_back(std::make_pair(X, findIncomingValue(X, IncomingValues)));
+        else 
+          WorkList.push_back(std::make_pair(X, BVP.second));
+   
         HandledBBs.insert(X);
       }
     }
-
-    // We've found a retain on this path.
-    if (Kind.first == FindRetainKind::Found) 
-      EpilogueRetainInsts.push_back(Kind.second);
   }
 
   // For every block with retain, we need to check the transitive
@@ -573,6 +610,11 @@ ConsumedReturnValueToEpilogueRetainMatcher::RetainKindValue
 ConsumedReturnValueToEpilogueRetainMatcher::
 findMatchingRetainsInner(SILBasicBlock *BB, SILValue V) {
   for (auto II = BB->rbegin(), IE = BB->rend(); II != IE; ++II) {
+    // Handle self-recursion.
+    if (ApplyInst *AI = dyn_cast<ApplyInst>(&*II))
+      if (AI->getCalleeFunction() == BB->getParent()) 
+        return std::make_pair(FindRetainKind::Recursion, AI);
+    
     // If we do not have a retain_value or strong_retain...
     if (!isa<RetainValueInst>(*II) && !isa<StrongRetainInst>(*II)) {
       // we can ignore it if it can not decrement the reference count of the

lib/SILOptimizer/IPO/FunctionSignatureOpts.cpp

@@ -41,7 +41,7 @@ using namespace swift;
 STATISTIC(NumFunctionSignaturesOptimized, "Total func sig optimized");
 STATISTIC(NumDeadArgsEliminated, "Total dead args eliminated");
 STATISTIC(NumOwnedConvertedToGuaranteed, "Total owned args -> guaranteed args");
-STATISTIC(NumOwnedConvertedToGuaranteedReturnValue, "Total owned args -> guaranteed return value");
+STATISTIC(NumOwnedConvertedToNotOwnedResult, "Total owned result -> not owned result");
 STATISTIC(NumCallSitesOptimized, "Total call sites optimized");
 STATISTIC(NumSROAArguments, "Total SROA arguments optimized");
 
@@ -75,6 +75,7 @@ static bool hasIdenticalReleases(ReleaseList LHS, ReleaseList RHS) {
   return true;
 }
 
+
 /// Returns .Some(I) if I is a release that is the only non-debug instruction
 /// with side-effects in the use-def graph originating from Arg. Returns
 /// .Some(nullptr), if all uses from the arg were either debug insts or do not
@@ -654,7 +655,7 @@ bool SignatureAnalyzer::analyze() {
     if (!Retains.empty()) {
       RI.CalleeRetain = Retains;
       ShouldOptimize = true;
-      ++NumOwnedConvertedToGuaranteedReturnValue;
+      ++NumOwnedConvertedToNotOwnedResult;
     }
   }
 
@@ -781,11 +782,21 @@ SILFunction *SignatureOptimizer::createEmptyFunctionWithOptimizedSig(
 static void addRetainsForConvertedDirectResults(SILBuilder &Builder,
                                                 SILLocation Loc,
                                                 SILValue ReturnValue,
+                                                SILInstruction *AI,
                                      ArrayRef<ResultDescriptor> DirectResults) {
   for (auto I : indices(DirectResults)) {
     auto &RV = DirectResults[I];
     if (RV.CalleeRetain.empty()) continue;
 
+    bool IsSelfRecursionEpilogueRetain = false;
+    for (auto &X : RV.CalleeRetain) {
+      IsSelfRecursionEpilogueRetain |= (AI == X);
+    }
+
+    // We do not create an retain if this ApplyInst is a self-recursion.
+    if (IsSelfRecursionEpilogueRetain)
+      continue;
+
     // Extract the return value if necessary.
     SILValue SpecificResultValue = ReturnValue;
     if (DirectResults.size() != 1)
@@ -804,6 +815,7 @@ static void addRetainsForConvertedDirectResults(SILBuilder &Builder,
 static void rewriteApplyInstToCallNewFunction(SignatureOptimizer &Optimizer,
                                               SILFunction *NewF,
                                               const ApplyList &CallSites) {
+  llvm::DenseSet<SILInstruction *> ApplysToRemove;
   for (auto FAS : CallSites) {
     auto *AI = FAS.getInstruction();
 
@@ -865,15 +877,20 @@ static void rewriteApplyInstToCallNewFunction(SignatureOptimizer &Optimizer,
 
     // If we have converted the return value from @owned to @guaranteed,
     // insert a retain_value at the callsite.
-    addRetainsForConvertedDirectResults(Builder, Loc, ReturnValue,
+    addRetainsForConvertedDirectResults(Builder, Loc, ReturnValue, AI,
                                         Optimizer.getResultDescList());
-      
-    // Erase the old apply and its callee.
-    recursivelyDeleteTriviallyDeadInstructions(AI, true,
-                                               [](SILInstruction *) {});
 
+    // Make sure we remove this apply in the end.
+    ApplysToRemove.insert(AI);
     ++NumCallSitesOptimized;
   }
+
+  // Lastly, we have rewritten all the callsites, erase the old applys and
+  // its callee.
+  for (auto *AI : ApplysToRemove) {
+    recursivelyDeleteTriviallyDeadInstructions(AI, true,
+                                               [](SILInstruction *) {});
+  }
 }
 
 static void createThunkBody(SILBasicBlock *BB, SILFunction *NewF,
@@ -937,7 +954,7 @@ static void createThunkBody(SILBasicBlock *BB, SILFunction *NewF,
   }
 
   // Handle @owned to @unowned return value conversion.
-  addRetainsForConvertedDirectResults(Builder, Loc, ReturnValue,
+  addRetainsForConvertedDirectResults(Builder, Loc, ReturnValue, nullptr,
                                       Optimizer.getResultDescList());
 
   // Function that are marked as @NoReturn must be followed by an 'unreachable'
@@ -1056,6 +1073,8 @@ static bool optimizeFunctionSignature(llvm::BumpPtrAllocator &BPA,
   // to unowned conversion.
   for (ResultDescriptor &RD : Optimizer.getResultDescList()) {
     for (auto &X : RD.CalleeRetain) {
+      if (!isa<StrongRetainInst>(X) && !isa<RetainValueInst>(X))
+        continue;
       X->eraseFromParent();
     }
   }

test/SILOptimizer/functionsigopts.sil

@@ -112,6 +112,25 @@ bb0(%0 : $boo):
   return %0 : $boo
 }
 
+
+// CHECK-LABEL: sil [thunk] @single_owned_return_value_with_self_recursion
+// CHECK: function_ref @_TTSf4n_g_single_owned_return_value
+// CHECK: [[RET:%.*]] = apply
+// CHECK: retain_value [[RET]]
+sil @single_owned_return_value_with_self_recursion : $@convention(thin) (@owned boo) -> @owned boo {
+bb0(%0 : $boo):
+  cond_br undef, bb1, bb2
+bb1:
+  retain_value %0 : $boo
+  br bb3(%0 : $boo)
+bb2:
+  %2 = function_ref @single_owned_return_value_with_self_recursion : $@convention(thin) (@owned boo) -> @owned boo
+  %3 = apply %2(%0) : $@convention(thin) (@owned boo) -> @owned boo
+  br bb3 (%3 : $boo)
+bb3(%4 : $boo):
+  return %4 : $boo
+}
+
 // CHECK-LABEL: sil [thunk] @single_owned_return_value_with_interfering_release
 // CHECK: bb0([[INPUT_ARG0:%[0-9]+]] : $boo):
 // CHECK: [[IN1:%.*]] = function_ref @_TTSf4s__single_owned_return_value_with_interfering_release
@@ -245,6 +264,22 @@ bb0(%0 : $boo):
   return %5 : $()
 }
 
+// Make sure we pull out the retain_value from the callee.
+//
+// The retain and release can be got rid of easily by arc this way.
+//
+// CHECK-LABEL: sil @single_owned_return_value_with_self_recursion_callsite
+// CHECK: [[RET:%.*]] = apply
+// CHECK: retain_value [[RET]]
+// CHECK: release_value [[RET]]
+sil @single_owned_return_value_with_self_recursion_callsite : $@convention(thin) (@owned boo) -> () {
+bb0(%0 : $boo):
+  %2 = function_ref @single_owned_return_value_with_self_recursion : $@convention(thin) (@owned boo) -> @owned boo 
+  %4 = apply %2(%0) : $@convention(thin) (@owned boo) -> @owned boo 
+  release_value %4 : $boo
+  %5 = tuple()
+  return %5 : $()
+}
 
 sil [fragile] @exploded_release_to_dead_param_callsite : $@convention(thin) (@owned boo) -> () {
 bb0(%0 : $boo):
@@ -648,6 +683,13 @@ bb0(%0 : $Builtin.Int32):
 // CHECK-NOT: retain_value
 // CHECK: return
 
+// There should not be a single retain in this function.
+//
+// CHECK-LABEL: sil @_TTSf4n_g_single_owned_return_value_with_self_recursion : $@convention(thin) (@owned boo) -> boo
+// CHECK: bb0
+// CHECK-NOT: retain_value
+// CHECK: return
+
 // CHECK-LABEL: sil @_TTSf4s__single_owned_return_value_with_interfering_release : $@convention(thin) (@owned foo, @owned foo, Int) -> @owned boo {
 // CHECK: retain_value
 // CHECK: release_value