Requestify Extension Type Validation

Author: CodaFi

include/swift/AST/TypeCheckRequests.h

@@ -1104,6 +1104,40 @@ class AbstractGenericSignatureRequest :
   }
 };
 
+class ExtendedTypeRequest
+    : public SimpleRequest<ExtendedTypeRequest,
+                           Type(ExtensionDecl *),
+                           CacheKind::SeparatelyCached> {
+public:
+  using SimpleRequest::SimpleRequest;
+
+private:
+  friend SimpleRequest;
+
+  // Evaluation.
+  llvm::Expected<GenericSignature *>
+  evaluate(Evaluator &evaluator,
+           GenericSignature *baseSignature,
+           SmallVector<GenericTypeParamType *, 2> addedParameters,
+           SmallVector<Requirement, 2> addedRequirements) const;
+
+public:
+  // Separate caching.
+  bool isCached() const;
+
+  /// Abstract generic signature requests never have source-location info.
+  SourceLoc getNearestLoc() const {
+    return SourceLoc();
+  }
+  llvm::Expected<Type> evaluate(Evaluator &eval, ExtensionDecl *) const;
+
+public:
+  // Separate caching
+  bool isCached() const { return true; }
+  Optional<Type> getCachedResult() const;
+  void cacheResult(Type value) const;
+};
+
 // Allow AnyValue to compare two Type values, even though Type doesn't
 // support ==.
 template<>

include/swift/AST/TypeCheckerTypeIDZone.def

@@ -58,3 +58,4 @@ SWIFT_TYPEID(EmittedMembersRequest)
 SWIFT_TYPEID(IsImplicitlyUnwrappedOptionalRequest)
 SWIFT_TYPEID(ClassAncestryFlagsRequest)
 SWIFT_TYPEID(AbstractGenericSignatureRequest)
+SWIFT_TYPEID(ExtendedTypeRequest)

lib/AST/TypeCheckRequests.cpp

@@ -824,3 +824,20 @@ void IsImplicitlyUnwrappedOptionalRequest::cacheResult(bool value) const {
   auto *decl = std::get<0>(getStorage());
   decl->setImplicitlyUnwrappedOptional(value);
 }
+
+//----------------------------------------------------------------------------//
+// ExtendedTypeRequest computation.
+//----------------------------------------------------------------------------//
+
+Optional<Type> ExtendedTypeRequest::getCachedResult() const {
+  auto *ext = std::get<0>(getStorage());
+  if (auto type = ext->getExtendedTypeLoc().getType()) {
+    return type;
+  }
+  return None;
+}
+
+void ExtendedTypeRequest::cacheResult(Type value) const {
+  auto *ext = std::get<0>(getStorage());
+  ext->getExtendedTypeLoc().setType(value);
+}

lib/Sema/TypeCheckDecl.cpp

@@ -4414,7 +4414,8 @@ static bool isNonGenericTypeAliasType(Type type) {
   return false;
 }
 
-static Type validateExtendedType(ExtensionDecl *ext) {
+llvm::Expected<Type>
+ExtendedTypeRequest::evaluate(Evaluator &eval, ExtensionDecl *ext) const {
   auto error = [&ext]() {
     ext->setInvalid();
     return ErrorType::get(ext->getASTContext());
@@ -4430,7 +4431,6 @@ static Type validateExtendedType(ExtensionDecl *ext) {
   auto tr = TypeResolution::forStructural(ext->getDeclContext());
   auto extendedType = tr.resolveType(ext->getExtendedTypeLoc().getTypeRepr(),
                                      options);
-  ext->getExtendedTypeLoc().setType(extendedType);
 
   if (extendedType->hasError())
     return error();
@@ -4483,7 +4483,9 @@ void TypeChecker::validateExtension(ExtensionDecl *ext) {
 
   DeclValidationRAII IBV(ext);
 
-  auto extendedType = validateExtendedType(ext);
+  auto extendedType = evaluateOrDefault(Context.evaluator,
+                                        ExtendedTypeRequest{ext},
+                                        ErrorType::get(ext->getASTContext()));
 
   if (auto *nominal = ext->getExtendedNominal()) {
     // If this extension was not already bound, it means it is either in an
@@ -4496,8 +4498,6 @@ void TypeChecker::validateExtension(ExtensionDecl *ext) {
     // Validate the nominal type declaration being extended.
     validateDecl(nominal);
 
-    ext->getExtendedTypeLoc().setType(extendedType);
-
     if (auto *genericParams = ext->getGenericParams()) {
       GenericEnvironment *env =
         checkExtensionGenericParams(*this, ext, extendedType, genericParams);

lib/Serialization/Deserialization.cpp

@@ -3811,7 +3811,8 @@ class swift::DeclDeserializer {
     MF.configureGenericEnvironment(extension, genericEnvID);
 
     auto baseTy = MF.getType(baseID);
-    extension->getExtendedTypeLoc().setType(baseTy);
+    ctx.evaluator.cacheOutput(ExtendedTypeRequest{extension},
+                              std::move(baseTy));
     auto nominal = extension->getExtendedNominal();
 
     if (isImplicit)