Parameterized extensions

include/swift/AST/ASTDemangler.h

@@ -153,7 +153,7 @@ class ASTBuilder {
   bool validateParentType(TypeDecl *decl, Type parent);
   CanGenericSignature demangleGenericSignature(
       NominalTypeDecl *nominalDecl,
-      NodePointer node);
+      NodePointer node, bool isParameterizedExtension);
   DeclContext *findDeclContext(NodePointer node);
   ModuleDecl *findModule(NodePointer node);
   Demangle::NodePointer findModuleNode(NodePointer node);

include/swift/AST/Decl.h

@@ -1216,6 +1216,7 @@ class ExtensionDecl final : public GenericContext, public Decl,
 
   ExtensionDecl(SourceLoc extensionLoc, TypeRepr *extendedType,
                 MutableArrayRef<TypeLoc> inherited,
+                GenericParamList *genericParams,
                 DeclContext *parent,
                 TrailingWhereClause *trailingWhereClause);
 
@@ -1241,6 +1242,7 @@ class ExtensionDecl final : public GenericContext, public Decl,
   static ExtensionDecl *create(ASTContext &ctx, SourceLoc extensionLoc,
                                TypeRepr *extendedType,
                                MutableArrayRef<TypeLoc> inherited,
+                               GenericParamList *genericParams,
                                DeclContext *parent,
                                TrailingWhereClause *trailingWhereClause,
                                ClangNode clangNode = ClangNode());
@@ -1345,6 +1347,10 @@ class ExtensionDecl final : public GenericContext, public Decl,
   /// resiliently moved into the original protocol itself.
   bool isEquivalentToExtendedContext() const;
 
