Add PreCheckFunctionBuilderRequest

Drop a cache out of the type checker.

Author: CodaFi

include/swift/AST/ASTTypeIDZone.def

@@ -17,6 +17,7 @@
 
 SWIFT_TYPEID(AncestryFlags)
 SWIFT_TYPEID(CtorInitializerKind)
+SWIFT_TYPEID(FunctionBuilderClosurePreCheck)
 SWIFT_TYPEID(GenericSignature)
 SWIFT_TYPEID(ImplicitMemberAction)
 SWIFT_TYPEID(ParamSpecifier)
@@ -28,6 +29,7 @@ SWIFT_TYPEID(Type)
 SWIFT_TYPEID(TypePair)
 SWIFT_TYPEID(TypeWitnessAndDecl)
 SWIFT_TYPEID(Witness)
+SWIFT_TYPEID_NAMED(ClosureExpr *, ClosureExpr)
 SWIFT_TYPEID_NAMED(ConstructorDecl *, ConstructorDecl)
 SWIFT_TYPEID_NAMED(CustomAttr *, CustomAttr)
 SWIFT_TYPEID_NAMED(Decl *, Decl)

include/swift/AST/ASTTypeIDs.h

@@ -23,10 +23,12 @@ namespace swift {
 
 class AbstractFunctionDecl;
 class BraceStmt;
+class ClosureExpr;
 class ConstructorDecl;
 class CustomAttr;
 class Decl;
 class EnumDecl;
+enum class FunctionBuilderClosurePreCheck : uint8_t;
 class GenericParamList;
 class GenericSignature;
 class GenericTypeParamType;

include/swift/AST/Expr.h

@@ -5407,6 +5407,9 @@ Expr *packSingleArgument(ASTContext &ctx, SourceLoc lParenLoc,
                               [](const Expr *E) -> Type {
                                 return E->getType();
                               });
+
+void simple_display(llvm::raw_ostream &out, const ClosureExpr *CE);
+
 } // end namespace swift
 
 #endif

include/swift/AST/TypeCheckRequests.h

@@ -1745,6 +1745,36 @@ class ValueWitnessRequest
   void cacheResult(Witness value) const;
 };
 
+enum class FunctionBuilderClosurePreCheck : uint8_t {
+  /// There were no problems pre-checking the closure.
+  Okay,
+
+  /// There was an error pre-checking the closure.
+  Error,
+
+  /// The closure has a return statement.
+  HasReturnStmt,
+};
+
+class PreCheckFunctionBuilderRequest
+    : public SimpleRequest<PreCheckFunctionBuilderRequest,
+                           FunctionBuilderClosurePreCheck(ClosureExpr *),
+                           CacheKind::Cached> {
+public:
+  using SimpleRequest::SimpleRequest;
+
+private:
+  friend SimpleRequest;
+
+  // Evaluation.
+  llvm::Expected<FunctionBuilderClosurePreCheck>
+  evaluate(Evaluator &evaluator, ClosureExpr *closure) const;
+
+public:
+  // Separate caching.
+  bool isCached() const { return true; }
+};
+
 // Allow AnyValue to compare two Type values, even though Type doesn't
 // support ==.
 template<>
