[TypeChecker] Incremental multi-statement closure type-checking (disabled by default)

include/swift/AST/ASTNode.h

@@ -17,6 +17,7 @@
 #ifndef SWIFT_AST_AST_NODE_H
 #define SWIFT_AST_AST_NODE_H
 
+#include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/PointerUnion.h"
 #include "swift/Basic/Debug.h"
 #include "swift/AST/TypeAlignments.h"
@@ -35,13 +36,18 @@ namespace swift {
   class SourceLoc;
   class SourceRange;
   class ASTWalker;
+  class StmtConditionElement;
+  class CaseLabelItem;
   enum class ExprKind : uint8_t;
   enum class DeclKind : uint8_t;
   enum class PatternKind : uint8_t;
   enum class StmtKind;
 
-  struct ASTNode : public llvm::PointerUnion<Expr *, Stmt *, Decl *, Pattern *,
-                                             TypeRepr *> {
+  using StmtCondition = llvm::MutableArrayRef<StmtConditionElement>;
+
+  struct ASTNode
+      : public llvm::PointerUnion<Expr *, Stmt *, Decl *, Pattern *, TypeRepr *,
+                                  StmtCondition *, CaseLabelItem *> {
     // Inherit the constructors from PointerUnion.
     using PointerUnion::PointerUnion;
 

include/swift/AST/Decl.h

@@ -1805,7 +1805,12 @@ class PatternBindingDecl final : public Decl,
     return getPatternList()[i].getPattern();
   }
 
-  void setPattern(unsigned i, Pattern *Pat, DeclContext *InitContext);
+  void setPattern(unsigned i, Pattern *Pat, DeclContext *InitContext,
+                  bool isFullyValidated = false);
+
+  bool isFullyValidated(unsigned i) const {
+    return getPatternList()[i].isFullyValidated();
+  }
 
   DeclContext *getInitContext(unsigned i) const {
     return getPatternList()[i].getInitContext();

include/swift/AST/Stmt.h

@@ -579,6 +579,10 @@ class LabeledConditionalStmt : public LabeledStmt {
   StmtCondition getCond() const { return Cond; }
   void setCond(StmtCondition e);
 
+  /// FIXME: Find a better way to implement this. Allows conditions to be
+  ///        stored in \c ASTNode.
+  StmtCondition *getCondPointer() { return &Cond; }
+
   static bool classof(const Stmt *S) {
     return S->getKind() >= StmtKind::First_LabeledConditionalStmt &&
            S->getKind() <= StmtKind::Last_LabeledConditionalStmt;

include/swift/AST/TypeAlignments.h

@@ -60,6 +60,7 @@ namespace swift {
   class TypeDecl;
   class TypeRepr;
   class ValueDecl;
+  class CaseLabelItem;
 
   /// We frequently use three tag bits on all of these types.
   constexpr size_t AttrAlignInBits = 3;
@@ -151,6 +152,8 @@ LLVM_DECLARE_TYPE_ALIGNMENT(swift::AttributeBase, swift::AttrAlignInBits)
 
 LLVM_DECLARE_TYPE_ALIGNMENT(swift::TypeRepr, swift::TypeReprAlignInBits)
 
+LLVM_DECLARE_TYPE_ALIGNMENT(swift::CaseLabelItem, swift::PatternAlignInBits);
+
 static_assert(alignof(void*) >= 2, "pointer alignment is too small");
 
 #endif

include/swift/Basic/LangOptions.h

@@ -662,6 +662,10 @@ namespace swift {
     /// parameters of closures.
     bool EnableOneWayClosureParameters = false;
 
+    /// Enable experimental support for type inference through multi-statement
+    /// closures.
+    bool EnableMultiStatementClosureInference = false;
+
     /// See \ref FrontendOptions.PrintFullConvention
     bool PrintFullConvention = false;
   };

include/swift/Option/FrontendOptions.td

@@ -767,6 +767,10 @@ def experimental_one_way_closure_params :
   Flag<["-"], "experimental-one-way-closure-params">,
   HelpText<"Enable experimental support for one-way closure parameters">;
 
+def experimental_multi_statement_closures :
+  Flag<["-"], "experimental-multi-statement-closures">,
+  HelpText<"Enable experimental support for type inference in multi-statement closures">;
+
 def prebuilt_module_cache_path :
   Separate<["-"], "prebuilt-module-cache-path">,
   HelpText<"Directory of prebuilt modules for loading module interfaces">;

include/swift/Sema/Constraint.h

@@ -18,10 +18,13 @@
 #ifndef SWIFT_SEMA_CONSTRAINT_H
 #define SWIFT_SEMA_CONSTRAINT_H
 
+#include "swift/AST/ASTNode.h"
 #include "swift/AST/FunctionRefKind.h"
 #include "swift/AST/Identifier.h"
 #include "swift/AST/Type.h"
+#include "swift/AST/TypeLoc.h"
 #include "swift/Basic/Debug.h"
+#include "swift/Sema/ConstraintLocator.h"
 #include "swift/Sema/OverloadChoice.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/ilist.h"
@@ -47,6 +50,23 @@ class ConstraintLocator;
 class ConstraintSystem;
 enum class TrailingClosureMatching;
 
+/// Describes contextual type information about a particular element
+/// (expression, statement etc.) within a constraint system.
+struct ContextualTypeInfo {
+  TypeLoc typeLoc;
+  ContextualTypePurpose purpose;
+
+  ContextualTypeInfo() : typeLoc(TypeLoc()), purpose(CTP_Unused) {}
+
+  ContextualTypeInfo(Type contextualTy, ContextualTypePurpose purpose)
+      : typeLoc(TypeLoc::withoutLoc(contextualTy)), purpose(purpose) {}
+
+  ContextualTypeInfo(TypeLoc typeLoc, ContextualTypePurpose purpose)
+      : typeLoc(typeLoc), purpose(purpose) {}
+
+  Type getType() const { return typeLoc.getType(); }
+};
+
 /// Describes the kind of constraint placed on one or more types.
 enum class ConstraintKind : char {
   /// The two types must be bound to the same type. This is the only
@@ -129,6 +149,9 @@ enum class ConstraintKind : char {
   /// A disjunction constraint that specifies that one or more of the
   /// stored constraints must hold.
   Disjunction,
+  /// A conjunction constraint that specifies that all of the stored
+  /// constraints must hold.
+  Conjunction,
   /// The first type is an optional type whose object type is the second
   /// type, preserving lvalue-ness.
   OptionalObject,
@@ -192,6 +215,10 @@ enum class ConstraintKind : char {
   /// inferred from a conversion, so the check is more relax comparing to
   /// `ConformsTo`.
   TransitivelyConformsTo,
+  /// Represents an AST node contained in a body of a closure. It has only
+  /// one type - type variable representing type of a node, other side is
+  /// the AST node to infer the type for.
+  ClosureBodyElement,
 };
 
 /// Classification of the different kinds of constraints.
@@ -208,7 +235,13 @@ enum class ConstraintClassification : char {
   TypeProperty,
 
   /// A disjunction constraint.
-  Disjunction
+  Disjunction,
+
+  /// A conjunction constraint.
+  Conjunction,
+
+  /// An element of a closure body.
+  ClosureElement,
 };
 
 /// Specifies a restriction on the kind of conversion that should be
@@ -338,6 +371,10 @@ class Constraint final : public llvm::ilist_node<Constraint>,
   /// in its disjunction.
   unsigned IsFavored : 1;
 
+  /// Whether or not this constraint should be solved in isolation from
+  /// the rest of the constraint system. Currently only applies to conjunctions.
+  unsigned IsIsolated : 1;
+
   /// The number of type variables referenced by this constraint.
   ///
   /// The type variables themselves are tail-allocated.
@@ -399,6 +436,13 @@ class Constraint final : public llvm::ilist_node<Constraint>,
       /// The DC in which the use appears.
       DeclContext *UseDC;
     } Overload;
+
+    struct {
+      /// The node itself.
+      ASTNode Element;
+      /// Contextual information associated with the element (if any).
+      ContextualTypeInfo Context;
+    } ClosureElement;
   };
 
   /// The locator that describes where in the expression this
@@ -410,7 +454,7 @@ class Constraint final : public llvm::ilist_node<Constraint>,
   void *operator new(size_t) = delete;
 
   Constraint(ConstraintKind kind, ArrayRef<Constraint *> constraints,
-             ConstraintLocator *locator,
+             bool isIsolated, ConstraintLocator *locator,
              SmallPtrSetImpl<TypeVariableType *> &typeVars);
 
   /// Construct a new constraint.
@@ -450,6 +494,11 @@ class Constraint final : public llvm::ilist_node<Constraint>,
              ConstraintLocator *locator,
              SmallPtrSetImpl<TypeVariableType *> &typeVars);
 
+  /// Construct a closure body element constraint.
+  Constraint(ASTNode node, ContextualTypeInfo context,
+             ConstraintLocator *locator,
+             SmallPtrSetImpl<TypeVariableType *> &typeVars);
+
   /// Retrieve the type variables buffer, for internal mutation.
   MutableArrayRef<TypeVariableType *> getTypeVariablesBuffer() {
     return { getTrailingObjects<TypeVariableType *>(), NumTypeVariables };
@@ -518,12 +567,31 @@ class Constraint final : public llvm::ilist_node<Constraint>,
                                        RememberChoice_t shouldRememberChoice
                                          = ForgetChoice);
 
+  /// Create a new conjunction constraint.
+  ///
+  /// \param isIsolated - Indicates whether given constraint should be
+  /// solved in isolation from the rest of the constraint system i.e.
+  /// by removing all of the unrelated type variables and constraints.
+  static Constraint *
+  createConjunction(ConstraintSystem &cs, ArrayRef<Constraint *> constraints,
+                    bool isIsolated, ConstraintLocator *locator,
+                    ArrayRef<TypeVariableType *> referencedVars = {});
+
   /// Create a new Applicable Function constraint.
   static Constraint *createApplicableFunction(
       ConstraintSystem &cs, Type argumentFnType, Type calleeType,
       Optional<TrailingClosureMatching> trailingClosureMatching,
       ConstraintLocator *locator);
 
+  static Constraint *createClosureBodyElement(ConstraintSystem &cs,
+                                              ASTNode node,
+                                              ConstraintLocator *locator);
+
+  static Constraint *createClosureBodyElement(ConstraintSystem &cs,
+                                              ASTNode node,
+                                              ContextualTypeInfo context,
+                                              ConstraintLocator *locator);
+
   /// Determine the kind of constraint.
   ConstraintKind getKind() const { return Kind; }
 
@@ -629,6 +697,12 @@ class Constraint final : public llvm::ilist_node<Constraint>,
 
     case ConstraintKind::Disjunction:
       return ConstraintClassification::Disjunction;
+
+    case ConstraintKind::Conjunction:
+      return ConstraintClassification::Conjunction;
+
+    case ConstraintKind::ClosureBodyElement:
+      return ConstraintClassification::ClosureElement;
     }
 
     llvm_unreachable("Unhandled ConstraintKind in switch.");
@@ -640,6 +714,9 @@ class Constraint final : public llvm::ilist_node<Constraint>,
     case ConstraintKind::Disjunction:
       llvm_unreachable("disjunction constraints have no type operands");
 
+    case ConstraintKind::Conjunction:
+      llvm_unreachable("conjunction constraints have no type operands");
+
     case ConstraintKind::BindOverload:
       return Overload.First;
 
@@ -648,6 +725,9 @@ class Constraint final : public llvm::ilist_node<Constraint>,
     case ConstraintKind::ValueWitness:
       return Member.First;
 
+    case ConstraintKind::ClosureBodyElement:
+      llvm_unreachable("closure body element constraint has no type operands");
+
     default:
       return Types.First;
     }
@@ -657,7 +737,9 @@ class Constraint final : public llvm::ilist_node<Constraint>,
   Type getSecondType() const {
     switch (getKind()) {
     case ConstraintKind::Disjunction:
+    case ConstraintKind::Conjunction:
     case ConstraintKind::BindOverload:
+    case ConstraintKind::ClosureBodyElement:
       llvm_unreachable("constraint has no second type");
 
     case ConstraintKind::ValueMember:
@@ -710,7 +792,8 @@ class Constraint final : public llvm::ilist_node<Constraint>,
 
   /// Retrieve the set of constraints in a disjunction.
   ArrayRef<Constraint *> getNestedConstraints() const {
-    assert(Kind == ConstraintKind::Disjunction);
+    assert(Kind == ConstraintKind::Disjunction ||
+           Kind == ConstraintKind::Conjunction);
     return Nested;
   }
 
@@ -735,6 +818,10 @@ class Constraint final : public llvm::ilist_node<Constraint>,
   /// e.g. coercion constraint "as X" which forms a disjunction.
   bool isExplicitConversion() const;
 
+  /// Determine whether this constraint should be solved in isolation
+  /// from the rest of the constraint system.
+  bool isIsolated() const { return IsIsolated; }
+
   /// Whether this is a one-way constraint.
   bool isOneWayConstraint() const {
     return Kind == ConstraintKind::OneWayEqual ||
@@ -761,6 +848,16 @@ class Constraint final : public llvm::ilist_node<Constraint>,
     return Member.UseDC;
   }
 
+  ASTNode getClosureElement() const {
+    assert(Kind == ConstraintKind::ClosureBodyElement);
+    return ClosureElement.Element;
+  }
+
+  ContextualTypeInfo getElementContext() const {
+    assert(Kind == ConstraintKind::ClosureBodyElement);
+    return ClosureElement.Context;
+  }
+
   /// For an applicable function constraint, retrieve the trailing closure
   /// matching rule.
   Optional<TrailingClosureMatching> getTrailingClosureMatching() const;

include/swift/Sema/ConstraintLocator.h

@@ -72,8 +72,15 @@ enum ContextualTypePurpose : uint8_t {
                              ///< result type.
   CTP_Condition,        ///< Condition expression of various statements e.g.
                         ///< `if`, `for`, `while` etc.
+  CTP_CaseStmt,         ///< A single case statement associated with a `switch` or
+                        ///  a `do-catch` statement. It has to be convertible
+                        ///  to a type of a switch subject or an `Error` type.
   CTP_ForEachStmt,      ///< "expression/sequence" associated with 'for-in' loop
                         ///< is expected to conform to 'Sequence' protocol.
+  CTP_ForEachSequence,  ///< Sequence expression associated with `for-in` loop,
+                        ///  this element acts slightly differently compared to
+                        ///  \c CTP_ForEachStmt in a sence that it would
+                        ///  produce conformance constraints.
   CTP_WrappedProperty,  ///< Property type expected to match 'wrappedValue' type
   CTP_ComposedPropertyWrapper, ///< Composed wrapper type expected to match
                                ///< former 'wrappedValue' type
@@ -953,6 +960,52 @@ class LocatorPathElt::KeyPathType final
   }
 };
 
+class LocatorPathElt::ClosureBodyElement final
+    : public StoredPointerElement<void> {
+public:
+  ClosureBodyElement(ASTNode element)
+      : StoredPointerElement(PathElementKind::ClosureBodyElement,
+                             element.getOpaqueValue()) {
+    assert(element);
+  }
+
+  ASTNode getElement() const {
+    // Unfortunately \c getFromOpaqueValue doesn't produce an ASTNode.
+    auto node = ASTNode::getFromOpaqueValue(getStoredPointer());
+    if (auto *expr = node.dyn_cast<Expr *>())
+      return expr;
+
+    if (auto *stmt = node.dyn_cast<Stmt *>())
+      return stmt;
+
+    if (auto *decl = node.dyn_cast<Decl *>())
+      return decl;
+
+    if (auto *pattern = node.dyn_cast<Pattern *>())
+      return pattern;
+
+    if (auto *repr = node.dyn_cast<TypeRepr *>())
+      return repr;
+
+    if (auto *cond = node.dyn_cast<StmtCondition *>())
+      return cond;
+
+    if (auto *caseItem = node.dyn_cast<CaseLabelItem *>())
+      return caseItem;
+
+    llvm_unreachable("unhandled ASTNode element kind");
+  }
+
+  Stmt *asStmt() const {
+    auto node = ASTNode::getFromOpaqueValue(getStoredPointer());
+    return node.get<Stmt *>();
+  }
+
+  static bool classof(const LocatorPathElt *elt) {
+    return elt->getKind() == PathElementKind::ClosureBodyElement;
+  }
+};
+
 namespace details {
   template <typename CustomPathElement>
   class PathElement {