+  /// Whether this extension contains unique generic parameters that differ from
+  /// the extended nominal's generic signature.
+  bool isParameterized() const;
+
   // Implement isa/cast/dyncast/etc.
   static bool classof(const Decl *D) {
     return D->getKind() == DeclKind::Extension;

include/swift/AST/DeclContext.h

@@ -448,6 +448,15 @@ class alignas(1 << DeclContextAlignInBits) DeclContext {
         const_cast<DeclContext *>(this)->getInnermostSkippedFunctionContext();
   }
 
+  /// Returns the outermost context based off syntactic depth, right below the
+  /// module context. E.g. for a nested type in a nested type in an extension,
+  /// this returns that extension.
+  LLVM_READONLY
+  DeclContext *getOutermostSyntacticContext();
+  const DeclContext *getOutermostSyntacticContext() const {
+    return const_cast<DeclContext *>(this)->getOutermostSyntacticContext();
+  }
+
   /// Returns the semantic parent of this context.  A context has a
   /// parent if and only if it is not a module context.
   DeclContext *getParent() const {

include/swift/AST/DiagnosticsSema.def

@@ -1787,6 +1787,10 @@ ERROR(inferred_opaque_type,none,
 // Extensions
 ERROR(non_nominal_extension,none,
       "non-nominal type %0 cannot be extended", (Type))
+ERROR(generic_param_extension,none,
+      "cannot extend generic parameter type %0", (Identifier))
+ERROR(concrete_generic_extension,none,
+      "cannot have generic parameters when extending a concrete type", ())
 WARNING(composition_in_extended_type,none,
         "extending a protocol composition is not supported; extending %0 "
         "instead", (Type))
@@ -1797,9 +1801,6 @@ ERROR(extension_access_with_conformances,none,
       "protocol conformances", (DeclAttribute))
 ERROR(extension_metatype,none,
       "cannot extend a metatype %0", (Type))
-ERROR(extension_specialization,none,
-      "constrained extension must be declared on the unspecialized generic "
-      "type %0 with constraints specified by a 'where' clause", (Identifier))
 ERROR(extension_stored_property,none,
       "extensions must not contain stored properties", ())
 NOTE(extension_stored_property_fixit,none,

include/swift/Demangling/DemangleNodes.def

@@ -297,6 +297,7 @@ NODE(NoncanonicalSpecializedGenericTypeMetadataCache)
 NODE(GlobalVariableOnceFunction)
 NODE(GlobalVariableOnceToken)
 NODE(GlobalVariableOnceDeclList)
+CONTEXT_NODE(GenericExtension)
 
 #undef CONTEXT_NODE
 #undef NODE

include/swift/Demangling/Demangler.h

@@ -503,7 +503,7 @@ class Demangler : public NodeFactory {
   NodePointer popTypeAndGetAnyGeneric();
   NodePointer demangleBuiltinType();
   NodePointer demangleAnyGenericType(Node::Kind kind);
-  NodePointer demangleExtensionContext();
+  NodePointer demangleExtensionContext(bool isParameterized);
   NodePointer demanglePlainFunction();
   NodePointer popFunctionType(Node::Kind kind);
   NodePointer popFunctionParams(Node::Kind kind);

include/swift/Demangling/TypeDecoder.h

@@ -1078,6 +1078,7 @@ class TypeDecoder {
       case Node::Kind::Module:
         break;
       case Node::Kind::Extension:
+      case Node::Kind::GenericExtension:
         // Decode the type being extended.
         if (parentContext->getNumChildren() < 2)
           return MAKE_NODE_TYPE_ERROR(parentContext,

lib/AST/ASTDemangler.cpp

@@ -840,13 +840,33 @@ ASTBuilder::getForeignModuleKind(NodePointer node) {
 
 CanGenericSignature ASTBuilder::demangleGenericSignature(
     NominalTypeDecl *nominalDecl,
-    NodePointer node) {
+    NodePointer node,
+    bool isParameterizedExtension) {
+  SmallVector<GenericTypeParamType *, 2> gpTypes;
   SmallVector<Requirement, 2> requirements;
 
   for (auto &child : *node) {
-    if (child->getKind() ==
-          Demangle::Node::Kind::DependentGenericParamCount)
+    if (child->getKind() == Demangle::Node::Kind::DependentGenericParamCount) {
+      // If this is not a parameterized extension, just ignore this node.
+      if (!isParameterizedExtension) {
+        continue;
+      }
+
+      // However, if it is a parameterized extension, this node tells us exactly
+      // how many generic parameters the extension introduced. Create that many
+      // generic parameter types for the generic signature.
+
+      // The extension's generic depth is 1 higher than the context it's
+      // extending.
+      unsigned depth = nominalDecl->getGenericContextDepth() + 1;
+
+      for (uint64_t i = 0; i != child->getIndex(); i += 1) {
+        auto gpTy = GenericTypeParamType::get(depth, i, Ctx);
+        gpTypes.push_back(gpTy);
+      }
+
       continue;
+    }
 
     if (child->getNumChildren() != 2)
       return CanGenericSignature();
@@ -930,7 +950,7 @@ CanGenericSignature ASTBuilder::demangleGenericSignature(
   return evaluateOrDefault(Ctx.evaluator,
                            AbstractGenericSignatureRequest{
                                nominalDecl->getGenericSignature().getPointer(),
-                               {},
+                               std::move(gpTypes),
                                std::move(requirements)},
                            GenericSignature())
       .getCanonicalSignature();
@@ -1018,7 +1038,8 @@ ASTBuilder::findDeclContext(NodePointer node) {
   case Demangle::Node::Kind::Global:
     return findDeclContext(node->getChild(0));
 
-  case Demangle::Node::Kind::Extension: {
+  case Demangle::Node::Kind::Extension:
+  case Demangle::Node::Kind::GenericExtension: {
     auto *moduleDecl = dyn_cast_or_null<ModuleDecl>(
         findDeclContext(node->getChild(0)));
     if (!moduleDecl)
@@ -1029,9 +1050,14 @@ ASTBuilder::findDeclContext(NodePointer node) {
     if (!nominalDecl)
       return nullptr;
 
+    bool isParameterized = false;
+    if (node->getKind() == Demangle::Node::Kind::GenericExtension)
+      isParameterized = true;
+
     CanGenericSignature genericSig;
     if (node->getNumChildren() > 2)
-      genericSig = demangleGenericSignature(nominalDecl, node->getChild(2));
+      genericSig = demangleGenericSignature(nominalDecl, node->getChild(2),
+                                            isParameterized);
 
     for (auto *ext : nominalDecl->getExtensions()) {
       if (ext->getParentModule() != moduleDecl)

lib/AST/ASTDumper.cpp

@@ -630,6 +630,9 @@ namespace {
 
     void visitExtensionDecl(ExtensionDecl *ED) {
       printCommon(ED, "extension_decl", ExtensionColor);
+      if (ED->isParameterized()) {
+        printGenericParameters(OS, ED->getParsedGenericParams());
+      }
       OS << ' ';
       ED->getExtendedType().print(OS);
       printCommonPost(ED);

lib/AST/ASTMangler.cpp

@@ -1920,13 +1920,13 @@ void ASTMangler::appendContext(const DeclContext *ctx, StringRef useModuleName)
     return;
 
   case DeclContextKind::ExtensionDecl: {
-    auto ExtD = cast<ExtensionDecl>(ctx);
-    auto decl = ExtD->getExtendedNominal();
+    auto ext = cast<ExtensionDecl>(ctx);
+    auto decl = ext->getExtendedNominal();
     // Recover from erroneous extension.
     if (!decl)
-      return appendContext(ExtD->getDeclContext(), useModuleName);
+      return appendContext(ext->getDeclContext(), useModuleName);
 
-    if (!ExtD->isEquivalentToExtendedContext()) {
+    if (!ext->isEquivalentToExtendedContext()) {
     // Mangle the extension if:
     // - the extension is defined in a different module from the original
     //   nominal type decl,
@@ -1936,17 +1936,22 @@ void ASTMangler::appendContext(const DeclContext *ctx, StringRef useModuleName)
     // "extension is to a protocol" would no longer be a reason to use the
     // extension mangling, because an extension method implementation could be
     // resiliently moved into the original protocol itself.
-      auto sig = ExtD->getGenericSignature();
+      auto sig = ext->getGenericSignature();
       // If the extension is constrained, mangle the generic signature that
       // constrains it.
       appendAnyGenericType(decl);
-      appendModule(ExtD->getParentModule(), useModuleName);
-      if (sig && ExtD->isConstrainedExtension()) {
-        Mod = ExtD->getModuleContext();
-        auto nominalSig = ExtD->getSelfNominalTypeDecl()
+      appendModule(ext->getParentModule(), useModuleName);
+      if (sig && ext->isConstrainedExtension()) {
+        Mod = ext->getModuleContext();
+        auto nominalSig = ext->getSelfNominalTypeDecl()
                             ->getGenericSignatureOfContext();
         appendGenericSignature(sig, nominalSig);
       }
+
+      if (ext->isParameterized()) {
+        return appendOperator("J");
+      }
+
       return appendOperator("E");
     }
     return appendAnyGenericType(decl);

lib/AST/ASTPrinter.cpp

@@ -2315,9 +2315,15 @@ void PrintAST::printExtension(ExtensionDecl *decl) {
   if (Options.BracketOptions.shouldOpenExtension(decl)) {
     printDocumentationComment(decl);
     printAttributes(decl);
-    Printer << "extension ";
+    Printer << tok::kw_extension;
+
+    if (decl->isParameterized()) {
+      printGenericDeclGenericParams(decl);
+    }
+
+    Printer << " ";
+
     recordDeclLoc(decl, [&]{
-      // We cannot extend sugared types.
       Type extendedType = decl->getExtendedType();
       if (!extendedType) {
         // Fallback to TypeRepr.

lib/AST/Decl.cpp

@@ -1103,9 +1103,10 @@ NominalTypeDecl::takeConformanceLoaderSlow() {
 ExtensionDecl::ExtensionDecl(SourceLoc extensionLoc,
                              TypeRepr *extendedType,
                              MutableArrayRef<TypeLoc> inherited,
+                             GenericParamList *genericParams,
                              DeclContext *parent,
                              TrailingWhereClause *trailingWhereClause)
-  : GenericContext(DeclContextKind::ExtensionDecl, parent, nullptr),
+  : GenericContext(DeclContextKind::ExtensionDecl, parent, genericParams),
     Decl(DeclKind::Extension, parent),
     IterableDeclContext(IterableDeclContextKind::ExtensionDecl),
     ExtensionLoc(extensionLoc),
@@ -1120,6 +1121,7 @@ ExtensionDecl::ExtensionDecl(SourceLoc extensionLoc,
 ExtensionDecl *ExtensionDecl::create(ASTContext &ctx, SourceLoc extensionLoc,
                                      TypeRepr *extendedType,
                                      MutableArrayRef<TypeLoc> inherited,
+                                     GenericParamList *genericParams,
                                      DeclContext *parent,
                                      TrailingWhereClause *trailingWhereClause,
                                      ClangNode clangNode) {
@@ -1130,8 +1132,8 @@ ExtensionDecl *ExtensionDecl::create(ASTContext &ctx, SourceLoc extensionLoc,
 
   // Construct the extension.
   auto result = ::new (declPtr) ExtensionDecl(extensionLoc, extendedType,
-                                              inherited, parent,
-                                              trailingWhereClause);
+                                              inherited, genericParams,
+                                              parent, trailingWhereClause);
   if (clangNode)
     result->setClangNode(clangNode);
 
@@ -1229,6 +1231,20 @@ bool ExtensionDecl::isEquivalentToExtendedContext() const {
     && !getDeclaredInterfaceType()->isExistentialType();
 }
 
+bool ExtensionDecl::isParameterized() const {
+  if (!hasBeenBound())
+    return getParsedGenericParams();
+
+  auto nominal = getExtendedNominal();
+  if (!nominal && getParsedGenericParams())
+    return true;
+
+  if (!nominal)
+    return false;
+
+  return getGenericContextDepth() != nominal->getGenericContextDepth();
+}
+
 AccessLevel ExtensionDecl::getDefaultAccessLevel() const {
   ASTContext &ctx = getASTContext();
   return evaluateOrDefault(ctx.evaluator,

lib/AST/DeclContext.cpp

@@ -96,7 +96,12 @@ GenericTypeParamType *DeclContext::getProtocolSelfType() const {
   if (auto proto = dyn_cast<ProtocolDecl>(this)) {
     genericParams = proto->getGenericParams();
   } else {
-    genericParams = cast<ExtensionDecl>(this)->getGenericParams();
+    auto ext = cast<ExtensionDecl>(this);
+    genericParams = ext->getGenericParams();
+
+    if (ext->isParameterized()) {
+      genericParams = genericParams->getOuterParameters();
+    }
   }
 
   if (genericParams == nullptr)
@@ -242,6 +247,17 @@ DeclContext *DeclContext::getInnermostSkippedFunctionContext() {
   return nullptr;
 }
 
+DeclContext *DeclContext::getOutermostSyntacticContext() {
+  auto depth = getSyntacticDepth() - 1;
+  auto dc = this;
+
+  for (auto i = 0; i != depth; i += 1) {
+    dc = dc->getParent();
+  }
+
+  return dc;
+}
+
 DeclContext *DeclContext::getParentForLookup() const {
   if (isa<ProtocolDecl>(this) || isa<ExtensionDecl>(this)) {
     // If we are inside a protocol or an extension, skip directly

lib/AST/NameLookup.cpp

@@ -2382,6 +2382,13 @@ createExtensionGenericParams(ASTContext &ctx,
                              NominalTypeDecl *nominal) {
   // Collect generic parameters from all outer contexts.
   SmallVector<GenericParamList *, 2> allGenericParams;
+
+  // If this is a parameterized extension, set those generic params as innermost
+  // parameters.
+  if (auto gpList = ext->getParsedGenericParams()) {
+    allGenericParams.push_back(gpList);
+  }
+
   nominal->forEachGenericContext([&](GenericParamList *gpList) {
     allGenericParams.push_back(gpList->clone(ext));
   });
@@ -2405,19 +2412,36 @@ GenericParamListRequest::evaluate(Evaluator &evaluator, GenericContext *value) c
     if (!nominal) {
       return nullptr;
     }
+
+    bool isParameterized = false;
+    if (ext->getParsedGenericParams()) {
+      isParameterized = true;
+    }
+
     auto *genericParams = createExtensionGenericParams(ctx, ext, nominal);
 
     // Protocol extensions need an inheritance clause due to how name lookup
     // is implemented.
     if (auto *proto = ext->getExtendedProtocolDecl()) {
       auto protoType = proto->getDeclaredInterfaceType();
       TypeLoc selfInherited[1] = { TypeLoc::withoutLoc(protoType) };
-      genericParams->getParams().front()->setInherited(
+      auto protoParams = genericParams;
+
+      if (isParameterized) {
+        protoParams = protoParams->getOuterParameters();
+      }
+
+      protoParams->getParams().front()->setInherited(
         ctx.AllocateCopy(selfInherited));
     }
 
     // Set the depth of every generic parameter.
     unsigned depth = nominal->getGenericContextDepth();
+
+    if (isParameterized) {
+      depth += 1;
+    }
+
     for (auto *outerParams = genericParams;
          outerParams != nullptr;
          outerParams = outerParams->getOuterParameters())

lib/AST/NameLookupRequests.cpp

@@ -212,6 +212,13 @@ void GetDestructorRequest::cacheResult(DestructorDecl *value) const {
 
 Optional<GenericParamList *> GenericParamListRequest::getCachedResult() const {
   auto *decl = std::get<0>(getStorage());
+
+  // If this is an extension and we can get the parsed generic params, go ahead
+  // and go through the request evaluation because we still need to clone the
+  // nominal's generic param list.
+  if (isa<ExtensionDecl>(decl) && decl->getParsedGenericParams())
+    return None;
+
   if (auto *params = decl->GenericParamsAndBit.getPointer())
     return params;