Merge pull request #60065 from xedin/result-builder-ast-transform-under-flag

[TypeChecker] Implement result builder transform via AST modification under a flag

Author: xedin

include/swift/AST/AnyFunctionRef.h

@@ -150,6 +150,23 @@ class AnyFunctionRef {
     return cast<AutoClosureExpr>(ACE)->getBody();
   }
 
+  void setParsedBody(BraceStmt *stmt, bool isSingleExpression) {
+    if (auto *AFD = TheFunction.dyn_cast<AbstractFunctionDecl *>()) {
+      AFD->setBody(stmt, AbstractFunctionDecl::BodyKind::Parsed);
+      AFD->setHasSingleExpressionBody(isSingleExpression);
+      return;
+    }
+
+    auto *ACE = TheFunction.get<AbstractClosureExpr *>();
+    if (auto *CE = dyn_cast<ClosureExpr>(ACE)) {
+      CE->setBody(stmt, isSingleExpression);
+      CE->setBodyState(ClosureExpr::BodyState::ReadyForTypeChecking);
+      return;
+    }
+
+    llvm_unreachable("autoclosures don't have statement bodies");
+  }
+
   void setTypecheckedBody(BraceStmt *stmt, bool isSingleExpression) {
     if (auto *AFD = TheFunction.dyn_cast<AbstractFunctionDecl *>()) {
       AFD->setBody(stmt, AbstractFunctionDecl::BodyKind::TypeChecked);
@@ -159,7 +176,9 @@ class AnyFunctionRef {
 
     auto *ACE = TheFunction.get<AbstractClosureExpr *>();
     if (auto *CE = dyn_cast<ClosureExpr>(ACE)) {
-      return CE->setBody(stmt, isSingleExpression);
+      CE->setBody(stmt, isSingleExpression);
+      CE->setBodyState(ClosureExpr::BodyState::TypeCheckedWithSignature);
+      return;
     }
 
     llvm_unreachable("autoclosures don't have statement bodies");

include/swift/AST/Expr.h

@@ -363,6 +363,11 @@ class alignas(8) Expr : public ASTAllocated<Expr> {
     : NumPadBits,
     NumElements : 32
   );
+
+  SWIFT_INLINE_BITFIELD_FULL(TypeJoinExpr, Expr, 32,
+    : NumPadBits,
+    NumElements : 32
+  );
   } Bits;
 
 private:
@@ -5940,6 +5945,55 @@ class PackExpr final : public Expr,
   static bool classof(const Expr *E) { return E->getKind() == ExprKind::Pack; }
 };
 
+class TypeJoinExpr final : public Expr,
+                           private llvm::TrailingObjects<TypeJoinExpr, Expr *> {
+  friend TrailingObjects;
+
+  DeclRefExpr *Var;
+
+  size_t numTrailingObjects() const {
+    return getNumElements();
+  }
+
+  MutableArrayRef<Expr *> getMutableElements() {
+    return { getTrailingObjects<Expr *>(), getNumElements() };
+  }
+
+  TypeJoinExpr(DeclRefExpr *var, ArrayRef<Expr *> elements);
+
+public:
+  static TypeJoinExpr *create(ASTContext &ctx, DeclRefExpr *var,
+                              ArrayRef<Expr *> exprs);
+
+  SourceLoc getLoc() const { return SourceLoc(); }
+  SourceRange getSourceRange() const { return SourceRange(); }
+
+  DeclRefExpr *getVar() const { return Var; }
+
+  void setVar(DeclRefExpr *var) {
+    assert(var && "cannot set variable reference to null");
+    Var = var;
+  }
+
+  ArrayRef<Expr *> getElements() const {
+    return { getTrailingObjects<Expr *>(), getNumElements() };
+  }
+
+  Expr *getElement(unsigned i) const {
+    return getElements()[i];
+  }
+
+  void setElement(unsigned i, Expr *E) {
+    getMutableElements()[i] = E;
+  }
+
+  unsigned getNumElements() const { return Bits.TypeJoinExpr.NumElements; }
+
+  static bool classof(const Expr *E) {
+    return E->getKind() == ExprKind::TypeJoin;
+  }
+};
+
 inline bool Expr::isInfixOperator() const {
   return isa<BinaryExpr>(this) || isa<IfExpr>(this) ||
          isa<AssignExpr>(this) || isa<ExplicitCastExpr>(this);

include/swift/AST/ExprNodes.def

@@ -205,7 +205,8 @@ UNCHECKED_EXPR(KeyPathDot, Expr)
 UNCHECKED_EXPR(OneWay, Expr)
 EXPR(Tap, Expr)
 EXPR(Pack, Expr)
-LAST_EXPR(Pack)
+UNCHECKED_EXPR(TypeJoin, Expr)
+LAST_EXPR(TypeJoin)
 
 #undef EXPR_RANGE
 #undef LITERAL_EXPR

include/swift/AST/Stmt.h

@@ -159,6 +159,12 @@ class BraceStmt final : public Stmt,
                            SourceLoc rbloc,
                            Optional<bool> implicit = None);
 
+  static BraceStmt *createImplicit(ASTContext &ctx,
+                                   ArrayRef<ASTNode> elements) {
+    return create(ctx, /*lbloc=*/SourceLoc(), elements, /*rbloc=*/SourceLoc(),
+                  /*implicit=*/true);
+  }
+
   SourceLoc getLBraceLoc() const { return LBLoc; }
   SourceLoc getRBraceLoc() const { return RBLoc; }
 
@@ -552,6 +558,9 @@ class DoStmt : public LabeledStmt {
                   labelInfo),
       DoLoc(doLoc), Body(body) {}
 
+  static DoStmt *createImplicit(ASTContext &C, LabeledStmtInfo labelInfo,
+                                ArrayRef<ASTNode> body);
+
   SourceLoc getDoLoc() const { return DoLoc; }
 
   SourceLoc getStartLoc() const { return getLabelLocOrKeywordLoc(DoLoc); }
@@ -1001,6 +1010,11 @@ class CaseStmt final
 
   unsigned getNumCaseLabelItems() const { return Bits.CaseStmt.NumPatterns; }
 
+  FallthroughStmt *getFallthroughStmt() const {
+    return hasFallthroughDest() ? *getTrailingObjects<FallthroughStmt *>()
+                                : nullptr;
+  }
+
   NullablePtr<CaseStmt> getFallthroughDest() const {
     return const_cast<CaseStmt &>(*this).getFallthroughDest();
   }
@@ -1030,6 +1044,8 @@ class CaseStmt final
   }
   SourceLoc getEndLoc() const { return getBody()->getEndLoc(); }
 
