[caller-analysis] Add a new SIL utility primitive: findLocalApplySites.

This utility works by taking in a function_ref and then traverses the transitive
uses of the function_ref until it finds either a use it does not understand
"escape" or an "apply" instruction. It returns a result structure that contains
the final found applications and more importantly a bool telling the caller if
we found any "escaping" uses.

This is intended to be an inverse operation to ApplySite::getCalleeOrigin(). As
such it has a bunch of assertions in it that check that the two stay in sync.

rdar://41146023

Author: gottesmm

include/swift/SILOptimizer/Utils/Local.h

@@ -643,6 +643,52 @@ class StaticInitCloner : public SILCloner<StaticInitCloner> {
   }
 };
 
+/// Move only data structure that is the result of findLocalApplySite.
+///
+/// NOTE: Generally it is not suggested to have move only types that contain
+/// small vectors. Since our small vectors contain one element or a std::vector
+/// like data structure , this is ok since we will either just copy the single
+/// element when we do the move or perform a move of the vector type.
+struct LLVM_LIBRARY_VISIBILITY FindLocalApplySitesResult {
+  /// Contains the list of local non fully applied partial apply sites that we
+  /// found.
+  SmallVector<ApplySite, 1> partialApplySites;
+
+  /// Contains the list of full apply sites that we found.
+  SmallVector<FullApplySite, 1> fullApplySites;
+
+  /// Set to true if the function_ref escapes into a use that our analysis does
+  /// not understand. Set to false if we found a use that had an actual
+  /// escape. Set to None if we did not find any call sites, but also didn't
+  /// find any "escaping uses" as well.
+  ///
+  /// The none case is so that we can distinguish in between saying that a value
+  /// did escape and saying that we did not find any conservative information.
+  bool escapes;
+
+  FindLocalApplySitesResult() = default;
+  FindLocalApplySitesResult(const FindLocalApplySitesResult &) = delete;
+  FindLocalApplySitesResult &
+  operator=(const FindLocalApplySitesResult &) = delete;
+  FindLocalApplySitesResult(FindLocalApplySitesResult &&) = default;
+  FindLocalApplySitesResult &operator=(FindLocalApplySitesResult &&) = default;
+  ~FindLocalApplySitesResult() = default;
+
+  /// Treat this function ref as escaping only if we found an actual user we
+  /// didn't understand. Do not treat it as escaping if we did not find any
+  /// users at all.
+  bool isEscaping() const { return escapes; }
+};
+
+/// Returns .some(FindLocalApplySitesResult) if we found any interesting
+/// information for the given function_ref. Otherwise, returns None.
+///
+/// We consider "interesting information" to mean inclusively that:
+///
+/// 1. We discovered that the function_ref never escapes.
+/// 2. We were able to find either a partial apply or a full apply site.
+Optional<FindLocalApplySitesResult> findLocalApplySites(FunctionRefInst *FRI);
+
 } // end namespace swift
 
 #endif

lib/SILOptimizer/Utils/Local.cpp

@@ -1586,4 +1586,72 @@ StaticInitCloner::clone(SingleValueInstruction *InitVal) {
   return cast<SingleValueInstruction>(ValueMap[InitVal]);
 }
 
+Optional<FindLocalApplySitesResult>
+swift::findLocalApplySites(FunctionRefInst *FRI) {
+  SmallVector<Operand *, 32> worklist(FRI->use_begin(), FRI->use_end());
 
+  FindLocalApplySitesResult f;
+
+  // Optimistically state that we have no escapes before our def-use dataflow.
+  f.escapes = false;
+
+  while (!worklist.empty()) {
+    auto *op = worklist.pop_back_val();
+    auto *user = op->getUser();
+
+    // If we have a full apply site as our user.
+    if (auto apply = FullApplySite::isa(user)) {
+      if (apply.getCallee() == op->get()) {
+        f.fullApplySites.push_back(apply);
+        continue;
+      }
+    }
+
+    // If we have a partial apply as a user, start tracking it, but also look at
+    // its users.
+    if (auto *pai = dyn_cast<PartialApplyInst>(user)) {
+      if (pai->getCallee() == op->get()) {
+        // Track the partial apply that we saw so we can potentially eliminate
+        // dead closure arguments.
+        f.partialApplySites.push_back(pai);
+        // Look to see if we can find a full application of this partial apply
+        // as well.
+        copy(pai->getUses(), std::back_inserter(worklist));
+        continue;
+      }
+    }
+
+    // Otherwise, see if we have any function casts to look through...
+    switch (user->getKind()) {
+    case SILInstructionKind::ThinToThickFunctionInst:
+    case SILInstructionKind::ConvertFunctionInst:
+    case SILInstructionKind::ConvertEscapeToNoEscapeInst:
+      copy(cast<SingleValueInstruction>(user)->getUses(),
+           std::back_inserter(worklist));
+      continue;
+
+    // Look through any reference count instructions since these are not
+    // escapes:
+    case SILInstructionKind::CopyValueInst:
+      copy(cast<CopyValueInst>(user)->getUses(), std::back_inserter(worklist));
+      continue;
+    case SILInstructionKind::StrongRetainInst:
+    case SILInstructionKind::StrongReleaseInst:
+    case SILInstructionKind::RetainValueInst:
+    case SILInstructionKind::ReleaseValueInst:
+    case SILInstructionKind::DestroyValueInst:
+      continue;
+    default:
+      break;
+    }
+
+    // But everything else is considered an escape.
+    f.escapes = true;
+  }
+
+  // If we did escape and didn't find any apply sites, then we have no
+  // information for our users that is interesting.
+  if (f.escapes && f.partialApplySites.empty() && f.fullApplySites.empty())
+    return None;
+  return f;
+}