Sema: Ban forward references of captured values from 'defer' body

While 'defer' is implemented as a local function, it doesn't
behave as one. In particular, since SILGen runs it after
destroying all local bindings that appear after the 'defer'
definition, the body of a 'defer' cannot forward reference
captured bindings the way that local functions can.

Note that I had to remove a SILGen test case for an older,
related issue. The new diagnostic in Sema catches these cases
earlier.

Fixes rdar://problem/75088379.

Author: slavapestov

include/swift/AST/AnyFunctionRef.h

@@ -179,12 +179,6 @@ class AnyFunctionRef {
     return TheFunction.dyn_cast<AbstractClosureExpr*>();
   }
 
-  bool isDeferBody() const {
-    if (auto *fd = dyn_cast_or_null<FuncDecl>(getAbstractFunctionDecl()))
-      return fd->isDeferBody();
-    return false;
-  }
-
   /// Return true if this closure is passed as an argument to a function and is
   /// known not to escape from that function.  In this case, captures can be
   /// more efficient.

include/swift/AST/DiagnosticsSIL.def

@@ -103,8 +103,6 @@ ERROR(objc_selector_malformed,none,"the type ObjectiveC.Selector is malformed",
 // Capture before declaration diagnostics.
 ERROR(capture_before_declaration,none,
       "closure captures %0 before it is declared", (Identifier))
-ERROR(capture_before_declaration_defer,none,
-      "'defer' block captures %0 before it is declared", (Identifier))
 NOTE(captured_value_declared_here,none,
      "captured value declared here", ())
 

lib/SILGen/SILGenFunction.cpp

@@ -245,22 +245,16 @@ void SILGenFunction::emitCaptures(SILLocation loc,
       auto &Diags = getASTContext().Diags;
 
       SourceLoc loc;
-      bool isDeferBody;
       if (closure.kind == SILDeclRef::Kind::DefaultArgGenerator) {
         auto *param = getParameterAt(closure.getDecl(),
                                      closure.defaultArgIndex);
         loc = param->getLoc();
-        isDeferBody = false;
       } else {
         auto f = *closure.getAnyFunctionRef();
         loc = f.getLoc();
-        isDeferBody = f.isDeferBody();
       }
 
-      Diags.diagnose(loc,
-                     isDeferBody
-                     ? diag::capture_before_declaration_defer
-                     : diag::capture_before_declaration,
+      Diags.diagnose(loc, diag::capture_before_declaration,
                      vd->getBaseIdentifier());
       Diags.diagnose(vd->getLoc(), diag::captured_value_declared_here);
       Diags.diagnose(capture.getLoc(), diag::value_captured_here);

lib/Sema/PreCheckExpr.cpp

@@ -328,6 +328,34 @@ static bool isMemberChainTail(Expr *expr, Expr *parent) {
   return parent == nullptr || !isMemberChainMember(parent);
 }
 
+static bool isValidForwardReference(ValueDecl *D, DeclContext *DC,
+                                    ValueDecl **localDeclAfterUse) {
+  *localDeclAfterUse = nullptr;
+
+  // References to variables injected by lldb are always valid.
+  if (isa<VarDecl>(D) && cast<VarDecl>(D)->isDebuggerVar())
+    return true;
+
+  // If we find something in the current context, it must be a forward
+  // reference, because otherwise if it was in scope, it would have
+  // been returned by the call to ASTScope::lookupLocalDecls() above.
+  if (D->getDeclContext()->isLocalContext()) {
+    do {
+      if (D->getDeclContext() == DC) {
+        *localDeclAfterUse = D;
+        return false;
+      }
+
+      // If we're inside of a 'defer' context, walk up to the parent
+      // and check again. We don't want 'defer' bodies to forward
+      // reference bindings in the immediate outer scope.
+    } while (isa<FuncDecl>(DC) &&
+             cast<FuncDecl>(DC)->isDeferBody() &&
+             (DC = DC->getParent()));
+  }
+  return true;
+}
+
 /// Bind an UnresolvedDeclRefExpr by performing name lookup and
 /// returning the resultant expression. Context is the DeclContext used
 /// for the lookup.
@@ -398,24 +426,12 @@ Expr *TypeChecker::resolveDeclRefExpr(UnresolvedDeclRefExpr *UDRE,
     Lookup = TypeChecker::lookupUnqualified(DC, LookupName, Loc, lookupOptions);
 
     ValueDecl *localDeclAfterUse = nullptr;
-    auto isValid = [&](ValueDecl *D) {
-      // References to variables injected by lldb are always valid.
-      if (isa<VarDecl>(D) && cast<VarDecl>(D)->isDebuggerVar())
-        return true;
-
-      // If we find something in the current context, it must be a forward
-      // reference, because otherwise if it was in scope, it would have
-      // been returned by the call to ASTScope::lookupLocalDecls() above.
-      if (D->getDeclContext()->isLocalContext() &&
-          D->getDeclContext() == DC) {
-        localDeclAfterUse = D;
-        return false;
-      }
-      return true;
-    };
     AllDeclRefs =
         findNonMembers(Lookup.innerResults(), UDRE->getRefKind(),
-                       /*breakOnMember=*/true, ResultValues, isValid);
+                       /*breakOnMember=*/true, ResultValues,
+                       [&](ValueDecl *D) {
+                         return isValidForwardReference(D, DC, &localDeclAfterUse);
+                       });
 
     // If local declaration after use is found, check outer results for
     // better matching candidates.
@@ -435,7 +451,10 @@ Expr *TypeChecker::resolveDeclRefExpr(UnresolvedDeclRefExpr *UDRE,
         localDeclAfterUse = nullptr;
         AllDeclRefs =
             findNonMembers(Lookup.innerResults(), UDRE->getRefKind(),
-                           /*breakOnMember=*/true, ResultValues, isValid);
+                           /*breakOnMember=*/true, ResultValues,
+                           [&](ValueDecl *D) {
+                             return isValidForwardReference(D, DC, &localDeclAfterUse);
+                           });
       }
     }
   }