+  SourceLoc getItemTerminatorLoc() const { return ItemTerminatorLoc; }
+
   SourceRange getLabelItemsRange() const {
     switch (ParentKind) {
     case CaseParentKind::Switch:

include/swift/Basic/Features.def

@@ -101,6 +101,7 @@ EXPERIMENTAL_FEATURE(FlowSensitiveConcurrencyCaptures)
 EXPERIMENTAL_FEATURE(MoveOnly)
 EXPERIMENTAL_FEATURE(OneWayClosureParameters)
 EXPERIMENTAL_FEATURE(TypeWitnessSystemInference)
+EXPERIMENTAL_FEATURE(ResultBuilderASTTransform)
 
 /// Whether to enable experimental differentiable programming features:
 /// `@differentiable` declaration attribute, etc.

include/swift/Sema/ConstraintSystem.h

@@ -134,6 +134,8 @@ struct ResultBuilder {
     return BuilderType->getAnyNominal();
   }
 
+  VarDecl *getBuilderSelf() const { return BuilderSelf; }
+
   Identifier getBuildOptionalId() const { return BuildOptionalId; }
 
   bool supports(Identifier fnBaseName, ArrayRef<Identifier> argLabels = {},
@@ -647,6 +649,10 @@ T *getAsPattern(ASTNode node) {
   return nullptr;
 }
 
+template <typename T = Stmt> T *castToStmt(ASTNode node) {
+  return cast<T>(node.get<Stmt *>());
+}
+
 SourceLoc getLoc(ASTNode node);
 SourceRange getSourceRange(ASTNode node);
 
@@ -910,6 +916,10 @@ struct AppliedBuilderTransform {
   /// converted. Opaque types should be unopened.
   Type bodyResultType;
 
+  /// The version of the original body with result builder applied
+  /// as AST transformation.
+  NullablePtr<BraceStmt> transformedBody;
+
   /// An expression whose value has been recorded for later use.
   struct RecordedExpr {
     /// The temporary value that captures the value of the expression, if
@@ -1059,6 +1069,7 @@ class SolutionApplicationTargetsKey {
     pattern,
     patternBindingEntry,
     varDecl,
+    functionRef,
   };
 
 private:
@@ -1079,6 +1090,8 @@ class SolutionApplicationTargetsKey {
     } patternBindingEntry;
 
     const VarDecl *varDecl;
+
+    const DeclContext *functionRef;
   } storage;
 
 public:
@@ -1124,6 +1137,11 @@ class SolutionApplicationTargetsKey {
     storage.varDecl = varDecl;
   }
 
+  SolutionApplicationTargetsKey(const AnyFunctionRef functionRef) {
+    kind = Kind::functionRef;
+    storage.functionRef = functionRef.getAsDeclContext();
+  }
+
   friend bool operator==(
       SolutionApplicationTargetsKey lhs, SolutionApplicationTargetsKey rhs) {
     if (lhs.kind != rhs.kind)
@@ -1155,6 +1173,9 @@ class SolutionApplicationTargetsKey {
 
     case Kind::varDecl:
       return lhs.storage.varDecl == rhs.storage.varDecl;
+
+    case Kind::functionRef:
+      return lhs.storage.functionRef == rhs.storage.functionRef;
     }
     llvm_unreachable("invalid SolutionApplicationTargetsKey kind");
   }
@@ -1206,6 +1227,11 @@ class SolutionApplicationTargetsKey {
       return hash_combine(
           DenseMapInfo<unsigned>::getHashValue(static_cast<unsigned>(kind)),
           DenseMapInfo<void *>::getHashValue(storage.varDecl));
+
+    case Kind::functionRef:
+      return hash_combine(
+          DenseMapInfo<unsigned>::getHashValue(static_cast<unsigned>(kind)),
+          DenseMapInfo<void *>::getHashValue(storage.functionRef));
     }
     llvm_unreachable("invalid statement kind");
   }
@@ -2710,6 +2736,14 @@ class ConstraintSystem {
   /// diagnostics when result builder has multiple overloads.
   llvm::SmallDenseSet<AnyFunctionRef> InvalidResultBuilderBodies;
 
+  /// The *global* set of all functions that have a particular result builder
+  /// applied.
+  ///
+  /// The value here is `$__builderSelf` variable and a transformed body.
+  llvm::DenseMap<std::pair<AnyFunctionRef, NominalTypeDecl *>,
+                 std::pair<VarDecl *, BraceStmt *>>
+      BuilderTransformedBodies;
+
   /// Arguments after the code completion token that were thus ignored (i.e.
   /// assigned fresh type variables) for type checking.
   llvm::SetVector<Expr *> IgnoredArguments;
@@ -3714,6 +3748,35 @@ class ConstraintSystem {
     return known->second;
   }
 
+  Optional<AppliedBuilderTransform>
+  getAppliedResultBuilderTransform(AnyFunctionRef fn) const {
+    auto transformed = resultBuilderTransformed.find(fn);
+    if (transformed != resultBuilderTransformed.end())
+      return transformed->second;
+    return None;
+  }
+
+  void setBuilderTransformedBody(AnyFunctionRef fn, NominalTypeDecl *builder,
+                                 NullablePtr<VarDecl> builderSelf,
+                                 NullablePtr<BraceStmt> body) {
+    assert(builder->getAttrs().hasAttribute<ResultBuilderAttr>());
+    assert(body);
+    assert(builderSelf);
+
+    auto existing = BuilderTransformedBodies.insert(
+        {{fn, builder}, {builderSelf.get(), body.get()}});
+    assert(existing.second && "Duplicate result builder transform");
+    (void)existing;
+  }
+
+  Optional<std::pair<VarDecl *, BraceStmt *>>
+  getBuilderTransformedBody(AnyFunctionRef fn, NominalTypeDecl *builder) const {
+    auto result = BuilderTransformedBodies.find({fn, builder});
+    if (result == BuilderTransformedBodies.end())
+      return None;
+    return result->second;
+  }
+
   void setCaseLabelItemInfo(const CaseLabelItem *item, CaseLabelItemInfo info) {
     assert(item != nullptr);
     assert(caseLabelItems.count(item) == 0);
@@ -4792,6 +4855,16 @@ class ConstraintSystem {
   LLVM_NODISCARD
   bool generateConstraints(ClosureExpr *closure);
 
+  /// Generate constraints for the body of the given function.
+  ///
+  /// \param fn The function or closure expression
+  /// \param body The body of the given function that should be
+  /// used for constraint generation.
+  ///
+  /// \returns \c true if constraint generation failed, \c false otherwise
+  LLVM_NODISCARD
+  bool generateConstraints(AnyFunctionRef fn, BraceStmt *body);
+
   /// Generate constraints for the given (unchecked) expression.
   ///
   /// \returns a possibly-sanitized expression, or null if an error occurred.
@@ -5640,14 +5713,14 @@ class ConstraintSystem {
   ///
   ///
   /// \param solution The solution to apply.
-  /// \param closure The closure to which the solution is being applied.
+  /// \param fn The function or closure to which the solution is being applied.
   /// \param currentDC The declaration context in which transformations
   /// will be applied.
   /// \param rewriteTarget Function that performs a rewrite of any
   /// solution application target within the context.
   ///
   /// \returns true if solution cannot be applied.
-  bool applySolutionToBody(Solution &solution, ClosureExpr *closure,
+  bool applySolutionToBody(Solution &solution, AnyFunctionRef fn,
                            DeclContext *&currentDC,
                            std::function<Optional<SolutionApplicationTarget>(
                                SolutionApplicationTarget)>
@@ -6485,6 +6558,13 @@ bool isOperatorDisjunction(Constraint *disjunction);
 /// or nested declarations).
 ASTNode findAsyncNode(ClosureExpr *closure);
 
+/// Check whether the given binding represents a placeholder variable that
+/// has to get its type inferred at a first use site.
+///
+/// \returns The currently assigned type if it's a placeholder,
+/// empty type otherwise.
+Type isPlaceholderVar(PatternBindingDecl *PB);
+
 } // end namespace constraints
 
 template<typename ...Args>

lib/AST/ASTDumper.cpp

@@ -2905,6 +2905,21 @@ class PrintExpr : public ExprVisitor<PrintExpr> {
     printRec(E->getBody(), E->getVar()->getDeclContext()->getASTContext());
     PrintWithColorRAII(OS, ParenthesisColor) << ')';
   }
+
+  void visitTypeJoinExpr(TypeJoinExpr *E) {
+    printCommon(E, "type_join_expr");
+
+    PrintWithColorRAII(OS, DeclColor) << " var=";
+    printRec(E->getVar());
+    OS << '\n';
+
+    for (auto *member : E->getElements()) {
+      printRec(member);
+      OS << '\n';
+    }
+
+    PrintWithColorRAII(OS, ParenthesisColor) << ')';
+  }
 };
 
 } // end anonymous namespace

lib/AST/ASTPrinter.cpp

@@ -3037,6 +3037,10 @@ static bool usesFeatureOneWayClosureParameters(Decl *decl) {
   return false;
 }
 
+static bool usesFeatureResultBuilderASTTransform(Decl *decl) {
+  return false;
+}
+
 static bool usesFeatureTypeWitnessSystemInference(Decl *decl) {
   return false;
 }
@@ -4551,6 +4555,9 @@ void PrintAST::visitPackExpr(PackExpr *expr) {
 void PrintAST::visitReifyPackExpr(ReifyPackExpr *expr) {
 }
 
+void PrintAST::visitTypeJoinExpr(TypeJoinExpr *expr) {
+}
+
 void PrintAST::visitAssignExpr(AssignExpr *expr) {
   visit(expr->getDest());
   Printer << " = ";