[BuilderTransform] Cache previously applied transforms

Since result builder is just an AST transformation, the result
of a successful transform could be cache and reused with a different
`$__builderSelf` type.

Author: xedin

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 = {},
@@ -2734,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;
@@ -3746,6 +3756,27 @@ class ConstraintSystem {
     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);

lib/Sema/BuilderTransform.cpp

@@ -942,6 +942,8 @@ class ResultBuilderTransform
     return castToStmt<BraceStmt>(newBody.get());
   }
 
+  VarDecl *getBuilderSelf() const { return builder.getBuilderSelf(); }
+
 protected:
   NullablePtr<Stmt> failTransform(UnsupportedElt unsupported) {
     recordUnsupported(unsupported);
@@ -2426,6 +2428,7 @@ ConstraintSystem::matchResultBuilder(AnyFunctionRef fn, Type builderType,
                                      Type bodyResultType,
                                      ConstraintKind bodyResultConstraintKind,
                                      ConstraintLocatorBuilder locator) {
+  builderType = simplifyType(builderType);
   auto builder = builderType->getAnyNominal();
   assert(builder && "Bad result builder type");
   assert(builder->getAttrs().hasAttribute<ResultBuilderAttr>());
@@ -2477,58 +2480,72 @@ ConstraintSystem::matchResultBuilder(AnyFunctionRef fn, Type builderType,
     return None;
   }
 
-  if (Context.TypeCheckerOpts.ResultBuilderASTTransform) {
-    ResultBuilderTransform transform(*this, fn.getAsDeclContext(), builderType,
-                                     bodyResultType);
-    auto *body = transform.apply(fn.getBody());
+  if (Context.LangOpts.hasFeature(Feature::ResultBuilderASTTransform)) {
+    auto transformedBody = getBuilderTransformedBody(fn, builder);
+    // If this builder transform has not yet been applied to this function,
+    // let's do it and cache the result.
+    if (!transformedBody) {
+      ResultBuilderTransform transform(*this, fn.getAsDeclContext(),
+                                       builderType, bodyResultType);
+      auto *body = transform.apply(fn.getBody());
 
-    if (auto unsupported = transform.getUnsupportedElement()) {
-      assert(!body);
+      if (auto unsupported = transform.getUnsupportedElement()) {
+        assert(!body);
 
-      // If we aren't supposed to attempt fixes, fail.
-      if (!shouldAttemptFixes()) {
-        return getTypeMatchFailure(locator);
-      }
+        // If we aren't supposed to attempt fixes, fail.
+        if (!shouldAttemptFixes()) {
+          return getTypeMatchFailure(locator);
+        }
 
-      // If we're solving for code completion and the body contains the code
-      // completion location, skipping it won't get us to a useful solution so
-      // just bail.
-      if (isForCodeCompletion() && containsCodeCompletionLoc(fn.getBody())) {
-        return getTypeMatchFailure(locator);
-      }
+        // If we're solving for code completion and the body contains the code
+        // completion location, skipping it won't get us to a useful solution so
+        // just bail.
+        if (isForCodeCompletion() && containsCodeCompletionLoc(fn.getBody())) {
+          return getTypeMatchFailure(locator);
+        }
 
-      // Record the first unhandled construct as a fix.
-      if (recordFix(SkipUnhandledConstructInResultBuilder::create(
-              *this, unsupported, builder, getConstraintLocator(locator)))) {
-        return getTypeMatchFailure(locator);
-      }
+        // Record the first unhandled construct as a fix.
+        if (recordFix(SkipUnhandledConstructInResultBuilder::create(
+                *this, unsupported, builder, getConstraintLocator(locator)))) {
+          return getTypeMatchFailure(locator);
+        }
 
-      if (auto *closure = getAsExpr<ClosureExpr>(fn.getAbstractClosureExpr())) {
-        auto closureTy = getClosureType(closure);
-        simplifyType(closureTy).visit([&](Type componentTy) {
-          if (auto *typeVar = componentTy->getAs<TypeVariableType>()) {
-            assignFixedType(typeVar,
-                            PlaceholderType::get(getASTContext(), typeVar));
-          }
-        });
+        if (auto *closure =
+                getAsExpr<ClosureExpr>(fn.getAbstractClosureExpr())) {
+          auto closureTy = getClosureType(closure);
+          simplifyType(closureTy).visit([&](Type componentTy) {
+            if (auto *typeVar = componentTy->getAs<TypeVariableType>()) {
+              assignFixedType(typeVar,
+                              PlaceholderType::get(getASTContext(), typeVar));
+            }
+          });
+        }
+
+        return getTypeMatchSuccess();
       }
 
-      return getTypeMatchSuccess();
+      transformedBody = std::make_pair(transform.getBuilderSelf(), body);
+      // Record the transformation so it could be re-used if needed.
+      setBuilderTransformedBody(fn, builder, transformedBody->first,
+                                transformedBody->second);
     }
 
+    // Set the type of `$__builderSelf` variable before constraint generation.
+    setType(transformedBody->first, MetatypeType::get(builderType));
+
     if (isDebugMode()) {
       auto &log = llvm::errs();
       auto indent = solverState ? solverState->depth * 2 : 0;
       log.indent(indent) << "------- Transfomed Body -------\n";
-      body->dump(log);
+      transformedBody->second->dump(log);
       log << '\n';
     }
 
     AppliedBuilderTransform transformInfo;
 
     transformInfo.builderType = builderType;
     transformInfo.bodyResultType = bodyResultType;
-    transformInfo.transformedBody = body;
+    transformInfo.transformedBody = transformedBody->second;
 
     // Record the transformation.
     assert(
@@ -2541,7 +2558,7 @@ ConstraintSystem::matchResultBuilder(AnyFunctionRef fn, Type builderType,
     resultBuilderTransformed.insert(
         std::make_pair(fn, std::move(transformInfo)));
 
-    if (generateConstraints(fn, body))
+    if (generateConstraints(fn, transformInfo.transformedBody.get()))
       return getTypeMatchFailure(locator);
 
     return getTypeMatchSuccess();