test/SILGen/capture_order.swift

@@ -119,25 +119,6 @@ func captureInClosure() {
 
 /// Regression tests
 
-func sr3210_crash() {
-  defer { // expected-error {{'defer' block captures 'b' before it is declared}}
-    print("\(b)") // expected-note {{captured here}}
-  }
-
-  return
-
-  let b = 2 // expected-note {{captured value declared here}}
-  // expected-warning@-1 {{code after 'return' will never be executed}}
-}
-
-func sr3210() {
-  defer {
-    print("\(b)")
-  }
-
-  let b = 2
-}
-
 class SR4812 {
   public func foo() {
     let bar = { [weak self] in

test/stmt/defer.swift

@@ -129,3 +129,35 @@ class SomeDerivedClass: SomeTestClass {
     }
   }
 }
+
+// rdar://75088379 -- 'defer' should not be able to forward-reference captures
+func badForwardReference() {
+  defer {
+    _ = x2 // expected-error {{use of local variable 'x2' before its declaration}}
+
+    let x1 = 0
+    let y1 = 0
+
+    defer {
+      _ = x1
+      _ = x2 // expected-error {{use of local variable 'x2' before its declaration}}
+      _ = y1
+      _ = y2 // expected-error {{use of local variable 'y2' before its declaration}}
+    }
+
+    let y2 = 0 // expected-note {{'y2' declared here}}
+  }
+
+  let x2 = 0 // expected-note 2{{'x2' declared here}}
+
+  func localFunc() {
+    defer {
+      _ = z1 // expected-error {{use of local variable 'z1' before its declaration}}
+      _ = z2
+    }
+
+    let z1 = 0 // expected-note {{'z1' declared here}}
+  }
+
+  let z2 = 0
+}