AST and Serialization support for generic subscripts

include/swift/AST/ASTContext.h

@@ -58,13 +58,13 @@ namespace swift {
   class ExtensionDecl;
   class ForeignRepresentationInfo;
   class FuncDecl;
+  class GenericContext;
   class InFlightDiagnostic;
   class IterableDeclContext;
-  class LazyAbstractFunctionData;
-  class LazyGenericTypeData;
   class LazyContextData;
-  class LazyMemberLoader;
+  class LazyGenericContextData;
   class LazyIterableDeclContextData;
+  class LazyMemberLoader;
   class LazyResolver;
   class PatternBindingDecl;
   class PatternBindingInitializer;
@@ -722,25 +722,16 @@ class ASTContext {
   /// \param lazyLoader If non-null, the lazy loader to use when creating the
   /// lazy data. The pointer must either be null or be consistent
   /// across all calls for the same \p func.
-  LazyContextData *getOrCreateLazyContextData(const Decl *decl,
+  LazyContextData *getOrCreateLazyContextData(const DeclContext *decl,
                                               LazyMemberLoader *lazyLoader);
 
-  /// Get the lazy function data for the given generic type.
-  ///
-  /// \param lazyLoader If non-null, the lazy loader to use when creating the
-  /// generic type data. The pointer must either be null or be consistent
-  /// across all calls for the same \p type.
-  LazyGenericTypeData *getOrCreateLazyGenericTypeData(
-                                                  const GenericTypeDecl *type,
-                                                  LazyMemberLoader *lazyLoader);
-
-  /// Get the lazy function data for the given abstract function.
+  /// Get the lazy function data for the given generic context.
   ///
   /// \param lazyLoader If non-null, the lazy loader to use when creating the
   /// function data. The pointer must either be null or be consistent
   /// across all calls for the same \p func.
-  LazyAbstractFunctionData *getOrCreateLazyFunctionContextData(
-                                              const AbstractFunctionDecl *func,
+  LazyGenericContextData *getOrCreateLazyGenericContextData(
+                                              const GenericContext *dc,
                                               LazyMemberLoader *lazyLoader);
 
   /// Get the lazy iterable context for the given iterable declaration context.

include/swift/AST/Decl.h

@@ -4527,18 +4527,20 @@ enum class ObjCSubscriptKind {
 /// A given type can have multiple subscript declarations, so long as the
 /// signatures (indices and element type) are distinct.
 ///
-class SubscriptDecl : public AbstractStorageDecl, public DeclContext {
+class SubscriptDecl : public AbstractStorageDecl, public GenericContext {
   SourceLoc ArrowLoc;
   ParameterList *Indices;
   TypeLoc ElementTy;
 
 public:
   SubscriptDecl(DeclName Name, SourceLoc SubscriptLoc, ParameterList *Indices,
-                SourceLoc ArrowLoc, TypeLoc ElementTy, DeclContext *Parent)
+                SourceLoc ArrowLoc, TypeLoc ElementTy, DeclContext *Parent,
+                GenericParamList *GenericParams)
     : AbstractStorageDecl(DeclKind::Subscript, Parent, Name, SubscriptLoc),
-      DeclContext(DeclContextKind::SubscriptDecl, Parent),
+      GenericContext(DeclContextKind::SubscriptDecl, Parent),
       ArrowLoc(ArrowLoc), Indices(nullptr), ElementTy(ElementTy) {
     setIndices(Indices);
+    setGenericParams(GenericParams);
   }
 
   SourceLoc getSubscriptLoc() const { return getNameLoc(); }

include/swift/AST/LazyResolver.h

@@ -191,22 +191,15 @@ class LazyContextData {
   LazyMemberLoader *loader;
 };
 
-/// Context data for abstract function declarations.
-class LazyAbstractFunctionData : public LazyContextData {
-public:
-  /// The context data used for loading the generic environment.
-  uint64_t genericEnvData = 0;
-};
-
-/// Context data for generic type declarations.
-class LazyGenericTypeData : public LazyContextData {
+/// Context data for generic contexts.
+class LazyGenericContextData : public LazyContextData {
 public:
   /// The context data used for loading the generic environment.
   uint64_t genericEnvData = 0;
 };
 
 /// Context data for iterable decl contexts.
-class LazyIterableDeclContextData : public LazyGenericTypeData {
+class LazyIterableDeclContextData : public LazyGenericContextData {
 public:
   /// The context data used for loading all of the members of the iterable
   /// context.
@@ -254,7 +247,7 @@ class alignas(void*) LazyMemberLoader {
   }
 
   /// Returns the generic environment.
-  virtual GenericEnvironment *loadGenericEnvironment(const Decl *decl,
+  virtual GenericEnvironment *loadGenericEnvironment(const DeclContext *decl,
                                                      uint64_t contextData) {
     llvm_unreachable("unimplemented");
   }

include/swift/Serialization/ModuleFile.h

@@ -507,8 +507,7 @@ class ModuleFile : public LazyMemberLoader {
   /// Set up a (potentially lazy) generic environment for the given type,
   /// function or extension.
   void configureGenericEnvironment(
-                   llvm::PointerUnion3<GenericTypeDecl *, ExtensionDecl *,
-                                       AbstractFunctionDecl *> genericDecl,
+                   GenericContext *genericDecl,
                    serialization::GenericEnvironmentID envID);
 
   /// Populates the vector with members of a DeclContext from \c DeclTypeCursor.
@@ -711,7 +710,7 @@ class ModuleFile : public LazyMemberLoader {
   virtual void finishNormalConformance(NormalProtocolConformance *conformance,
                                        uint64_t contextData) override;
 
-  GenericEnvironment *loadGenericEnvironment(const Decl *decl,
+  GenericEnvironment *loadGenericEnvironment(const DeclContext *decl,
                                              uint64_t contextData) override;
 
   Optional<StringRef> getGroupNameById(unsigned Id) const;

include/swift/Serialization/ModuleFormat.h

@@ -54,7 +54,7 @@ const uint16_t VERSION_MAJOR = 0;
 /// in source control, you should also update the comment to briefly
 /// describe what change you made. The content of this comment isn't important;
 /// it just ensures a conflict if two people change the module format.
-const uint16_t VERSION_MINOR = 318; // Last change: SIL open_exist. access kind
+const uint16_t VERSION_MINOR = 319; // Last change: generic subscripts
 
 using DeclID = PointerEmbeddedInt<unsigned, 31>;
 using DeclIDField = BCFixed<31>;
@@ -976,6 +976,7 @@ namespace decls_block {
     BCFixed<1>,  // implicit?
     BCFixed<1>,  // objc?
     StorageKindField,   // StorageKind
+    GenericEnvironmentIDField, // generic environment
     TypeIDField, // interface type
     DeclIDField, // getter
     DeclIDField, // setter
@@ -988,7 +989,9 @@ namespace decls_block {
     AccessibilityKindField, // accessibility
     AccessibilityKindField, // setter accessibility, if applicable
     BCArray<IdentifierIDField> // name components
-    // The indices pattern trails the record.
+    // Trailed by:
+    // - generic parameters, if any
+    // - the indices pattern
   >;
 
   using ExtensionLayout = BCRecordLayout<

lib/AST/ASTContext.cpp

@@ -206,7 +206,7 @@ struct ASTContext::Implementation {
     DelayedConformanceDiags;
 
   /// Stores information about lazy deserialization of various declarations.
-  llvm::DenseMap<const Decl *, LazyContextData *> LazyContexts;
+  llvm::DenseMap<const DeclContext *, LazyContextData *> LazyContexts;
 
   /// Stored generic signature builders for canonical generic signatures.
   llvm::DenseMap<std::pair<GenericSignature *, ModuleDecl *>,
@@ -1477,9 +1477,9 @@ ASTContext::getInheritedConformance(Type type, ProtocolConformance *inherited) {
 }
 
 LazyContextData *ASTContext::getOrCreateLazyContextData(
-                                                const Decl *decl,
+                                                const DeclContext *dc,
                                                 LazyMemberLoader *lazyLoader) {
-  auto known = Impl.LazyContexts.find(decl);
+  auto known = Impl.LazyContexts.find(dc);
   if (known != Impl.LazyContexts.end()) {
     // Make sure we didn't provide an incompatible lazy loader.
     assert(!lazyLoader || lazyLoader == known->second->loader);
@@ -1488,30 +1488,18 @@ LazyContextData *ASTContext::getOrCreateLazyContextData(
 
   // Create new lazy iterable context data with the given loader.
   assert(lazyLoader && "Queried lazy data for non-lazy iterable context");
-  if (isa<NominalTypeDecl>(decl) || isa<ExtensionDecl>(decl)) {
+  if (isa<NominalTypeDecl>(dc) || isa<ExtensionDecl>(dc)) {
     auto *contextData = Allocate<LazyIterableDeclContextData>();
     contextData->loader = lazyLoader;
-    Impl.LazyContexts[decl] = contextData;
+    Impl.LazyContexts[dc] = contextData;
     return contextData;
   }
 
-  // Create new lazy generic type data with the given loader.
-  if (isa<GenericTypeDecl>(decl)) {
-    auto *contextData = Allocate<LazyGenericTypeData>();
-    contextData->loader = lazyLoader;
-    Impl.LazyContexts[decl] = contextData;
-    return contextData;
-  }
-
-  // Create new lazy function context data with the given loader.
-  if (isa<AbstractFunctionDecl>(decl)) {
-    auto *contextData = Allocate<LazyAbstractFunctionData>();
-    contextData->loader = lazyLoader;
-    Impl.LazyContexts[decl] = contextData;
-    return contextData;
-  }
-
-  llvm_unreachable("unhandled lazy context");
+  // Create new lazy generic context data with the given loader.
+  auto *contextData = Allocate<LazyGenericContextData>();
+  contextData->loader = lazyLoader;
+  Impl.LazyContexts[dc] = contextData;
+  return contextData;
 }
 
 LazyIterableDeclContextData *ASTContext::getOrCreateLazyIterableContextData(
@@ -1527,18 +1515,13 @@ LazyIterableDeclContextData *ASTContext::getOrCreateLazyIterableContextData(
                                                                    lazyLoader);
 }
 
-LazyAbstractFunctionData *ASTContext::getOrCreateLazyFunctionContextData(
-                                               const AbstractFunctionDecl *func,
+LazyGenericContextData *ASTContext::getOrCreateLazyGenericContextData(
+                                               const GenericContext *dc,
                                                LazyMemberLoader *lazyLoader) {
-  return (LazyAbstractFunctionData *)getOrCreateLazyContextData(func,
-                                                                lazyLoader);
+  return (LazyGenericContextData *)getOrCreateLazyContextData(dc,
+                                                              lazyLoader);
 }
 
-LazyGenericTypeData *ASTContext::getOrCreateLazyGenericTypeData(
-                                               const GenericTypeDecl *type,
-                                               LazyMemberLoader *lazyLoader) {
-  return (LazyGenericTypeData *)getOrCreateLazyContextData(type, lazyLoader);
-}
 void ASTContext::addDelayedConformanceDiag(
        NormalProtocolConformance *conformance,
        DelayedConformanceDiag fn) {

lib/AST/Decl.cpp

@@ -592,26 +592,10 @@ GenericContext::getLazyGenericEnvironmentSlow() const {
   assert(GenericSigOrEnv.is<GenericSignature *>() &&
          "not a lazily deserialized generic environment");
 
-  GenericEnvironment *genericEnv;
-
-  if (auto *extensionDecl = dyn_cast<ExtensionDecl>(this)) {
-    auto contextData = getASTContext().getOrCreateLazyIterableContextData(
-      extensionDecl, nullptr);
-    genericEnv = contextData->loader->loadGenericEnvironment(
-      extensionDecl, contextData->genericEnvData);
-  } else if (auto *typeDecl = dyn_cast<GenericTypeDecl>(this)) {
-    auto contextData = getASTContext().getOrCreateLazyGenericTypeData(
-      typeDecl, nullptr);
-    genericEnv = contextData->loader->loadGenericEnvironment(
-      typeDecl, contextData->genericEnvData);
-  } else if (auto *funcDecl = dyn_cast<AbstractFunctionDecl>(this)) {
-    auto contextData = getASTContext().getOrCreateLazyFunctionContextData(
-      funcDecl, nullptr);
-    genericEnv = contextData->loader->loadGenericEnvironment(
-      funcDecl, contextData->genericEnvData);
-  } else {
-    llvm_unreachable("Bad GenericContext kind");
-  }
+  auto contextData = getASTContext().getOrCreateLazyGenericContextData(
+    this, nullptr);
+  auto *genericEnv = contextData->loader->loadGenericEnvironment(
+    this, contextData->genericEnvData);
 
   const_cast<GenericContext *>(this)->setGenericEnvironment(genericEnv);
   ++NumLazyGenericEnvironmentsLoaded;
@@ -624,24 +608,9 @@ void GenericContext::setLazyGenericEnvironment(LazyMemberLoader *lazyLoader,
   assert(GenericSigOrEnv.isNull() && "already have a generic signature");
   GenericSigOrEnv = genericSig;
 
-  if (auto *extensionDecl = dyn_cast<ExtensionDecl>(this)) {
-    auto contextData =
-      getASTContext().getOrCreateLazyIterableContextData(
-        extensionDecl, lazyLoader);
-    contextData->genericEnvData = genericEnvData;
-  } else if (auto *typeDecl = dyn_cast<GenericTypeDecl>(this)) {
-    auto contextData =
-      getASTContext().getOrCreateLazyGenericTypeData(
-        typeDecl, lazyLoader);
-    contextData->genericEnvData = genericEnvData;
-  } else if (auto *funcDecl = dyn_cast<AbstractFunctionDecl>(this)) {
-    auto contextData =
-      getASTContext().getOrCreateLazyFunctionContextData(
-        funcDecl, lazyLoader);
-    contextData->genericEnvData = genericEnvData;
-  } else {
-    llvm_unreachable("Bad GenericContext kind");
-  }
+  auto contextData =
+    getASTContext().getOrCreateLazyGenericContextData(this, lazyLoader);
+  contextData->genericEnvData = genericEnvData;
 
   ++NumLazyGenericEnvironments;
 }

lib/AST/DeclContext.cpp

@@ -193,11 +193,15 @@ GenericParamList *DeclContext::getGenericParamsOfContext() const {
     case DeclContextKind::SerializedLocal:
     case DeclContextKind::Initializer:
     case DeclContextKind::AbstractClosureExpr:
-    case DeclContextKind::SubscriptDecl:
       // Closures and initializers can't themselves be generic, but they
       // can occur in generic contexts.
       continue;
 
+    case DeclContextKind::SubscriptDecl:
+      if (auto GP = cast<SubscriptDecl>(dc)->getGenericParams())
+        return GP;
+      continue;
+
     case DeclContextKind::AbstractFunctionDecl:
       if (auto GP = cast<AbstractFunctionDecl>(dc)->getGenericParams())
         return GP;
@@ -228,11 +232,13 @@ GenericSignature *DeclContext::getGenericSignatureOfContext() const {
     case DeclContextKind::Initializer:
     case DeclContextKind::SerializedLocal:
     case DeclContextKind::AbstractClosureExpr:
-    case DeclContextKind::SubscriptDecl:
       // Closures and initializers can't themselves be generic, but they
       // can occur in generic contexts.
       continue;
 
+    case DeclContextKind::SubscriptDecl:
+      return cast<SubscriptDecl>(dc)->getGenericSignature();
+
     case DeclContextKind::AbstractFunctionDecl:
       return cast<AbstractFunctionDecl>(dc)->getGenericSignature();
 
@@ -257,11 +263,13 @@ GenericEnvironment *DeclContext::getGenericEnvironmentOfContext() const {
     case DeclContextKind::Initializer:
     case DeclContextKind::SerializedLocal:
     case DeclContextKind::AbstractClosureExpr:
-    case DeclContextKind::SubscriptDecl:
       // Closures and initializers can't themselves be generic, but they
       // can occur in generic contexts.
       continue;
 
+    case DeclContextKind::SubscriptDecl:
+      return cast<SubscriptDecl>(dc)->getGenericEnvironment();
+
     case DeclContextKind::AbstractFunctionDecl:
       return cast<AbstractFunctionDecl>(dc)->getGenericEnvironment();
 
@@ -435,7 +443,12 @@ bool DeclContext::isGenericContext() const {
     case DeclContextKind::Initializer:
     case DeclContextKind::AbstractClosureExpr:
     case DeclContextKind::SerializedLocal:
+      // Check parent context.
+      continue;
+
     case DeclContextKind::SubscriptDecl:
+      if (cast<SubscriptDecl>(dc)->getGenericParams())
+        return true;
       // Check parent context.
       continue;
 
@@ -516,12 +529,14 @@ ResilienceExpansion DeclContext::getResilienceExpansion() const {
 /// Determine whether the innermost context is generic.
 bool DeclContext::isInnermostContextGeneric() const {
   switch (getContextKind()) {
+  case DeclContextKind::SubscriptDecl:
+    return cast<SubscriptDecl>(this)->isGeneric();
   case DeclContextKind::AbstractFunctionDecl:
-    return (cast<AbstractFunctionDecl>(this)->getGenericParams() != nullptr);
+    return cast<AbstractFunctionDecl>(this)->isGeneric();
   case DeclContextKind::ExtensionDecl:
-    return (cast<ExtensionDecl>(this)->getGenericParams() != nullptr);
+    return cast<ExtensionDecl>(this)->isGeneric();
   case DeclContextKind::GenericTypeDecl:
-    return (cast<GenericTypeDecl>(this)->getGenericParams() != nullptr);
+    return cast<GenericTypeDecl>(this)->isGeneric();
   default:
     return false;
   }

lib/ClangImporter/ImportDecl.cpp

@@ -5984,16 +5984,18 @@ SwiftDeclConverter::importSubscript(Decl *decl,
   auto subscript = Impl.createDeclWithClangNode<SubscriptDecl>(
       getter->getClangNode(), getOverridableAccessibility(dc), name,
       decl->getLoc(), bodyParams, decl->getLoc(),
-      TypeLoc::withoutLoc(elementContextTy), dc);
+      TypeLoc::withoutLoc(elementContextTy), dc,
+      /*GenericParams=*/nullptr);
 
   /// Record the subscript as an alternative declaration.
   Impl.addAlternateDecl(associateWithSetter ? setter : getter, subscript);
 
+  subscript->setGenericEnvironment(dc->getGenericEnvironmentOfContext());
+
   subscript->makeComputed(SourceLoc(), getterThunk, setterThunk, nullptr,
                           SourceLoc());
   auto indicesType = bodyParams->getType(C);
 
-  // TODO: no good when generics are around
   auto fnType = FunctionType::get(indicesType, elementTy);
   subscript->setInterfaceType(fnType);