[Sema] Catch use-before-declarations in nested closures

Previously we would allow these in Sema and diagnose them in SILGen,
but allowing them in Sema is unsound because it means the constraint
system ends up kicking interface type requests for declarations that
should be type-checked as part of the closure itself. Adjust the
name lookup logic to look through parent closures when detecting
invalid forward references.

For now we don't fallback to an outer result on encountering a
use-before-declaration to preserve the current behavior. I'm planning
on changing that in the next commit though.

Author: hamishknight

include/swift/AST/Decl.h

@@ -6725,13 +6725,18 @@ class VarDecl : public AbstractStorageDecl {
     Bits.VarDecl.IsDebuggerVar = IsDebuggerVar;
   }
 
-  /// Visit all auxiliary declarations to this VarDecl.
+  /// Visit all auxiliary variables for this VarDecl.
   ///
-  /// An auxiliary declaration is a declaration synthesized by the compiler to support
-  /// this VarDecl, such as synthesized property wrapper variables.
+  /// An auxiliary variable is one that is synthesized by the compiler to
+  /// support this VarDecl, such as synthesized property wrapper variables.
   ///
-  /// \note this function only visits auxiliary decls that are not part of the AST.
-  void visitAuxiliaryDecls(llvm::function_ref<void(VarDecl *)>) const;
+  /// \param forNameLookup If \c true, will only visit auxiliary variables that
+  /// may appear in name lookup results.
+  ///
+  /// \note this function only visits auxiliary variables that are not part of
+  /// the AST.
+  void visitAuxiliaryVars(bool forNameLookup,
+                          llvm::function_ref<void(VarDecl *)>) const;
 
   /// Is this the synthesized storage for a 'lazy' property?
   bool isLazyStorageProperty() const {

include/swift/AST/Expr.h

@@ -506,6 +506,10 @@ class alignas(8) Expr : public ASTAllocated<Expr> {
   /// walk into the body of it (unless it is single-expression).
   void forEachChildExpr(llvm::function_ref<Expr *(Expr *)> callback);
 
+  /// Apply the specified function to all variables referenced in any
+  /// child UnresolvedPatternExprs.
+  void forEachUnresolvedVariable(llvm::function_ref<void(VarDecl *)> f) const;
+
   /// Determine whether this expression refers to a type by name.
   ///
   /// This distinguishes static references to types, like Int, from metatype

include/swift/AST/Pattern.h

@@ -192,7 +192,7 @@ class alignas(8) Pattern : public ASTAllocated<Pattern> {
   /// Collect the set of variables referenced in the given pattern.
   void collectVariables(SmallVectorImpl<VarDecl *> &variables) const;
 
-  /// apply the specified function to all variables referenced in this
+  /// Apply the specified function to all variables referenced in this
   /// pattern.
   void forEachVariable(llvm::function_ref<void(VarDecl *)> f) const;
 

include/swift/Sema/ConstraintSystem.h

@@ -2131,6 +2131,7 @@ class SyntacticElementTargetRewriter {
 
   virtual void addLocalDeclToTypeCheck(Decl *D) = 0;
 
+  [[nodiscard]]
   virtual std::optional<SyntacticElementTarget>
   rewriteTarget(SyntacticElementTarget target) = 0;
 

lib/AST/Decl.cpp

@@ -507,7 +507,7 @@ void Decl::visitAuxiliaryDecls(
     }
   }
 
-  // FIXME: fold VarDecl::visitAuxiliaryDecls into this.
+  // FIXME: fold VarDecl::visitAuxiliaryVars into this.
 }
 
 void Decl::forEachAttachedMacro(MacroRole role,
@@ -8767,14 +8767,22 @@ bool VarDecl::hasStorageOrWrapsStorage() const {
   return false;
 }
 
-void VarDecl::visitAuxiliaryDecls(llvm::function_ref<void(VarDecl *)> visit) const {
+void VarDecl::visitAuxiliaryVars(
+    bool forNameLookup, llvm::function_ref<void(VarDecl *)> visit) const {
   if (getDeclContext()->isTypeContext() ||
       (isImplicit() && !isa<ParamDecl>(this)))
     return;
 
-  if (getAttrs().hasAttribute<LazyAttr>()) {
-    if (auto *backingVar = getLazyStorageProperty())
-      visit(backingVar);
+  // The backing storage for a lazy property is not accessible to name lookup.
+  // This is not only a matter of hiding implementation details, but also
+  // correctness since triggering LazyStoragePropertyRequest currently eagerly
+  // requests the interface type for the var, which could result in incorrectly
+  // type-checking a lazy local independently of its surrounding closure.
+  if (!forNameLookup) {
+    if (getAttrs().hasAttribute<LazyAttr>()) {
+      if (auto *backingVar = getLazyStorageProperty())
+        visit(backingVar);
+    }
   }
 
   if (getAttrs().hasAttribute<CustomAttr>() || hasImplicitPropertyWrapper()) {

lib/AST/Pattern.cpp

@@ -192,8 +192,7 @@ namespace {
     const std::function<void(VarDecl*)> &fn;
   public:
 
-    WalkToVarDecls(const std::function<void(VarDecl*)> &fn)
-    : fn(fn) {}
+    WalkToVarDecls(const std::function<void(VarDecl*)> &fn) : fn(fn) {}
 
     /// Walk everything that's available; there shouldn't be macro expansions
     /// that matter anyway.
@@ -239,6 +238,9 @@ namespace {
   };
 } // end anonymous namespace
 
+void Expr::forEachUnresolvedVariable(llvm::function_ref<void(VarDecl *)> f) const {
+  const_cast<Expr *>(this)->walk(WalkToVarDecls(f));
+}
 
 /// apply the specified function to all variables referenced in this
 /// pattern.

lib/AST/UnqualifiedLookup.cpp

@@ -708,12 +708,13 @@ bool ASTScopeDeclConsumerForUnqualifiedLookup::consume(
       bool foundMatch = false;
       if (auto *varDecl = dyn_cast<VarDecl>(value)) {
         // Check if the name matches any auxiliary decls not in the AST
-        varDecl->visitAuxiliaryDecls([&](VarDecl *auxiliaryVar) {
-          if (auxiliaryVar->ValueDecl::getName().matchesRef(fullName)) {
-            value = auxiliaryVar;
-            foundMatch = true;
-          }
-        });
+        varDecl->visitAuxiliaryVars(
+            /*forNameLookup*/ true, [&](VarDecl *auxiliaryVar) {
+              if (auxiliaryVar->ValueDecl::getName().matchesRef(fullName)) {
+                value = auxiliaryVar;
+                foundMatch = true;
+              }
+            });
       }
 
       if (!foundMatch)
@@ -918,13 +919,14 @@ class ASTScopeDeclConsumerForLocalLookup
       bool foundMatch = false;
       if (auto *varDecl = dyn_cast<VarDecl>(value)) {
         // Check if the name matches any auxiliary decls not in the AST
-        varDecl->visitAuxiliaryDecls([&](VarDecl *auxiliaryVar) {
-          if (name.isSimpleName(auxiliaryVar->getName())
-                && hasCorrectABIRole(auxiliaryVar)) {
-            results.push_back(auxiliaryVar);
-            foundMatch = true;
-          }
-        });
+        varDecl->visitAuxiliaryVars(
+            /*forNameLookup*/ true, [&](VarDecl *auxiliaryVar) {
+              if (name.isSimpleName(auxiliaryVar->getName()) &&
+                  hasCorrectABIRole(auxiliaryVar)) {
+                results.push_back(auxiliaryVar);
+                foundMatch = true;
+              }
+            });
       }
 
       if (!foundMatch && value->getName().matchesRef(name)

lib/SILGen/SILGenDecl.cpp

@@ -1716,7 +1716,7 @@ void SILGenFunction::visitVarDecl(VarDecl *D) {
   }
 
   // Emit lazy and property wrapper backing storage.
-  D->visitAuxiliaryDecls([&](VarDecl *var) {
+  D->visitAuxiliaryVars(/*forNameLookup*/ false, [&](VarDecl *var) {
     if (auto *patternBinding = var->getParentPatternBinding())
       visitPatternBindingDecl(patternBinding);
 

lib/SILGen/SILGenProlog.cpp

@@ -1064,7 +1064,7 @@ class ArgumentInitHelper {
   void emitParam(ParamDecl *PD) {
     // Register any auxiliary declarations for the parameter to be
     // visited later.
-    PD->visitAuxiliaryDecls([&](VarDecl *localVar) {
+    PD->visitAuxiliaryVars(/*forNameLookup*/ false, [&](VarDecl *localVar) {
       SGF.LocalAuxiliaryDecls.push_back(localVar);
     });
 

lib/Sema/CSApply.cpp

@@ -8949,7 +8949,8 @@ namespace {
     PreWalkResult<Expr *> walkToExprPre(Expr *expr) override {
       // For closures, update the parameter types and check the body.
       if (auto closure = dyn_cast<ClosureExpr>(expr)) {
-        rewriteFunction(closure);
+        if (rewriteFunction(closure))
+          return Action::Stop();
 
         if (AnyFunctionRef(closure).hasExternalPropertyWrapperParameters()) {
           auto *thunkTy = Rewriter.cs.getType(closure)->castTo<FunctionType>();
@@ -8962,19 +8963,22 @@ namespace {
       }
 
       if (auto *SVE = dyn_cast<SingleValueStmtExpr>(expr)) {
-        rewriteSingleValueStmtExpr(SVE);
+        if (rewriteSingleValueStmtExpr(SVE))
+          return Action::Stop();
         return Action::SkipNode(SVE);
       }
 
       if (auto tap = dyn_cast_or_null<TapExpr>(expr)) {
-        rewriteTapExpr(tap);
+        if (rewriteTapExpr(tap))
+          return Action::Stop();
         return Action::SkipNode(tap);
       }
 
       if (auto captureList = dyn_cast<CaptureListExpr>(expr)) {
         // Rewrite captures.
         for (const auto &capture : captureList->getCaptureList()) {
-          (void)rewriteTarget(SyntacticElementTarget(capture.PBD));
+          if (!rewriteTarget(SyntacticElementTarget(capture.PBD)))
+            return Action::Stop();
         }
       }
 
@@ -9000,31 +9004,36 @@ namespace {
       return Action::SkipNode();
     }
 
-    NullablePtr<Pattern>
-    rewritePattern(Pattern *pattern, DeclContext *DC);
+    [[nodiscard]]
+    NullablePtr<Pattern> rewritePattern(Pattern *pattern, DeclContext *DC);
 
     /// Rewrite the target, producing a new target.
+    [[nodiscard]]
     std::optional<SyntacticElementTarget>
     rewriteTarget(SyntacticElementTarget target) override;
 
     /// Rewrite the function for the given solution.
     ///
     /// \returns true if an error occurred.
+    [[nodiscard]]
     bool rewriteFunction(AnyFunctionRef fn) {
       return Rewriter.cs.applySolution(fn, *this);
     }
 
+    [[nodiscard]]
     bool rewriteSingleValueStmtExpr(SingleValueStmtExpr *SVE) {
       return Rewriter.cs.applySolutionToSingleValueStmt(SVE, *this);
     }
 
-    void rewriteTapExpr(TapExpr *tap) {
+    [[nodiscard]]
+    bool rewriteTapExpr(TapExpr *tap) {
       // First, let's visit the tap expression itself
       // and set all of the inferred types.
-      Rewriter.visitTapExpr(tap);
+      if (!Rewriter.visitTapExpr(tap))
+        return true;
 
       // Now, let's apply solution to the body
-      (void)Rewriter.cs.applySolutionToBody(tap, *this);
+      return Rewriter.cs.applySolutionToBody(tap, *this);
     }
   };
 } // end anonymous namespace