SILOptimizer: Re-implement NPCR diagnostics in SIL pass

This fixes a test involving transitive captures of local functions,
as well as an infinite recursion possible with the old code.

Fixes <rdar://problem/34496304>.

Author: slavapestov

include/swift/AST/DiagnosticsSIL.def

@@ -137,6 +137,12 @@ NOTE(value_captured_here,none,"captured here", ())
 NOTE(value_captured_transitively,none,
      "captured indirectly by this call", ())
 
+ERROR(err_noescape_param_call,none,
+      "passing a %select{|closure which captures a }1non-escaping function "
+      "parameter %0 to a call to a non-escaping function parameter can allow "
+      "re-entrant modification of a variable",
+      (DeclName, unsigned))
+
 // Definite initialization diagnostics.
 NOTE(variable_defined_here,none,
      "%select{variable|constant}0 defined here", (bool))

lib/SILOptimizer/Mandatory/DiagnoseInvalidEscapingCaptures.cpp

@@ -329,11 +329,72 @@ static void checkPartialApply(ASTContext &Context, DeclContext *DC,
   }
 }
 
+static void checkApply(ASTContext &Context, FullApplySite site) {
+  auto isNoEscapeParam = [&](SILValue value) -> const ParamDecl * {
+    // If the value is an escaping, do not enforce any restrictions.
+    if (!value->getType().getASTType()->isNoEscape())
+      return nullptr;
+
+    // If the value is not a function parameter, do not enforce any restrictions.
+    return getParamDeclFromOperand(value);
+  };
+
+  // If the callee is not a no-escape parameter, there is nothing to check.
+  auto callee = site.getCalleeOrigin();
+  if (!isNoEscapeParam(callee))
+    return;
+
+  // See if any of our arguments are noescape parameters, or closures capturing
+  // noescape parameters.
+  SmallVector<std::pair<SILValue, bool>, 4> args;
+  llvm::SmallDenseSet<SILValue, 4> visited;
+  auto arglistInsert = [&](SILValue arg, bool capture) {
+    if (visited.insert(arg).second)
+      args.emplace_back(arg, capture);
+  };
+
+  for (auto arg : site.getArguments())
+    arglistInsert(arg, /*capture=*/false);
+
+  while (!args.empty()) {
+    auto pair = args.pop_back_val();
+    auto arg = pair.first;
+    bool capture = pair.second;
+
+    if (auto *CI = dyn_cast<ConversionInst>(arg)) {
+      arglistInsert(CI->getConverted(), /*capture=*/false);
+      continue;
+    }
+
+    // If one of our call arguments is a noescape parameter, diagnose the
+    // violation.
+    if (auto *param = isNoEscapeParam(arg)) {
+      diagnose(Context, site.getLoc(), diag::err_noescape_param_call,
+               param->getName(), capture);
+      return;
+    }
+
+    // If one of our call arguments is a closure, recursively visit all of
+    // the closure's captures.
+    if (auto *PAI = dyn_cast<PartialApplyInst>(arg)) {
+      ApplySite site(PAI);
+      for (auto arg : site.getArguments())
+        arglistInsert(arg, /*capture=*/true);
+      continue;
+    }
+  }
+}
+
 static void checkForViolationsAtInstruction(ASTContext &Context,
                                             DeclContext *DC,
                                             SILInstruction *I) {
   if (auto *PAI = dyn_cast<PartialApplyInst>(I))
     checkPartialApply(Context, DC, PAI);
+
+  if (isa<ApplyInst>(I) || isa<TryApplyInst>(I)) {
+    FullApplySite site(I);
+    checkApply(Context, site);
+  }
 }
 
 static void checkEscapingCaptures(SILFunction *F) {

test/SILOptimizer/noescape_param_exclusivity.swift

@@ -0,0 +1,48 @@
+// RUN: %target-swift-frontend -emit-sil %s -verify
+
+func test0(a: (() -> ()) -> (), b: () -> ()) {
+  a(b) // expected-error {{passing a non-escaping function parameter 'b' to a call to a non-escaping function parameter can allow re-entrant modification of a variable}}
+}
+
+func test0(fn: (() -> ()) -> ()) {
+  fn { fn {} } // expected-error {{passing a closure which captures a non-escaping function parameter 'fn' to a call to a non-escaping function parameter can allow re-entrant modification of a variable}}
+}
+
+func test1(fn: (() -> ()) -> ()) {
+  func foo() {
+    fn { fn {} } // expected-error {{passing a closure which captures a non-escaping function parameter 'fn' to a call to a non-escaping function parameter can allow re-entrant modification of a variable}}
+  }
+}
+
+func test2(x: inout Int, fn: (() -> ()) -> ()) {
+  func foo(myfn: () -> ()) {
+    x += 1
+    myfn()
+  }
+
+  // Make sure we only complain about calls to noescape parameters.
+  foo { fn {} }
+}
+
+func test3(fn: (() -> ()) -> ()) {
+  { myfn in myfn { fn {} } }(fn) // expected-error {{passing a closure which captures a non-escaping function parameter 'fn' to a call to a non-escaping function parameter can allow re-entrant modification of a variable}}
+}
+
+func test4(fn: (() -> ()) -> ()) {
+  func foo() {
+    fn {}
+  }
+
+  fn(foo) // expected-error {{passing a closure which captures a non-escaping function parameter 'fn' to a call to a non-escaping function parameter can allow re-entrant modification of a variable}}
+}
+
+// rdar://problem/34496304
+func test5(outer: (() throws -> Int) throws -> Int) throws -> Int {
+  func descend(_ inner: (() throws -> Int) throws -> Int) throws -> Int {
+    return try inner { // expected-error {{passing a closure which captures a non-escaping function parameter 'inner' to a call to a non-escaping function parameter can allow re-entrant modification of a variable}}
+      try descend(inner)
+    }
+  }
+
+  return try descend(outer)
+}