Sema: Accept if #_hasSymbol() conditions in closure contexts.

Resolves rdar://100129165

Author: tshortli

lib/AST/ASTWalker.cpp

@@ -1394,8 +1394,17 @@ class Traversal : public ASTVisitor<Traversal, Expr*, Stmt*,
     for (auto &elt : C) {
       switch (elt.getKind()) {
       case StmtConditionElement::CK_Availability:
-      case StmtConditionElement::CK_HasSymbol:
         break;
+      case StmtConditionElement::CK_HasSymbol: {
+        auto E = elt.getHasSymbolInfo()->getSymbolExpr();
+        if (!E)
+          return true;
+        E = doIt(E);
+        if (!E)
+          return true;
+        elt.getHasSymbolInfo()->setSymbolExpr(E);
+        break;
+      }
       case StmtConditionElement::CK_Boolean: {
         auto E = elt.getBoolean();
         // Walk an expression condition normally.

lib/Sema/CSApply.cpp

@@ -8923,9 +8923,23 @@ ExprWalker::rewriteTarget(SolutionApplicationTarget target) {
     for (auto &condElement : *stmtCondition) {
       switch (condElement.getKind()) {
       case StmtConditionElement::CK_Availability:
-      case StmtConditionElement::CK_HasSymbol:
         continue;
 
+      case StmtConditionElement::CK_HasSymbol: {
+        ConstraintSystem &cs = solution.getConstraintSystem();
+        auto target = *cs.getSolutionApplicationTarget(&condElement);
+        auto resolvedTarget = rewriteTarget(target);
+        if (!resolvedTarget)
+          return None;
+
+        auto info = condElement.getHasSymbolInfo();
+        auto rewrittenExpr = resolvedTarget->getAsExpr();
+        info->setSymbolExpr(rewrittenExpr);
+        info->setReferencedDecl(
+            TypeChecker::getReferencedDeclForHasSymbolCondition(rewrittenExpr));
+        continue;
+      }
+
       case StmtConditionElement::CK_Boolean: {
         auto condExpr = condElement.getBoolean();
         auto finalCondExpr = condExpr->walk(*this);

lib/Sema/CSGen.cpp

@@ -4369,10 +4369,15 @@ bool ConstraintSystem::generateConstraints(StmtCondition condition,
       continue;
 
     case StmtConditionElement::CK_HasSymbol: {
-      ASTContext &ctx = getASTContext();
-      ctx.Diags.diagnose(condElement.getStartLoc(),
-                         diag::has_symbol_unsupported_in_closures);
-      return true;
+      Expr *symbolExpr = condElement.getHasSymbolInfo()->getSymbolExpr();
+      auto target = SolutionApplicationTarget(symbolExpr, dc, CTP_Unused,
+                                              Type(), /*isDiscarded=*/false);
+
+      if (generateConstraints(target, FreeTypeVariableBinding::Disallow))
+        return true;
+
+      setSolutionApplicationTarget(&condElement, target);
+      continue;
     }
 
     case StmtConditionElement::CK_Boolean: {

lib/Sema/MiscDiagnostics.cpp

@@ -4302,6 +4302,50 @@ checkImplicitPromotionsInCondition(const StmtConditionElement &cond,
   }
 }
 
+/// Perform MiscDiagnostics for the conditions belonging to a \c
+/// LabeledConditionalStmt.
+static void checkLabeledStmtConditions(ASTContext &ctx,
+                                       const LabeledConditionalStmt *stmt,
+                                       DeclContext *DC) {
+  for (auto elt : stmt->getCond()) {
+    // Check for implicit optional promotions in stmt-condition patterns.
+    checkImplicitPromotionsInCondition(elt, ctx);
+
+    switch (elt.getKind()) {
+    case StmtConditionElement::CK_Boolean:
+    case StmtConditionElement::CK_PatternBinding:
+    case StmtConditionElement::CK_Availability:
+      break;
+
+    case StmtConditionElement::CK_HasSymbol: {
+      auto info = elt.getHasSymbolInfo();
+      auto symbolExpr = info->getSymbolExpr();
+      if (!symbolExpr)
+        break;
+
+      if (!symbolExpr->getType())
+        break;
+
+      if (auto decl = info->getReferencedDecl().getDecl()) {
+        // `if #_hasSymbol(someStronglyLinkedSymbol)` is functionally a no-op
+        // and may indicate the developer has mis-identified the declaration
+        // they want to check (or forgot to import the module weakly).
+        if (!decl->isWeakImported(DC->getParentModule())) {
+          ctx.Diags.diagnose(symbolExpr->getLoc(),
+                             diag::has_symbol_decl_must_be_weak,
+                             decl->getDescriptiveKind(), decl->getName());
+        }
+      } else {
+        // Diagnose because we weren't able to interpret the expression as one
+        // that uniquely identifies a single declaration.
+        ctx.Diags.diagnose(symbolExpr->getLoc(), diag::has_symbol_invalid_expr);
+      }
+      break;
+    }
+    }
+  }
+}
+
 static void diagnoseUnintendedOptionalBehavior(const Expr *E,
                                                const DeclContext *DC) {
   if (!E || isa<ErrorExpr>(E) || !E->getType())
@@ -5290,11 +5334,9 @@ void swift::performStmtDiagnostics(const Stmt *S, DeclContext *DC) {
     checkSwitch(ctx, switchStmt, DC);
 
   checkStmtConditionTrailingClosure(ctx, S);
-  
-  // Check for implicit optional promotions in stmt-condition patterns.
+
   if (auto *lcs = dyn_cast<LabeledConditionalStmt>(S))
-    for (const auto &elt : lcs->getCond())
-      checkImplicitPromotionsInCondition(elt, ctx);
+    checkLabeledStmtConditions(ctx, lcs, DC);
 
   if (!ctx.LangOpts.DisableAvailabilityChecking)
     diagnoseStmtAvailability(S, const_cast<DeclContext*>(DC));

lib/Sema/TypeCheckStmt.cpp

@@ -448,8 +448,7 @@ static Expr *getDeclRefProvidingExpressionForHasSymbol(Expr *E) {
   return E;
 }
 
-static ConcreteDeclRef
-getReferencedDeclForHasSymbolCondition(ASTContext &Context, Expr *E) {
+ConcreteDeclRef TypeChecker::getReferencedDeclForHasSymbolCondition(Expr *E) {
   // Match DotSelfExprs (e.g. `SomeStruct.self`) when the type is static.
   if (auto DSE = dyn_cast<DotSelfExpr>(E)) {
     if (DSE->isStaticallyDerivedMetatype())
@@ -461,7 +460,6 @@ getReferencedDeclForHasSymbolCondition(ASTContext &Context, Expr *E) {
       return CDR;
   }
 
-  Context.Diags.diagnose(E->getLoc(), diag::has_symbol_invalid_expr);
   return ConcreteDeclRef();
 }
 
@@ -506,20 +504,10 @@ bool TypeChecker::typeCheckStmtConditionElement(StmtConditionElement &elt,
     auto exprTy = TypeChecker::typeCheckExpression(E, dc);
     Info->setSymbolExpr(E);
 
-    if (!exprTy)
-      return true;
-
-    auto CDR = getReferencedDeclForHasSymbolCondition(Context, E);
-    if (!CDR)
-      return true;
+    if (exprTy)
+      Info->setReferencedDecl(getReferencedDeclForHasSymbolCondition(E));
 
-    auto decl = CDR.getDecl();
-    if (!decl->isWeakImported(dc->getParentModule())) {
-      Context.Diags.diagnose(E->getLoc(), diag::has_symbol_decl_must_be_weak,
-                             decl->getDescriptiveKind(), decl->getName());
-    }
-    Info->setReferencedDecl(CDR);
-    return false;
+    return !exprTy;
   }
 
   if (auto E = elt.getBooleanOrNull()) {

lib/Sema/TypeChecker.h

@@ -441,6 +441,10 @@ Expr *substituteInputSugarTypeForResult(ApplyExpr *E);
 bool typeCheckStmtConditionElement(StmtConditionElement &elt, bool &isFalsable,
                                    DeclContext *dc);
 
+/// Returns the unique decl ref identified by the expr according to the
+/// requirements of the \c #_hasSymbol() condition type.
+ConcreteDeclRef getReferencedDeclForHasSymbolCondition(Expr *E);
+
 void typeCheckASTNode(ASTNode &node, DeclContext *DC,
                       bool LeaveBodyUnchecked = false);
 

test/Sema/has_symbol.swift

@@ -93,6 +93,7 @@ func testNotWeakDeclDiagnostics(_ s: LocalStruct) {
 }
 
 func testInvalidExpressionsDiagnostics() {
+  if #_hasSymbol(unknownDecl) {} // expected-error {{cannot find 'unknownDecl' in scope}}
   if #_hasSymbol(noArgFunc()) {} // expected-error {{#_hasSymbol condition must refer to a declaration}}
   if #_hasSymbol(global - 1) {} // expected-error {{#_hasSymbol condition must refer to a declaration}}
   if #_hasSymbol(S.staticFunc()) {} // expected-error {{#_hasSymbol condition must refer to a declaration}}
@@ -101,15 +102,36 @@ func testInvalidExpressionsDiagnostics() {
   if #_hasSymbol(1 as S) {} // expected-error {{cannot convert value of type 'Int' to type 'S' in coercion}}
 }
 
-func testMultiStatementClosure() {
-  let _: () -> Void = { // expected-error {{unable to infer closure type in the current context}}
-    if #_hasSymbol(global) {} // expected-error 2 {{#_hasSymbol is not supported in closures}}
-  }
-  
-  let _: () -> Void = { // expected-error {{unable to infer closure type in the current context}}
-    if #_hasSymbol(global) {} // expected-error 2 {{#_hasSymbol is not supported in closures}}
-    localFunc()
-  }
+func testGuard() {
+  guard #_hasSymbol(global) else { return }
+  guard #_hasSymbol(unknownDecl) else { return } // expected-error {{cannot find 'unknownDecl' in scope}}
+  guard #_hasSymbol(localFunc) else { return } // expected-warning {{global function 'localFunc()' is not a weakly linked declaration}}
+}
+
+func testWhile() {
+  while #_hasSymbol(global) { break }
+  while #_hasSymbol(unknownDecl) { break } // expected-error {{cannot find 'unknownDecl' in scope}}
+  while #_hasSymbol(localFunc) { break } // expected-warning {{global function 'localFunc()' is not a weakly linked declaration}}
+}
+
+func doIt(_ closure: () -> ()) {
+  closure()
+}
+
+func testClosure() {
+  doIt { if #_hasSymbol(global) {} }
+  doIt { if #_hasSymbol(noArgFunc) {} }
+  doIt { if #_hasSymbol(ambiguousFunc as () -> Int) {} }
+  doIt { if #_hasSymbol(S.self) {} }
+  doIt { if #_hasSymbol(ambiguousFunc) {} } // expected-error {{ambiguous use of 'ambiguousFunc()'}}
+  doIt { if #_hasSymbol(localFunc) {} } // expected-warning {{global function 'localFunc()' is not a weakly linked declaration}}
+  doIt { if #_hasSymbol(unknownDecl) {} } // expected-error {{cannot find 'unknownDecl' in scope}}
+  doIt { if #_hasSymbol(noArgFunc()) {} } // expected-error {{#_hasSymbol condition must refer to a declaration}}
+  doIt { if #_hasSymbol(global - 1) {} } // expected-error {{#_hasSymbol condition must refer to a declaration}}
+  doIt { if #_hasSymbol(S.staticFunc()) {} } // expected-error {{#_hasSymbol condition must refer to a declaration}}
+  doIt { if #_hasSymbol(C.classFunc()) {} } // expected-error {{#_hasSymbol condition must refer to a declaration}}
+  doIt { if #_hasSymbol(1 as Int) {} } // expected-error {{#_hasSymbol condition must refer to a declaration}}
+  doIt { if #_hasSymbol(1 as S) {} } // expected-error {{cannot convert value of type 'Int' to type 'S' in coercion}}
 }
 
 protocol View {}
@@ -120,15 +142,23 @@ protocol View {}
   static func buildEither<Content>(second content: Content) -> Content where Content : View { fatalError() }
 }
 
-struct Image : View {
-}
+struct Image : View {}
 
 struct MyView {
-  @ViewBuilder var body: some View {
-    if #_hasSymbol(global) { // expected-error {{#_hasSymbol is not supported in closures}}
-      Image()
-    } else {
-      Image()
-    }
+  let image = Image()
+  
+  @ViewBuilder var globalView: some View {
+    if #_hasSymbol(global) { image }
+    else { image }
+  }
+  
+  @ViewBuilder var ambiguousFuncView: some View {
+    if #_hasSymbol(ambiguousFunc) { image } // expected-error {{ambiguous use of 'ambiguousFunc()'}}
+    else { image }
+  }
+    
+  @ViewBuilder var localFuncView: some View {
+    if #_hasSymbol(localFunc) { image } // expected-warning {{global function 'localFunc()' is not a weakly linked declaration}}
+    else { image }
   }
 }