Merge pull request swiftlang#29786 from DougGregor/function-builders-let-decls

Allow initialized let/var declarations in function builders.

Author: DougGregor

lib/Sema/BuilderTransform.cpp

@@ -235,6 +235,143 @@ class BuilderClosureVisitor
     return nullptr;                                        \
   }
 
+  /// Provide a type for each variable that occurs within the given pattern,
+  /// by matching the pattern structurally with its already-computed pattern
+  /// type. The variables will either get a concrete type (when present in
+  /// the pattern type) or a fresh type variable bound to that part of the
+  /// pattern via a one-way constraint.
+  void bindVariablesInPattern(Pattern *pattern, Type patternType,
+                              ConstraintLocator *locator) {
+    switch (pattern->getKind()) {
+    case PatternKind::Paren: {
+      // Parentheses don't affect the type, but unwrap a paren type if we have
+      // one.
+      Type subPatternType;
+      if (auto parenType = dyn_cast<ParenType>(patternType.getPointer()))
+        subPatternType = parenType->getUnderlyingType();
+      else
+        subPatternType = patternType;
+      return bindVariablesInPattern(
+          cast<ParenPattern>(pattern)->getSubPattern(),
+          subPatternType, locator);
+    }
+
+    case PatternKind::Var:
+      // Var doesn't affect the type.
+      return bindVariablesInPattern(cast<VarPattern>(pattern)->getSubPattern(),
+                                    patternType, locator);
+
+    case PatternKind::Any:
+      // Nothing to bind.
+      return;
+
+    case PatternKind::Named: {
+      auto var = cast<NamedPattern>(pattern)->getDecl();
+
+      /// Create a fresh type variable to describe the type of the
+      Type varType = cs->createTypeVariable(locator, TVO_CanBindToNoEscape);
+
+      auto ROK = ReferenceOwnership::Strong;
+      if (auto *OA = var->getAttrs().getAttribute<ReferenceOwnershipAttr>())
+        ROK = OA->get();
+      switch (optionalityOf(ROK)) {
+      case ReferenceOwnershipOptionality::Required:
+        // FIXME: Can we assert this rather than just checking it.
+        if (auto optPatternType =
+                dyn_cast<OptionalType>(patternType.getPointer())) {
+          // Add a one-way constraint from the type variable to the wrapped
+          // type of the optional.
+          cs->addConstraint(
+            ConstraintKind::OneWayEqual, varType, optPatternType->getBaseType(),
+                            locator);
+
+          // Make the variable type optional.
+          varType = TypeChecker::getOptionalType(var->getLoc(), varType);
+          break;
+        }
+
+        // Fall through to treat this normally.
+        LLVM_FALLTHROUGH;
+
+      case ReferenceOwnershipOptionality::Allowed:
+      case ReferenceOwnershipOptionality::Disallowed:
+        // Add the one-way constraint from the variable type to the pattern
+        // type.
+        cs->addConstraint(ConstraintKind::OneWayEqual, varType, patternType,
+                          locator);
+        break;
+      }
+
+      // Bind the type of the variable.
+      cs->setType(var, varType);
+      return;
+    }
+
+    case PatternKind::Typed: {
+      // Ignore the type itself; it's part of patternType now.
+      return bindVariablesInPattern(
+          cast<TypedPattern>(pattern)->getSubPattern(),
+          patternType, locator);
+    }
+
+    case PatternKind::Tuple: {
+      auto tuplePat = cast<TuplePattern>(pattern);
+      auto tupleType = patternType->castTo<TupleType>();
+      for (unsigned i = 0, e = tuplePat->getNumElements(); i != e; ++i) {
+        bindVariablesInPattern(tuplePat->getElement(i).getPattern(),
+                               tupleType->getElementType(i), locator);
+      }
+      return;
+    }
+
+    // FIXME: Refutable patterns will generate additional constraints.
+#define PATTERN(Id, Parent)
+#define REFUTABLE_PATTERN(Id, Parent) case PatternKind::Id:
+#include "swift/AST/PatternNodes.def"
+      llvm_unreachable("Refutable patterns are not supported here");
+    }
+
+    llvm_unreachable("Unhandled pattern kind");
+  }
+
+  void visitPatternBindingDecl(PatternBindingDecl *patternBinding) {
+    // If any of the entries lacks an initializer, don't handle this node.
+    if (!llvm::all_of(range(patternBinding->getNumPatternEntries()),
+                      [&](unsigned index) {
+            return patternBinding->isExplicitlyInitialized(index);
+        })) {
+      if (!unhandledNode)
+        unhandledNode = patternBinding;
+      return;
+    }
+
+    // If we aren't generating constraints, there's nothing to do.
+    if (!cs)
+      return;
+
+    /// Generate constraints for each pattern binding entry
+    for (unsigned index : range(patternBinding->getNumPatternEntries())) {
+      // Type check the pattern.
+      auto pattern = patternBinding->getPattern(index);
+      auto contextualPattern = ContextualPattern::forRawPattern(pattern, dc);
+      Type patternType = TypeChecker::typeCheckPattern(contextualPattern);
+
+      // Generate constraints for the initialization.
+      auto target = SolutionApplicationTarget::forInitialization(
+          patternBinding->getInit(index), dc, patternType, pattern);
+      if (cs->generateConstraints(target, FreeTypeVariableBinding::Disallow))
+        continue;
+
+      // Keep track of this binding entry.
+      applied.patternBindingEntries.insert({{patternBinding, index}, target});
+
+      // Bind the variables that occur in the pattern to the corresponding
+      // type entry for the pattern itself.
+      bindVariablesInPattern(pattern, cs->getType(pattern),
+                             cs->getConstraintLocator(target.getAsExpr()));
+    }
+  }
+
   VarDecl *visitBraceStmt(BraceStmt *braceStmt) {
     SmallVector<Expr *, 4> expressions;
     auto addChild = [&](VarDecl *childVar) {
@@ -270,6 +407,18 @@ class BuilderClosureVisitor
           continue;
         }
 
+        // Pattern bindings are okay so long as all of the entries are
+        // initialized.
+        if (auto patternBinding = dyn_cast<PatternBindingDecl>(decl)) {
+          visitPatternBindingDecl(patternBinding);
+          continue;
+        }
+
+        // Ignore variable declarations, because they're always handled within
+        // their enclosing pattern bindings.
+        if (isa<VarDecl>(decl))
+          continue;
+
         if (!unhandledNode)
           unhandledNode = decl;
 
@@ -744,6 +893,26 @@ class BuilderClosureRewriter
     elements.push_back(pbd);
   }
 
