[BuilderTransform] Extract support checks and builder info into ResultBuilder type

Author: xedin

include/swift/Sema/ConstraintSystem.h

@@ -105,6 +105,42 @@ enum class FreeTypeVariableBinding {
 
 namespace constraints {
 
+struct ResultBuilder {
+private:
+  DeclContext *DC;
+  Type BuilderType;
+  llvm::SmallDenseMap<DeclName, bool> SupportedOps;
+
+  Identifier BuildOptionalId;
+
+public:
+  ResultBuilder(DeclContext *DC, Type builderType)
+      : DC(DC), BuilderType(builderType) {
+    auto &ctx = DC->getASTContext();
+    // Use buildOptional(_:) if available, otherwise fall back to buildIf
+    // when available.
+    BuildOptionalId =
+        (supports(ctx.Id_buildOptional) || !supports(ctx.Id_buildIf))
+            ? ctx.Id_buildOptional
+            : ctx.Id_buildIf;
+  }
+
+  DeclContext *getDeclContext() const { return DC; }
+
+  Type getType() const { return BuilderType; }
+
+  NominalTypeDecl *getBuilderDecl() const {
+    return BuilderType->getAnyNominal();
+  }
+
+  Identifier getBuildOptionalId() const { return BuildOptionalId; }
+
+  bool supports(Identifier fnBaseName, ArrayRef<Identifier> argLabels = {},
+                bool checkAvailability = false);
+
+  bool supportsOptional() { return supports(getBuildOptionalId()); }
+};
+
 /// Describes the algorithm to use for trailing closure matching.
 enum class TrailingClosureMatching {
   /// Match a trailing closure to the first parameter that appears to work.

lib/Sema/BuilderTransform.cpp

@@ -68,10 +68,8 @@ class BuilderClosureVisitor
   ConstraintSystem *cs;
   DeclContext *dc;
   ASTContext &ctx;
-  Type builderType;
-  NominalTypeDecl *builder = nullptr;
-  Identifier buildOptionalId;
-  llvm::SmallDenseMap<DeclName, bool> supportedOps;
+
+  ResultBuilder builder;
 
   /// The variable used as a base for all `build*` operations added
   /// by this transform.
@@ -121,21 +119,6 @@ class BuilderClosureVisitor
     return CallExpr::createImplicit(ctx, memberRef, argList);
   }
 
-  /// Check whether the builder supports the given operation.
-  bool builderSupports(Identifier fnBaseName,
-                       ArrayRef<Identifier> argLabels = {},
-                       bool checkAvailability = false) {
-    DeclName name(dc->getASTContext(), fnBaseName, argLabels);
-    auto known = supportedOps.find(name);
-    if (known != supportedOps.end()) {
-      return known->second;
-    }
-
-    return supportedOps[name] = TypeChecker::typeSupportsBuilderOp(
-               builderType, dc, fnBaseName, argLabels, /*allResults*/ {},
-               checkAvailability);
-  }
-
   /// Build an implicit variable in this context.
   VarDecl *buildVar(SourceLoc loc) {
     // Create the implicit variable.
@@ -192,27 +175,19 @@ class BuilderClosureVisitor
 public:
   BuilderClosureVisitor(ASTContext &ctx, ConstraintSystem *cs, DeclContext *dc,
                         Type builderType, Type bodyResultType)
-      : cs(cs), dc(dc), ctx(ctx), builderType(builderType) {
-    builder = builderType->getAnyNominal();
-    applied.builderType = builderType;
+      : cs(cs), dc(dc), ctx(ctx),
+        builder(dc, cs ? cs->simplifyType(builderType) : builderType) {
+    applied.builderType = builder.getType();
     applied.bodyResultType = bodyResultType;
 
-    // Use buildOptional(_:) if available, otherwise fall back to buildIf
-    // when available.
-    if (builderSupports(ctx.Id_buildOptional) ||
-        !builderSupports(ctx.Id_buildIf))
-      buildOptionalId = ctx.Id_buildOptional;
-    else
-      buildOptionalId = ctx.Id_buildIf;
-
     // If we are about to generate constraints, let's establish builder
     // variable for the base of `build*` calls.
     if (cs) {
       builderVar = new (ctx) VarDecl(
           /*isStatic=*/false, VarDecl::Introducer::Let,
           /*nameLoc=*/SourceLoc(), ctx.Id_builderSelf, dc);
       builderVar->setImplicit();
-      cs->setType(builderVar, MetatypeType::get(cs->simplifyType(builderType)));
+      cs->setType(builderVar, MetatypeType::get(builder.getType()));
     }
   }
 
@@ -226,7 +201,7 @@ class BuilderClosureVisitor
 
     // If there is a buildFinalResult(_:), call it.
     ASTContext &ctx = cs->getASTContext();
-    if (builderSupports(ctx.Id_buildFinalResult, { Identifier() })) {
+    if (builder.supports(ctx.Id_buildFinalResult, {Identifier()})) {
       applied.returnExpr = buildCallIfWanted(
           applied.returnExpr->getLoc(), ctx.Id_buildFinalResult,
           { applied.returnExpr }, { Identifier() });
@@ -350,7 +325,7 @@ class BuilderClosureVisitor
       }
 
       auto expr = node.get<Expr *>();
-      if (cs && builderSupports(ctx.Id_buildExpression)) {
+      if (cs && builder.supports(ctx.Id_buildExpression)) {
         expr = buildCallIfWanted(expr->getLoc(), ctx.Id_buildExpression,
                                  { expr }, { Identifier() });
       }
@@ -366,11 +341,11 @@ class BuilderClosureVisitor
     // `buildPartialBlock(accumulated:next:)`, use this to combine
     // subexpressions pairwise.
     if (!expressions.empty() &&
-        builderSupports(ctx.Id_buildPartialBlock, {ctx.Id_first},
-                        /*checkAvailability*/ true) &&
-        builderSupports(ctx.Id_buildPartialBlock,
-                        {ctx.Id_accumulated, ctx.Id_next},
-                        /*checkAvailability*/ true)) {
+        builder.supports(ctx.Id_buildPartialBlock, {ctx.Id_first},
+                         /*checkAvailability*/ true) &&
+        builder.supports(ctx.Id_buildPartialBlock,
+                         {ctx.Id_accumulated, ctx.Id_next},
+                         /*checkAvailability*/ true)) {
       // NOTE: The current implementation uses one-way constraints in between
       // subexpressions. It's functionally equivalent to the following:
       //   let v0 = Builder.buildPartialBlock(first: arg_0)
@@ -391,11 +366,11 @@ class BuilderClosureVisitor
     // If `buildBlock` does not exist at this point, it could be the case that
     // `buildPartialBlock` did not have the sufficient availability for this
     // call site.  Diagnose it.
-    else if (!builderSupports(ctx.Id_buildBlock)) {
+    else if (!builder.supports(ctx.Id_buildBlock)) {
       ctx.Diags.diagnose(
           braceStmt->getStartLoc(),
           diag::result_builder_missing_available_buildpartialblock,
-          builderType);
+          builder.getType());
       return nullptr;
     }
     // Otherwise, call `buildBlock` on all subexpressions.
@@ -462,14 +437,14 @@ class BuilderClosureVisitor
       return false;
 
     // If there's a missing 'else', we need 'buildOptional' to exist.
-    if (isOptional && !builderSupports(buildOptionalId))
+    if (isOptional && !builder.supportsOptional())
       return false;
 
     // If there are multiple clauses, we need 'buildEither(first:)' and
     // 'buildEither(second:)' to both exist.
     if (numPayloads > 1) {
-      if (!builderSupports(ctx.Id_buildEither, {ctx.Id_first}) ||
-          !builderSupports(ctx.Id_buildEither, {ctx.Id_second}))
+      if (!builder.supports(ctx.Id_buildEither, {ctx.Id_first}) ||
+          !builder.supports(ctx.Id_buildEither, {ctx.Id_second}))
         return false;
     }
 
@@ -543,7 +518,7 @@ class BuilderClosureVisitor
     // buildLimitedAvailability(_:).
     auto availabilityCond = findAvailabilityCondition(ifStmt->getCond());
     bool supportsAvailability =
-        availabilityCond && builderSupports(ctx.Id_buildLimitedAvailability);
+        availabilityCond && builder.supports(ctx.Id_buildLimitedAvailability);
     if (supportsAvailability &&
         !availabilityCond->getAvailability()->isUnavailability()) {
       thenVarRefExpr = buildCallIfWanted(ifStmt->getThenStmt()->getEndLoc(),
@@ -592,10 +567,12 @@ class BuilderClosureVisitor
     // The operand should have optional type if we had optional results,
     // so we just need to call `buildIf` now, since we're at the top level.
     if (isOptional && isTopLevel) {
-      thenExpr = buildCallIfWanted(ifStmt->getEndLoc(), buildOptionalId,
-                                   thenExpr,  /*argLabels=*/{ });
-      elseExpr = buildCallIfWanted(ifStmt->getEndLoc(), buildOptionalId,
-                                   elseExpr,  /*argLabels=*/{ });
+      thenExpr =
+          buildCallIfWanted(ifStmt->getEndLoc(), builder.getBuildOptionalId(),
+                            thenExpr, /*argLabels=*/{});
+      elseExpr =
+          buildCallIfWanted(ifStmt->getEndLoc(), builder.getBuildOptionalId(),
+                            elseExpr, /*argLabels=*/{});
     }
 
     thenExpr = cs->generateConstraints(thenExpr, dc);
@@ -832,7 +809,7 @@ class BuilderClosureVisitor
   VarDecl *visitForEachStmt(ForEachStmt *forEachStmt) {
     // for...in statements are handled via buildArray(_:); bail out if the
     // builder does not support it.
-    if (!builderSupports(ctx.Id_buildArray)) {
+    if (!builder.supports(ctx.Id_buildArray)) {
       if (!unhandledNode)
         unhandledNode = forEachStmt;
       return nullptr;
@@ -2220,3 +2197,17 @@ void swift::printResultBuilderBuildFunction(
     printer << "}";
   }
 }
+
+bool ResultBuilder::supports(Identifier fnBaseName,
+                             ArrayRef<Identifier> argLabels,
+                             bool checkAvailability) {
+  DeclName name(DC->getASTContext(), fnBaseName, argLabels);
+  auto known = SupportedOps.find(name);
+  if (known != SupportedOps.end()) {
+    return known->second;
+  }
+
+  return SupportedOps[name] = TypeChecker::typeSupportsBuilderOp(
+             BuilderType, DC, fnBaseName, argLabels, /*allResults*/ {},
+             checkAvailability);
+}