From f748c84cf1a1712f43e1abcf452130a74f0c999c Mon Sep 17 00:00:00 2001
From: Josh Learn <joshua.learn@apple.com>
Date: Wed, 15 Dec 2021 12:16:22 -0800
Subject: [PATCH] [Sema][MiscDiag] Fix constantness diag to handle result
 builder patterns

We currently have a problem with how constantness diagnostics
traverse the AST to look for function calls to diagnose. We
special case closure bodies and don't check them (unless they're
single expression closures) because closure bodies are type-
checked separately and will be covered later. This poses a problem
in certain AST structures, such as what we see with result builders,
because the call expressions are rooted in declarations, which aren't
checked in the closure body type-checking covered by MiscDiag.

This patch fixes the problem by manually checking all closure bodies
and stopping misc diagnostics from checking the bodies separately.

rdar://85737300
---
 lib/Sema/ConstantnessSemaDiagnostics.cpp | 39 +++++++++++++-----------
 test/Sema/diag_constantness_check.swift  | 29 ++++++++++++++++++
 2 files changed, 50 insertions(+), 18 deletions(-)

diff --git a/lib/Sema/ConstantnessSemaDiagnostics.cpp b/lib/Sema/ConstantnessSemaDiagnostics.cpp
index ce786b1b517d8..79f7d0f0a8d87 100644
--- a/lib/Sema/ConstantnessSemaDiagnostics.cpp
+++ b/lib/Sema/ConstantnessSemaDiagnostics.cpp
@@ -334,9 +334,10 @@ void swift::diagnoseConstantArgumentRequirement(
     const Expr *expr, const DeclContext *declContext) {
   class ConstantReqCallWalker : public ASTWalker {
     DeclContext *DC;
+    bool insideClosure;
 
   public:
-    ConstantReqCallWalker(DeclContext *DC) : DC(DC) {}
+    ConstantReqCallWalker(DeclContext *DC) : DC(DC), insideClosure(false) {}
 
     // Descend until we find a call expressions. Note that the input expression
     // could be an assign expression or another expression that contains the
@@ -347,10 +348,15 @@ void swift::diagnoseConstantArgumentRequirement(
       if (auto *closureExpr = dyn_cast<ClosureExpr>(expr)) {
         return walkToClosureExprPre(closureExpr);
       }
+
       // Interpolated expressions' bodies will be type checked
       // separately so exit early to avoid duplicate diagnostics.
+      // The caveat is that they won't be checked inside closure
+      // bodies because we manually check all closures to avoid
+      // duplicate diagnostics. Therefore we must still descend into
+      // interpolated expressions if we are inside of a closure.
       if (!expr || isa<ErrorExpr>(expr) || !expr->getType() ||
-          isa<InterpolatedStringLiteralExpr>(expr))
+          (isa<InterpolatedStringLiteralExpr>(expr) && !insideClosure))
         return {false, expr};
       if (auto *callExpr = dyn_cast<CallExpr>(expr)) {
         diagnoseConstantArgumentRequirementOfCall(callExpr, DC->getASTContext());
@@ -359,33 +365,30 @@ void swift::diagnoseConstantArgumentRequirement(
     }
     
     std::pair<bool, Expr *> walkToClosureExprPre(ClosureExpr *closure) {
-      auto &ctx = DC->getASTContext();
-
-      if (closure->hasSingleExpressionBody() ||
-          ctx.TypeCheckerOpts.EnableMultiStatementClosureInference) {
-        // Closure bodies are not visited directly by the ASTVisitor,
-        // so we must descend into the body manuall and set the
-        // DeclContext to that of the closure.
-        DC = closure;
-        return {true, closure};
-      }
-      return {false, closure};
+      DC = closure;
+      insideClosure = true;
+      return {true, closure};
     }
     
     Expr *walkToExprPost(Expr *expr) override {
       if (auto *closureExpr = dyn_cast<ClosureExpr>(expr)) {
         // Reset the DeclContext to the outer scope if we descended
-        // into a closure expr.
+        // into a closure expr and check whether or not we are still
+        // within a closure context.
         DC = closureExpr->getParent();
+        insideClosure = isa<ClosureExpr>(DC);
       }
       return expr;
     }
-    
-    std::pair<bool, Stmt *> walkToStmtPre(Stmt *stmt) override {
-      return {true, stmt};
-    }
   };
 
+  // We manually check closure bodies from their outer contexts,
+  // so bail early if we are being called directly on expressions
+  // inside of a closure body.
+  if (isa<ClosureExpr>(declContext)) {
+    return;
+  }
+
   ConstantReqCallWalker walker(const_cast<DeclContext *>(declContext));
   const_cast<Expr *>(expr)->walk(walker);
 }
diff --git a/test/Sema/diag_constantness_check.swift b/test/Sema/diag_constantness_check.swift
index f03830d7adaa5..2969467d80c99 100644
--- a/test/Sema/diag_constantness_check.swift
+++ b/test/Sema/diag_constantness_check.swift
@@ -424,3 +424,32 @@ func testCallsWithinClosures(s: String, x: Int) {
     constantArgumentFunction("string with a single interpolation \(x)")
   }
 }
+
+@resultBuilder
+struct MyArrayBuilder {
+    typealias Component = [Int]
+    typealias Expression = Int
+    static func buildExpression(_ element: Expression) -> Component {
+        return [element]
+    }
+    static func buildBlock(_ components: Component...) -> Component {
+        return Array(components.joined())
+    }
+}
+
+struct MyArray {
+    public init(@MyArrayBuilder arr: () -> [Int]) {}
+}
+
+func testResultBuilder(x: Int, y: Int) -> MyArray {
+    let _: MyArray = MyArray {
+        constantArgumentFunctionReturningInt(x)
+          // expected-error@-1 {{argument must be an integer literal}}
+        constantArgumentFunctionReturningInt(y)
+          // expected-error@-1 {{argument must be an integer literal}}
+    }
+    let _: MyArray = MyArray {
+        constantArgumentFunctionReturningInt(x)
+          // expected-error@-1 {{argument must be an integer literal}}
+    }
+}