+  /// Produce a final type-checked pattern binding.
+  void finishPatternBindingDecl(PatternBindingDecl *patternBinding) {
+    for (unsigned index : range(patternBinding->getNumPatternEntries())) {
+      // Find the solution application target for this.
+      auto knownTarget =
+          builderTransform.patternBindingEntries.find({patternBinding, index});
+      assert(knownTarget != builderTransform.patternBindingEntries.end());
+
+      // Rewrite the target.
+      auto resultTarget = rewriteTarget(knownTarget->second);
+      if (!resultTarget)
+        continue;
+
+      patternBinding->setPattern(
+          index, resultTarget->getInitializationPattern(),
+          resultTarget->getDeclContext());
+      patternBinding->setInit(index, resultTarget->getAsExpr());
+    }
+  }
+
 public:
   BuilderClosureRewriter(
       const Solution &solution,
@@ -811,12 +980,29 @@ class BuilderClosureRewriter
       auto decl = node.get<Decl *>();
 
       // Skip #if declarations.
-      if (isa<IfConfigDecl>(decl))
+      if (isa<IfConfigDecl>(decl)) {
+        newElements.push_back(decl);
         continue;
+      }
 
       // Diagnose #warning / #error during application.
       if (auto poundDiag = dyn_cast<PoundDiagnosticDecl>(decl)) {
         TypeChecker::typeCheckDecl(poundDiag);
+        newElements.push_back(decl);
+        continue;
+      }
+
+      // Skip variable declarations; they're always part of a pattern
+      // binding.
+      if (isa<VarDecl>(decl)) {
+        newElements.push_back(decl);
+        continue;
+      }
+
+      // Handle pattern bindings.
+      if (auto patternBinding = dyn_cast<PatternBindingDecl>(decl)) {
+        finishPatternBindingDecl(patternBinding);
+        newElements.push_back(decl);
         continue;
       }
 

lib/Sema/ConstraintSystem.h

@@ -697,6 +697,13 @@ struct AppliedBuilderTransform {
 
   /// The return expression, capturing the last value to be emitted.
   Expr *returnExpr = nullptr;
+
+  using PatternEntry = std::pair<const PatternBindingDecl *, unsigned>;
+
+  /// Mapping from specific pattern binding entries to the solution application
+  /// targets capturing their initialization.
+  llvm::DenseMap<PatternEntry, SolutionApplicationTarget>
+      patternBindingEntries;
 };
 
 /// Describes the fixed score of a solution to the constraint system.

test/Constraints/function_builder.swift

@@ -476,3 +476,18 @@ func testIfConditions(cond: Bool, c1: Bool, i1: Int, i2: Int) {
 testIfConditions(cond: true, c1: true, i1: 1, i2: 1)
 // CHECK: testIfConditions
 // CHECK-SAME: hello
+
+// Use a "let" declaration within a function builder.
+tuplify(true) { c in
+  "testLetDeclarations"
+  let (a, b) = (c, c && true)
+  if a == b {
+    "hello"
+    b
+  }
+  a
+}
+// CHECK: testLetDeclarations"
+// CHECK-SAME: hello
+// CHECK-SAME: true
+

test/Constraints/function_builder_diags.swift

@@ -83,7 +83,8 @@ func testDiags() {
   // Declarations
   tuplify(true) { _ in
     17
-    let x = 17 // expected-error{{closure containing a declaration cannot be used with function builder 'TupleBuilder'}}
+    let x = 17
+    let y: Int // expected-error{{closure containing a declaration cannot be used with function builder 'TupleBuilder'}}
     x + 25
   }