Construct AutoClosureExpr With Deeper Parent Contexts

AutoClosureExprs created by the constraint system used to be constructed
with the decl context of the constraint system itself. This meant that
autoclosures in expressions nested in closures would initially be
parented onto any enclosing functions rather than the deepest closure
context. When we ran capture analysis and lookup from inside of the body
of these nascent values, we would fail to find declarations brought into
scope by those parent closures. This is especially relevant when
pre-typechecked code is involved since captures for those declarations
will be forced before their bodies have been recontextualized. See
issue #34230 for why we need to force things so early.

The attached test case demonstrates both bugs: The former a bogus lookup
through the parent context that would incorrectly reject this otherwise
well-formed code. The latter is a crash in SILGen when the capture
computation would fail to note $0.

Use the decl context of the solution application target, which is always
going to be the deepest user-written closure expression available to us,
and therefore the deepest scope that can introduce capturable variables.

rdar://79248469

Author: CodaFi

include/swift/Sema/ConstraintSystem.h

@@ -4433,7 +4433,23 @@ class ConstraintSystem {
 
   /// Build implicit autoclosure expression wrapping a given expression.
   /// Given expression represents computed result of the closure.
+  ///
+  /// The \p ClosureDC must be the deepest possible context that
+  /// contains this autoclosure expression. For example,
+  ///
+  /// func foo() {
+  ///   _ = { $0 || $1 || $2 }
+  /// }
+  ///
+  /// Even though the decl context of $1 (after solution application) is
+  /// `||`'s autoclosure parameter, we cannot know this until solution
+  /// application has finished because autoclosure expressions are expanded in
+  /// depth-first order then \c ContextualizeClosures comes around to clean up.
+  /// All that is required is that the explicit closure be the context since it
+  /// is the innermost context that can introduce potential new capturable
+  /// declarations.
   Expr *buildAutoClosureExpr(Expr *expr, FunctionType *closureType,
+                             DeclContext *ClosureDC,
                              bool isDefaultWrappedValue = false,
                              bool isAsyncLetWrapper = false);
 

lib/Sema/CSApply.cpp

@@ -1092,8 +1092,8 @@ namespace {
       FunctionType::ExtInfo closureInfo;
       auto *autoClosureType =
           FunctionType::get(outerParamTypes, fnType->getResult(), closureInfo);
-      auto *autoClosure = new (context) AutoClosureExpr(fnCall, autoClosureType,
-                                                        discriminator, cs.DC);
+      auto *autoClosure = new (context)
+          AutoClosureExpr(fnCall, autoClosureType, discriminator, dc);
       autoClosure->setParameterList(outerParams);
       autoClosure->setThunkKind(AutoClosureExpr::Kind::SingleCurryThunk);
       cs.cacheType(autoClosure);
@@ -1126,9 +1126,8 @@ namespace {
 
       auto resultTy = selfFnTy->getResult();
       auto discriminator = AutoClosureExpr::InvalidDiscriminator;
-      auto closure =
-          new (context) AutoClosureExpr(/*set body later*/nullptr, resultTy,
-                                        discriminator, cs.DC);
+      auto closure = new (context) AutoClosureExpr(/*set body later*/ nullptr,
+                                                   resultTy, discriminator, dc);
       closure->setParameterList(params);
       closure->setType(selfFnTy);
       closure->setThunkKind(AutoClosureExpr::Kind::SingleCurryThunk);
@@ -1693,8 +1692,7 @@ namespace {
                                        memberLocator);
 
         auto outerClosure =
-            new (context) AutoClosureExpr(closure, selfFnTy,
-                                          discriminator, cs.DC);
+            new (context) AutoClosureExpr(closure, selfFnTy, discriminator, dc);
         outerClosure->setThunkKind(AutoClosureExpr::Kind::DoubleCurryThunk);
 
         outerClosure->setParameterList(outerParams);
@@ -6056,14 +6054,15 @@ Expr *ExprRewriter::coerceCallArguments(
         bool isDefaultWrappedValue =
             target->propertyWrapperHasInitialWrappedValue();
         auto *placeholder = injectWrappedValuePlaceholder(
-            cs.buildAutoClosureExpr(arg, closureType, isDefaultWrappedValue),
+            cs.buildAutoClosureExpr(arg, closureType, dc,
+                                    isDefaultWrappedValue),
             /*isAutoClosure=*/true);
         arg = CallExpr::createImplicit(ctx, placeholder, {}, {});
         arg->setType(closureType->getResult());
         cs.cacheType(arg);
       }
 
-      convertedArg = cs.buildAutoClosureExpr(arg, closureType);
+      convertedArg = cs.buildAutoClosureExpr(arg, closureType, dc);
     } else {
       convertedArg = coerceToType(
           arg, paramType,
@@ -8324,7 +8323,7 @@ static Expr *wrapAsyncLetInitializer(
   auto closureType = FunctionType::get({ }, initializerType, extInfo);
   ASTContext &ctx = dc->getASTContext();
   Expr *autoclosureExpr = cs.buildAutoClosureExpr(
-      initializer, closureType, /*isDefaultWrappedValue=*/false,
+      initializer, closureType, dc, /*isDefaultWrappedValue=*/false,
       /*isAsyncLetWrapper=*/true);
 
   // Call the autoclosure so that the AST types line up. SILGen will ignore the
@@ -8792,7 +8791,8 @@ ExprWalker::rewriteTarget(SolutionApplicationTarget target) {
     // conversion.
     if (FunctionType *autoclosureParamType =
             target.getAsAutoclosureParamType()) {
-      resultExpr = cs.buildAutoClosureExpr(resultExpr, autoclosureParamType);
+      resultExpr = cs.buildAutoClosureExpr(resultExpr, autoclosureParamType,
+                                           target.getDeclContext());
     }
 
     solution.setExprTypes(resultExpr);

lib/Sema/ConstraintSystem.cpp

@@ -5098,6 +5098,7 @@ ConstraintSystem::isConversionEphemeral(ConversionRestrictionKind conversion,
 
 Expr *ConstraintSystem::buildAutoClosureExpr(Expr *expr,
                                              FunctionType *closureType,
+                                             DeclContext *ClosureContext,
                                              bool isDefaultWrappedValue,
                                              bool isAsyncLetWrapper) {
   auto &Context = DC->getASTContext();
@@ -5116,8 +5117,9 @@ Expr *ConstraintSystem::buildAutoClosureExpr(Expr *expr,
     newClosureType = closureType->withExtInfo(info.withNoEscape(false))
                          ->castTo<FunctionType>();
 
-  auto *closure = new (Context) AutoClosureExpr(
-      expr, newClosureType, AutoClosureExpr::InvalidDiscriminator, DC);
+  auto *closure = new (Context)
+      AutoClosureExpr(expr, newClosureType,
+                      AutoClosureExpr::InvalidDiscriminator, ClosureContext);
 
   closure->setParameterList(ParameterList::createEmpty(Context));
 

test/expr/capture/local_lazy.swift

@@ -0,0 +1,29 @@
+// RUN: %target-swift-frontend -emit-silgen -verify %s > /dev/null
+// RUN: %target-swift-frontend -dump-ast %s | %FileCheck %s
+
+struct S {
+  func foo() -> Int {
+    // Make sure the decl context for the autoclosure passed to ?? is deep
+    // enough that it can 'see' the capture of $0 from the outer closure.
+    // CHECK-LABEL: (lazy_initializer_expr
+    // CHECK: (closure_expr
+    // CHECK: location={{.*}}local_lazy.swift:[[@LINE+3]]
+    // CHECK: (autoclosure_expr implicit
+    // CHECK: captures=($0<direct>
+    lazy var nest: (Int) -> Int = { Optional<Int>.none ?? $0 }
+    return nest(1)
+  }
+}
+
+extension S {
+  func bar() -> Int {
+    // CHECK-LABEL: (lazy_initializer_expr
+    // CHECK: (closure_expr
+    // CHECK: location={{.*}}local_lazy.swift:[[@LINE+3]]
+    // CHECK: (autoclosure_expr implicit
+    // CHECK: captures=($0<direct>
+    lazy var nest: (Int) -> Int = { Optional<Int>.none ?? $0 }
+    return nest(1)
+  }
+}
+