Merge pull request #39575 from apple/set-types-for-case-vars

[ResultBuilders] Make sure that `case` body variables get types

Author: xedin

Diff for: lib/Sema/BuilderTransform.cpp

@@ -1462,6 +1462,14 @@ class BuilderClosureRewriter
         return nullptr;
     }
 
+    // Setup the types of our case body var decls.
+    for (auto *expected : caseStmt->getCaseBodyVariablesOrEmptyArray()) {
+      assert(expected->hasName());
+      auto prev = expected->getParentVarDecl();
+      auto type = solution.resolveInterfaceType(solution.getType(prev));
+      expected->setInterfaceType(type);
+    }
+
     // Transform the body of the case.
     auto body = cast<BraceStmt>(caseStmt->getBody());
     auto captured = takeCapturedStmt(body);
@@ -1721,7 +1729,16 @@ Optional<BraceStmt *> TypeChecker::applyResultBuilderBodyTransform(
           solutions.front(),
           SolutionApplicationTarget(func))) {
     performSyntacticDiagnosticsForTarget(*result, /*isExprStmt*/ false);
-    return result->getFunctionBody();
+    auto *body = result->getFunctionBody();
+
+    if (cs.isDebugMode()) {
+      auto &log = llvm::errs();
+      log << "--- Type-checked function body ---\n";
+      body->dump(log);
+      log << '\n';
+    }
+
+    return body;
   }
 
   return nullptr;

Diff for: test/Constraints/result_builder_switch_with_vars.swift

@@ -0,0 +1,88 @@
+// RUN: %target-swift-frontend -typecheck -dump-ast %s | %FileCheck %s
+
+enum Either<T,U> {
+  case first(T)
+  case second(U)
+}
+
+@resultBuilder
+struct TupleBuilder {
+  static func buildBlock<T1>(_ t1: T1) -> (T1) {
+    return (t1)
+  }
+
+  static func buildBlock<T1, T2>(_ t1: T1, _ t2: T2) -> (T1, T2) {
+    return (t1, t2)
+  }
+
+  static func buildBlock<T1, T2, T3>(_ t1: T1, _ t2: T2, _ t3: T3)
+      -> (T1, T2, T3) {
+    return (t1, t2, t3)
+  }
+
+  static func buildBlock<T1, T2, T3, T4>(_ t1: T1, _ t2: T2, _ t3: T3, _ t4: T4)
+      -> (T1, T2, T3, T4) {
+    return (t1, t2, t3, t4)
+  }
+
+  static func buildBlock<T1, T2, T3, T4, T5>(
+    _ t1: T1, _ t2: T2, _ t3: T3, _ t4: T4, _ t5: T5
+  ) -> (T1, T2, T3, T4, T5) {
+    return (t1, t2, t3, t4, t5)
+  }
+
+  static func buildEither<T,U>(first value: T) -> Either<T,U> {
+    return .first(value)
+  }
+  static func buildEither<T,U>(second value: U) -> Either<T,U> {
+    return .second(value)
+  }
+}
+
+enum E {
+case test(a: String, b: Int)
+}
+
+func tuplify<T>(@TupleBuilder body: (E) throws -> T) rethrows {
+}
+
+tuplify {
+  switch $0 {
+  // CHECK: (case_body_variables
+  // CHECK-NEXT: (var_decl implicit {{.*}} "a" type='String' interface type='String' let readImpl=stored immutable)
+  // CHECK-NEXT: (var_decl implicit {{.*}} "b" type='Int' interface type='Int' let readImpl=stored immutable)
+  case let .test(a, b):
+    a
+    b
+  }
+
+  switch $0 {
+    // CHECK: (case_body_variables
+    // CHECK-NEXT: (var_decl implicit {{.*}} "a" type='String' interface type='String' let readImpl=stored immutable)
+    // CHECK-NEXT: (var_decl implicit {{.*}} "b" type='Int' interface type='Int' let readImpl=stored immutable)
+  case .test(let a, let b):
+    a
+    b
+  }
+
+  switch $0 {
+    // CHECK: (case_body_variables
+    // CHECK-NEXT: (var_decl implicit {{.*}} "value" type='(a: String, b: Int)' interface type='(a: String, b: Int)' let readImpl=stored immutable)
+  case let .test((value)):
+    value.a
+  }
+
+  switch $0 {
+    // CHECK: (case_body_variables
+    // CHECK-NEXT: (var_decl implicit {{.*}} "value" type='(a: String, b: Int)' interface type='(a: String, b: Int)' let readImpl=stored immutable)
+  case let .test(value):
+    value.a
+  }
+
+  switch $0 {
+    // CHECK: (case_body_variables
+    // CHECK-NEXT: (var_decl implicit {{.*}} "value" type='(a: String, b: Int)' interface type='(a: String, b: Int)' let readImpl=stored immutable)
+  case .test(let value):
+    value.a
+  }
+}