@@ -1767,6 +1797,7 @@ AnyValue::Holder<GenericSignature>::equals(const HolderBase &other) const {
 void simple_display(llvm::raw_ostream &out, Type value);
 void simple_display(llvm::raw_ostream &out, const TypeRepr *TyR);
 void simple_display(llvm::raw_ostream &out, ImplicitMemberAction action);
+void simple_display(llvm::raw_ostream &out, FunctionBuilderClosurePreCheck pck);
 
 #define SWIFT_TYPEID_ZONE TypeChecker
 #define SWIFT_TYPEID_HEADER "swift/AST/TypeCheckerTypeIDZone.def"

include/swift/AST/TypeCheckerTypeIDZone.def

@@ -176,6 +176,9 @@ SWIFT_REQUEST(TypeChecker, HasUserDefinedDesignatedInitRequest,
               bool(NominalTypeDecl *), Cached, NoLocationInfo)
 SWIFT_REQUEST(TypeChecker, HasMemberwiseInitRequest,
               bool(StructDecl *), Cached, NoLocationInfo)
+SWIFT_REQUEST(TypeChecker, PreCheckFunctionBuilderRequest,
+              FunctionBuilderClosurePreCheck(ClosureExpr *),
+              Cached, NoLocationInfo)
 SWIFT_REQUEST(TypeChecker, ResolveImplicitMemberRequest,
               bool(NominalTypeDecl *, ImplicitMemberAction), Uncached,
               NoLocationInfo)

lib/AST/Expr.cpp

@@ -2246,6 +2246,19 @@ SourceLoc TapExpr::getEndLoc() const {
   return SourceLoc();
 }
 
+void swift::simple_display(llvm::raw_ostream &out, const ClosureExpr *CE) {
+  if (!CE) {
+    out << "(null)";
+    return;
+  }
+
+  if (CE->hasSingleExpressionBody()) {
+    out << "single expression closure";
+  } else {
+    out << "closure";
+  }
+}
+
 // See swift/Basic/Statistic.h for declaration: this enables tracing Exprs, is
 // defined here to avoid too much layering violation / circular linkage
 // dependency.

lib/AST/TypeCheckRequests.cpp

@@ -1113,3 +1113,22 @@ Optional<Witness> ValueWitnessRequest::getCachedResult() const {
 void ValueWitnessRequest::cacheResult(Witness type) const {
   // FIXME: Refactor this to be the thing that warms the cache.
 }
+
+//----------------------------------------------------------------------------//
+// PreCheckFunctionBuilderRequest computation.
+//----------------------------------------------------------------------------//
+
+void swift::simple_display(llvm::raw_ostream &out,
+                           FunctionBuilderClosurePreCheck value) {
+  switch (value) {
+  case FunctionBuilderClosurePreCheck::Okay:
+    out << "okay";
+    break;
+  case FunctionBuilderClosurePreCheck::HasReturnStmt:
+    out << "has return statement";
+    break;
+  case FunctionBuilderClosurePreCheck::Error:
+    out << "error";
+    break;
+  }
+}

lib/Sema/BuilderTransform.cpp

@@ -22,6 +22,7 @@
 #include "swift/AST/NameLookup.h"
 #include "swift/AST/NameLookupRequests.h"
 #include "swift/AST/ParameterList.h"
+#include "swift/AST/TypeCheckRequests.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/SmallVector.h"
 #include <iterator>
@@ -516,7 +517,9 @@ ConstraintSystem::TypeMatchResult ConstraintSystem::applyFunctionBuilder(
 
   // Pre-check the closure body: pre-check any expressions in it and look
   // for return statements.
-  switch (getTypeChecker().preCheckFunctionBuilderClosureBody(closure)) {
+  auto request = PreCheckFunctionBuilderRequest{closure};
+  switch (evaluateOrDefault(getASTContext().evaluator, request,
+                            FunctionBuilderClosurePreCheck::Error)) {
   case FunctionBuilderClosurePreCheck::Okay:
     // If the pre-check was okay, apply the function-builder transform.
     break;
@@ -670,21 +673,12 @@ class PreCheckFunctionBuilderClosure : public ASTWalker {
 
 }
 
-FunctionBuilderClosurePreCheck
-TypeChecker::preCheckFunctionBuilderClosureBody(ClosureExpr *closure) {
+llvm::Expected<FunctionBuilderClosurePreCheck>
+PreCheckFunctionBuilderRequest::evaluate(Evaluator &eval,
+                                         ClosureExpr *closure) const {
   // Single-expression closures should already have been pre-checked.
   if (closure->hasSingleExpressionBody())
     return FunctionBuilderClosurePreCheck::Okay;
 
-  // Check whether we've already done this analysis.
-  auto it = precheckedFunctionBuilderClosures.find(closure);
-  if (it != precheckedFunctionBuilderClosures.end())
-    return it->second;
-
-  auto result = PreCheckFunctionBuilderClosure(closure).run();
-
-  // Cache the result.
-  precheckedFunctionBuilderClosures.insert(std::make_pair(closure, result));
-
-  return result;
+  return PreCheckFunctionBuilderClosure(closure).run();
 }

lib/Sema/TypeChecker.h

@@ -530,17 +530,6 @@ enum class CheckedCastContextKind {
   EnumElementPattern,
 };
 
-enum class FunctionBuilderClosurePreCheck : uint8_t {
-  /// There were no problems pre-checking the closure.
-  Okay,
-
-  /// There was an error pre-checking the closure.
-  Error,
-
-  /// The closure has a return statement.
-  HasReturnStmt,
-};
-
 /// The Swift type checker, which takes a parsed AST and performs name binding,
 /// type checking, and semantic analysis to produce a type-annotated AST.
 class TypeChecker final {
@@ -600,11 +589,6 @@ class TypeChecker final {
   /// function bodies.
   bool SkipNonInlinableFunctionBodies = false;
 
-  /// Closure expressions whose bodies have already been prechecked as
-  /// part of trying to apply a function builder.
-  llvm::DenseMap<ClosureExpr *, FunctionBuilderClosurePreCheck>
-    precheckedFunctionBuilderClosures;
-
   TypeChecker(ASTContext &Ctx);
   friend class ASTContext;
   friend class constraints::ConstraintSystem;
@@ -1075,16 +1059,6 @@ class TypeChecker final {
   /// struct or class.
   static void addImplicitConstructors(NominalTypeDecl *typeDecl);
 
-  /// Pre-check the body of the given closure, which we are about to
-  /// generate constraints for.
-  ///
-  /// This mutates the body of the closure, but only in ways that should be
-  /// valid even if we end up not actually applying the function-builder
-  /// transform: it just does a normal pre-check of all the expressions in
-  /// the closure.
-  FunctionBuilderClosurePreCheck
-  preCheckFunctionBuilderClosureBody(ClosureExpr *closure);
-
   /// \name Name lookup
   ///
   /// Routines that perform name lookup.