Fix keypath-as-function crasher

include/swift/SIL/TypeLowering.h

@@ -972,6 +972,24 @@ class TypeConverter {
   CaptureInfo getLoweredLocalCaptures(SILDeclRef fn);
   bool hasLoweredLocalCaptures(SILDeclRef fn);
 
+#ifndef NDEBUG
+  /// If \c false, \c childDC is in a context it cannot capture variables from,
+  /// so it is expected that Sema may not have computed its \c CaptureInfo.
+  ///
+  /// This call exists for use in assertions; do not use it to skip capture
+  /// processing.
+  static bool canCaptureFromParent(DeclContext *childDC) {
+    // This call was added because Sema leaves the captures of functions that
+    // cannot capture anything uncomputed.
+    // TODO: Make Sema set them to CaptureInfo::empty() instead.
+
+    if (childDC)
+      if (auto decl = childDC->getAsDecl())
+         return decl->getDeclContext()->isLocalContext();
+    return true;
+  }
+#endif
+
   enum class ABIDifference : uint8_t {
     // No ABI differences, function can be trivially bitcast to result type.
     Trivial,

lib/AST/CaptureInfo.cpp

@@ -78,7 +78,7 @@ getLocalCaptures(SmallVectorImpl<CapturedValue> &Result) const {
   }
 }
 
-void CaptureInfo::dump() const {
+LLVM_ATTRIBUTE_USED void CaptureInfo::dump() const {
   print(llvm::errs());
   llvm::errs() << '\n';
 }

lib/SIL/TypeLowering.cpp

@@ -24,9 +24,11 @@
 #include "swift/AST/NameLookup.h"
 #include "swift/AST/ParameterList.h"
 #include "swift/AST/Pattern.h"
+#include "swift/AST/PrettyStackTrace.h"
 #include "swift/AST/PropertyWrappers.h"
 #include "swift/AST/Types.h"
 #include "swift/ClangImporter/ClangModule.h"
+#include "swift/SIL/PrettyStackTrace.h"
 #include "swift/SIL/SILArgument.h"
 #include "swift/SIL/SILBuilder.h"
 #include "swift/SIL/SILModule.h"
@@ -2111,6 +2113,9 @@ TypeConverter::hasLoweredLocalCaptures(SILDeclRef fn) {
 
 CaptureInfo
 TypeConverter::getLoweredLocalCaptures(SILDeclRef fn) {
+  PrettyStackTraceSILLocation stack("getting lowered local captures",
+                                    fn.getAsRegularLocation(), Context);
+
   fn.isForeign = 0;
   fn.isCurried = 0;
   fn.isDirectReference = 0;
@@ -2132,11 +2137,14 @@ TypeConverter::getLoweredLocalCaptures(SILDeclRef fn) {
   DynamicSelfType *capturesDynamicSelf = nullptr;
   OpaqueValueExpr *capturesOpaqueValue = nullptr;
 
-  std::function<void (CaptureInfo captureInfo)> collectCaptures;
+  std::function<void (CaptureInfo captureInfo, DeclContext *dc)> collectCaptures;
   std::function<void (AnyFunctionRef)> collectFunctionCaptures;
   std::function<void (SILDeclRef)> collectConstantCaptures;
 
-  collectCaptures = [&](CaptureInfo captureInfo) {
+  collectCaptures = [&](CaptureInfo captureInfo, DeclContext *dc) {
+    assert(captureInfo.hasBeenComputed() ||
+           !TypeConverter::canCaptureFromParent(dc));
+
     if (captureInfo.hasGenericParamCaptures())
       capturesGenericParams = true;
     if (captureInfo.hasDynamicSelfCapture())
@@ -2263,7 +2271,9 @@ TypeConverter::getLoweredLocalCaptures(SILDeclRef fn) {
     if (!visitedFunctions.insert(curFn).second)
       return;
 
-    collectCaptures(curFn.getCaptureInfo());
+    PrettyStackTraceAnyFunctionRef("lowering local captures", curFn);
+    auto dc = curFn.getAsDeclContext();
+    collectCaptures(curFn.getCaptureInfo(), dc);
 
     // A function's captures also include its default arguments, because
     // when we reference a function we don't track which default arguments
@@ -2274,17 +2284,22 @@ TypeConverter::getLoweredLocalCaptures(SILDeclRef fn) {
     if (auto *AFD = curFn.getAbstractFunctionDecl()) {
       for (auto *P : *AFD->getParameters()) {
         if (P->getDefaultValue())
-          collectCaptures(P->getDefaultArgumentCaptureInfo());
+          collectCaptures(P->getDefaultArgumentCaptureInfo(), dc);
       }
     }
   };
 
   collectConstantCaptures = [&](SILDeclRef curFn) {
     if (curFn.isDefaultArgGenerator()) {
+      PrettyStackTraceSILLocation stack("lowering local captures",
+                                        fn.getAsRegularLocation(), Context);
+
       if (auto *afd = dyn_cast<AbstractFunctionDecl>(curFn.getDecl())) {
         auto *param = getParameterAt(afd, curFn.defaultArgIndex);
-        if (param->getDefaultValue())
-          collectCaptures(param->getDefaultArgumentCaptureInfo());
+        if (param->getDefaultValue()) {
+          auto dc = afd->getInnermostDeclContext();
+          collectCaptures(param->getDefaultArgumentCaptureInfo(), dc);
+        }
         return;
       }
 

lib/SILGen/SILGenFunction.cpp

@@ -666,6 +666,31 @@ void SILGenFunction::emitArtificialTopLevel(ClassDecl *mainClass) {
   }
 }
 
+#ifndef NDEBUG
+/// If \c false, \c function is either a declaration that inherently cannot
+/// capture variables, or it is in a context it cannot capture variables from.
+/// In either case, it is expected that Sema may not have computed its
+/// \c CaptureInfo.
+///
+/// This call exists for use in assertions; do not use it to skip capture
+/// processing.
+static bool canCaptureFromParent(SILDeclRef function) {
+  switch (function.kind) {
+  case SILDeclRef::Kind::StoredPropertyInitializer:
+  case SILDeclRef::Kind::PropertyWrapperBackingInitializer:
+    return false;
+
+  default:
+    if (function.hasDecl()) {
+      if (auto dc = dyn_cast<DeclContext>(function.getDecl())) {
+        return TypeConverter::canCaptureFromParent(dc);
+      }
+    }
+    return false;
+  }
+}
+#endif
+
 void SILGenFunction::emitGeneratorFunction(SILDeclRef function, Expr *value,
                                            bool EmitProfilerIncrement) {
   auto *dc = function.getDecl()->getInnermostDeclContext();
@@ -706,6 +731,11 @@ void SILGenFunction::emitGeneratorFunction(SILDeclRef function, Expr *value,
   CaptureInfo captureInfo;
   if (function.getAnyFunctionRef())
     captureInfo = SGM.M.Types.getLoweredLocalCaptures(function);
+  else {
+    assert(!canCaptureFromParent(function));
+    captureInfo = CaptureInfo::empty();
+  }
+
   auto interfaceType = value->getType()->mapTypeOutOfContext();
   emitProlog(captureInfo, params, /*selfParam=*/nullptr,
              dc, interfaceType, /*throws=*/false, SourceLoc());

lib/SILGen/SILGenProlog.cpp

@@ -445,6 +445,9 @@ void SILGenFunction::emitProlog(CaptureInfo captureInfo,
 
   // Emit the capture argument variables. These are placed last because they
   // become the first curry level of the SIL function.
+  assert((captureInfo.hasBeenComputed() ||
+          !TypeConverter::canCaptureFromParent(DC)) &&
+         "can't emit prolog of function with uncomputed captures");
   for (auto capture : captureInfo.getCaptures()) {
     if (capture.isDynamicSelfMetadata()) {
       auto selfMetatype = MetatypeType::get(

lib/Sema/CSApply.cpp

@@ -4506,6 +4506,10 @@ namespace {
       outerClosure->setType(outerClosureTy);
       cs.cacheType(outerClosure);
 
+      // The inner closure at least will definitely have a capture.
+      cs.TC.ClosuresWithUncomputedCaptures.push_back(outerClosure);
+      cs.TC.ClosuresWithUncomputedCaptures.push_back(closure);
+
       // let outerApply = "\( outerClosure )( \(E) )"
       auto outerApply = CallExpr::createImplicit(ctx, outerClosure, {E}, {});
       outerApply->setType(closureTy);

lib/Sema/TypeCheckREPL.cpp

@@ -271,6 +271,7 @@ void REPLChecker::generatePrintOfExpression(StringRef NameStr, Expr *E) {
   BraceStmt *Body = builder.createBodyStmt(Loc, EndLoc);
   CE->setBody(Body, false);
   TC.typeCheckClosureBody(CE);
+  TC.ClosuresWithUncomputedCaptures.push_back(CE);
 
   auto *TheCall = CallExpr::createImplicit(Context, CE, { E }, { });
   TheCall->getArg()->setType(AnyFunctionType::composeInput(Context, args, false));

validation-test/compiler_crashers_2_fixed/0209-rdar56055600.swift

@@ -0,0 +1,2 @@
+// RUN: %target-swift-frontend %s -emit-silgen -o /dev/null
+let _: ([Int]) -> Int = \[Int].count