Merge pull request #32657 from aschwaighofer/wip_prespecialize_exported

Preliminary support for `_specialize(exported: true, ...)`

Author: aschwaighofer

docs/ABI/Mangling.rst

@@ -995,6 +995,7 @@ Function Specializations
 ::
 
   specialization ::= type '_' type* 'Tg' SPEC-INFO     // Generic re-abstracted specialization
+  specialization ::= type '_' type* 'Ts' SPEC-INFO     // Generic re-abstracted prespecialization
   specialization ::= type '_' type* 'TG' SPEC-INFO     // Generic not re-abstracted specialization
   specialization ::= type '_' type* 'Ti' SPEC-INFO     // Inlined function with generic substitutions.
 

include/swift/AST/Attr.def

@@ -584,6 +584,11 @@ SIMPLE_DECL_ATTR(globalActor, GlobalActor,
   APIStableToAdd | APIBreakingToRemove,
   104)
 
+SIMPLE_DECL_ATTR(_specializeExtension, SpecializeExtension,
+  OnExtension | UserInaccessible |
+  ABIStableToAdd | ABIStableToRemove | APIStableToAdd | APIStableToRemove,
+  105)
+
 #undef TYPE_ATTR
 #undef DECL_ATTR_ALIAS
 #undef CONTEXTUAL_DECL_ATTR_ALIAS

include/swift/AST/Attr.h

@@ -1408,7 +1408,12 @@ class SynthesizedProtocolAttr : public DeclAttribute {
 
 /// The @_specialize attribute, which forces specialization on the specified
 /// type list.
-class SpecializeAttr : public DeclAttribute {
+class SpecializeAttr final
+    : public DeclAttribute,
+      private llvm::TrailingObjects<SpecializeAttr, Identifier> {
+  friend class SpecializeAttrTargetDeclRequest;
+  friend TrailingObjects;
+
 public:
   // NOTE: When adding new kinds, you must update the inline bitfield macro.
   enum class SpecializationKind {
@@ -1420,18 +1425,48 @@ class SpecializeAttr : public DeclAttribute {
   TrailingWhereClause *trailingWhereClause;
   GenericSignature specializedSignature;
 
+  DeclNameRef targetFunctionName;
+  LazyMemberLoader *resolver = nullptr;
+  uint64_t resolverContextData;
+  size_t numSPIGroups;
+
   SpecializeAttr(SourceLoc atLoc, SourceRange Range,
                  TrailingWhereClause *clause, bool exported,
-                 SpecializationKind kind,
-                 GenericSignature specializedSignature);
+                 SpecializationKind kind, GenericSignature specializedSignature,
+                 DeclNameRef targetFunctionName,
+                 ArrayRef<Identifier> spiGroups);
 
 public:
   static SpecializeAttr *create(ASTContext &Ctx, SourceLoc atLoc,
                                 SourceRange Range, TrailingWhereClause *clause,
                                 bool exported, SpecializationKind kind,
+                                DeclNameRef targetFunctionName,
+                                ArrayRef<Identifier> spiGroups,
                                 GenericSignature specializedSignature
                                     = nullptr);
 
+  static SpecializeAttr *create(ASTContext &ctx, bool exported,
+                                SpecializationKind kind,
+                                ArrayRef<Identifier> spiGroups,
+                                GenericSignature specializedSignature,
+                                DeclNameRef replacedFunction);
+
+  static SpecializeAttr *create(ASTContext &ctx, bool exported,
+                                SpecializationKind kind,
+                                ArrayRef<Identifier> spiGroups,
+                                GenericSignature specializedSignature,
+                                DeclNameRef replacedFunction,
+                                LazyMemberLoader *resolver, uint64_t data);
+
+  /// Name of SPIs declared by the attribute.
+  ///
+  /// Note: A single SPI name per attribute is currently supported but this
+  /// may change with the syntax change.
+  ArrayRef<Identifier> getSPIGroups() const {
+    return { this->template getTrailingObjects<Identifier>(),
+             numSPIGroups };
+  }
+
   TrailingWhereClause *getTrailingWhereClause() const;
 
   GenericSignature getSpecializedSignature() const {
@@ -1458,6 +1493,13 @@ class SpecializeAttr : public DeclAttribute {
     return getSpecializationKind() == SpecializationKind::Partial;
   }
 
+  DeclNameRef getTargetFunctionName() const {
+    return targetFunctionName;
+  }
+
+  /// \p forDecl is the value decl that the attribute belongs to.
+  ValueDecl *getTargetFunctionDecl(const ValueDecl *forDecl) const;
+
   static bool classof(const DeclAttribute *DA) {
     return DA->getKind() == DAK_Specialize;
   }

include/swift/AST/Decl.h

@@ -2187,7 +2187,8 @@ class ValueDecl : public Decl {
   /// implementations for the requirements of a public protocol, even when
   /// the default implementations are not visible to name lookup.
   bool isAccessibleFrom(const DeclContext *DC,
-                        bool forConformance = false) const;
+                        bool forConformance = false,
+                        bool includeInlineable = false) const;
 
   /// Returns whether this declaration should be treated as \c open from
   /// \p useDC. This is very similar to #getFormalAccess, but takes

include/swift/AST/DeclContext.h

@@ -490,6 +490,10 @@ class alignas(1 << DeclContextAlignInBits) DeclContext {
   /// Determine whether the innermost context is generic.
   bool isInnermostContextGeneric() const;
 
+  /// Determine whether this or any parent context is a `@_specialize` extension
+  /// context.
+  bool isInSpecializeExtensionContext() const;
+
   /// Get the most optimal resilience expansion for code in this context.
   /// If the body is able to be inlined into functions in other resilience
   /// domains, this ensures that only sufficiently-conservative access patterns

include/swift/AST/DiagnosticsParse.def

@@ -626,6 +626,8 @@ ERROR(expected_sil_function_type, none,
       "sil function expected to have SIL function type", ())
 ERROR(sil_dynamically_replaced_func_not_found,none,
       "dynamically replaced function not found %0", (Identifier))
+ERROR(sil_specialize_target_func_not_found,none,
+      "_specialize target function not found %0", (Identifier))
 ERROR(sil_availability_expected_version,none,
       "expected version number in 'available' attribute", ())
 
@@ -1579,6 +1581,10 @@ ERROR(attr_specialize_parameter_already_defined,none,
 
 ERROR(attr_specialize_expected_partial_or_full,none,
       "expected 'partial' or 'full' as values of the 'kind' parameter in '_specialize' attribute", ())
+ERROR(attr_specialize_expected_function,none,
+      "expected a function name as the value of the 'target' parameter in '_specialize' attribute", ())
+ERROR(attr_specialize_expected_spi_name,none,
+      "expected an SPI identifier as the value of the 'spi' parameter in '_specialize' attribute", ())
 
 // _implements
 ERROR(attr_implements_expected_member_name,PointsToFirstBadToken,

include/swift/AST/DiagnosticsSema.def

@@ -5125,6 +5125,12 @@ ERROR(specialize_attr_missing_constraint,none,
       "Missing constraint for %0 in '_specialize' attribute", (DeclName))
 ERROR(specialize_attr_unsupported_kind_of_req,none,
       "Only same-type and layout requirements are supported by '_specialize' attribute", ())
+ERROR(specialize_target_function_not_found, none,
+     "target function %0 could not be found", (DeclNameRef))
+ERROR(specialize_target_function_of_type_not_found, none,
+      "target function %0 of type %1 could not be found", (DeclNameRef, Type))
+NOTE(specialize_found_function_of_type, none,
+      "found function %0 of type %1", (DeclName, Type))
 
 //------------------------------------------------------------------------------
 // MARK: Variable usage diagnostics

include/swift/AST/LazyResolver.h

@@ -114,6 +114,10 @@ class alignas(void*) LazyMemberLoader {
   /// Returns the type for a given @_typeEraser() attribute.
   virtual Type loadTypeEraserType(const TypeEraserAttr *TRA,
                                   uint64_t contextData) = 0;
+
+  // Returns the target parameter of the `@_specialize` attribute or null.
+  virtual ValueDecl *loadTargetFunctionDecl(const SpecializeAttr *attr,
+                                            uint64_t contextData) = 0;
 };
 
 /// A class that can lazily load conformances from a serialized format.

include/swift/AST/LookupKinds.h

@@ -50,6 +50,9 @@ enum NLOptions : unsigned {
   /// Include synonyms declared with @_implements()
   NL_IncludeAttributeImplements = 1 << 5,
 
+  // Include @usableFromInline and @inlinable
+  NL_IncludeUsableFromInlineAndInlineable = 1 << 6,
+
   /// The default set of options used for qualified name lookup.
   ///
   /// FIXME: Eventually, add NL_ProtocolMembers to this, once all of the

include/swift/AST/Module.h

@@ -578,6 +578,13 @@ class ModuleDecl : public DeclContext, public TypeDecl {
                          const ModuleDecl *importedModule,
                          llvm::SmallSetVector<Identifier, 4> &spiGroups) const;
 
+  // Is \p attr accessible as an explictly imported SPI from this module?
+  bool isImportedAsSPI(const SpecializeAttr *attr,
+                       const ValueDecl *targetDecl) const;
+
+  // Is \p spiGroup accessible as an explictly imported SPI from this module?
+  bool isImportedAsSPI(Identifier spiGroup, const ModuleDecl *fromModule) const;
+
   /// \sa getImportedModules
   enum class ImportFilterKind {
     /// Include imports declared with `@_exported`.

include/swift/AST/ModuleNameLookup.h

@@ -56,10 +56,13 @@ void simple_display(llvm::raw_ostream &out, ResolutionKind kind);
 /// \param moduleScopeContext The top-level context from which the lookup is
 ///        being performed, for checking access. This must be either a
 ///        FileUnit or a Module.
+/// \param options name lookup options. Currently only used to communicate the
+/// NL_IncludeUsableFromInlineAndInlineable option.
 void lookupInModule(const DeclContext *moduleOrFile,
                     DeclName name, SmallVectorImpl<ValueDecl *> &decls,
                     NLKind lookupKind, ResolutionKind resolutionKind,
-                    const DeclContext *moduleScopeContext);
+                    const DeclContext *moduleScopeContext,
+                    NLOptions options);
 
 /// Performs a qualified lookup into the given module and, if necessary, its
 /// reexports, observing proper shadowing rules.

include/swift/AST/NameLookup.h

@@ -225,6 +225,9 @@ enum class UnqualifiedLookupFlags {
   /// This lookup should include results from outside the innermost scope with
   /// results.
   IncludeOuterResults   = 1 << 4,
+  // This lookup should include results that are @inlinable or
+  // @usableFromInline.
+  IncludeInlineableAndUsableFromInline = 1 << 5,
 };
 
 using UnqualifiedLookupOptions = OptionSet<UnqualifiedLookupFlags>;

include/swift/AST/NameLookupRequests.h

@@ -437,11 +437,12 @@ using QualifiedLookupResult = SmallVector<ValueDecl *, 4>;
 
 /// Performs a lookup into a given module and its imports.
 class LookupInModuleRequest
-    : public SimpleRequest<LookupInModuleRequest,
-                           QualifiedLookupResult(
-                               const DeclContext *, DeclName, NLKind,
-                               namelookup::ResolutionKind, const DeclContext *),
-                           RequestFlags::Uncached | RequestFlags::DependencySink> {
+    : public SimpleRequest<
+          LookupInModuleRequest,
+          QualifiedLookupResult(const DeclContext *, DeclName, NLKind,
+                                namelookup::ResolutionKind, const DeclContext *,
+                                NLOptions),
+          RequestFlags::Uncached | RequestFlags::DependencySink> {
 public:
   using SimpleRequest::SimpleRequest;
 
@@ -452,7 +453,7 @@ class LookupInModuleRequest
   QualifiedLookupResult
   evaluate(Evaluator &evaluator, const DeclContext *moduleOrFile, DeclName name,
            NLKind lookupKind, namelookup::ResolutionKind resolutionKind,
-           const DeclContext *moduleScopeContext) const;
+           const DeclContext *moduleScopeContext, NLOptions options) const;
 
 public:
   // Incremental dependencies

include/swift/AST/TypeCheckRequests.h

@@ -2177,6 +2177,25 @@ class DynamicallyReplacedDeclRequest
   bool isCached() const { return true; }
 };
 
+class SpecializeAttrTargetDeclRequest
+    : public SimpleRequest<SpecializeAttrTargetDeclRequest,
+                           ValueDecl *(const ValueDecl *, SpecializeAttr *),
+                           RequestFlags::Cached> {
+public:
+  using SimpleRequest::SimpleRequest;
+
+private:
+  friend SimpleRequest;
+
+  // Evaluation.
+  ValueDecl *evaluate(Evaluator &evaluator, const ValueDecl *vd,
+                      SpecializeAttr *attr) const;
+
+public:
+  // Caching.
+  bool isCached() const { return true; }
+};
+
 class TypeCheckSourceFileRequest
     : public SimpleRequest<
           TypeCheckSourceFileRequest, evaluator::SideEffect(SourceFile *),

include/swift/AST/TypeCheckerTypeIDZone.def

@@ -67,6 +67,9 @@ SWIFT_REQUEST(TypeChecker, DynamicallyReplacedDeclRequest,
               Cached, NoLocationInfo)
 SWIFT_REQUEST(TypeChecker, SemanticMembersRequest,
               ArrayRef<Decl *>(IterableDeclContext *), Cached, NoLocationInfo)
+SWIFT_REQUEST(TypeChecker, SpecializeAttrTargetDeclRequest,
+              ValueDecl *(const ValueDecl *, SpecializeAttr *),
+              Cached, NoLocationInfo)
 SWIFT_REQUEST(TypeChecker, EnumRawValuesRequest,
               evaluator::SideEffect (EnumDecl *, TypeResolutionStage),
               SeparatelyCached, NoLocationInfo)

include/swift/Demangling/DemangleNodes.def

@@ -104,6 +104,7 @@ NODE(ResilientProtocolWitnessTable)
 NODE(GenericSpecialization)
 NODE(GenericSpecializationNotReAbstracted)
 NODE(GenericSpecializationParam)
+NODE(GenericSpecializationPrespecialized)
 NODE(InlinedGenericFunction)
 NODE(GenericTypeMetadataPattern)
 CONTEXT_NODE(Getter)

include/swift/Parse/Parser.h

@@ -1003,14 +1003,22 @@ class Parser {
   /// Parse the @_specialize attribute.
   /// \p closingBrace is the expected closing brace, which can be either ) or ]
   /// \p Attr is where to store the parsed attribute
-  bool parseSpecializeAttribute(swift::tok ClosingBrace, SourceLoc AtLoc,
-                                SourceLoc Loc, SpecializeAttr *&Attr);
+  bool parseSpecializeAttribute(
+      swift::tok ClosingBrace, SourceLoc AtLoc, SourceLoc Loc,
+      SpecializeAttr *&Attr,
+      llvm::function_ref<bool(Parser &)> parseSILTargetName =
+          [](Parser &) { return false; },
+      llvm::function_ref<bool(Parser &)> parseSILSIPModule =
+          [](Parser &) { return false; });
 
   /// Parse the arguments inside the @_specialize attribute
   bool parseSpecializeAttributeArguments(
       swift::tok ClosingBrace, bool &DiscardAttribute, Optional<bool> &Exported,
       Optional<SpecializeAttr::SpecializationKind> &Kind,
-      TrailingWhereClause *&TrailingWhereClause);
+      TrailingWhereClause *&TrailingWhereClause, DeclNameRef &targetFunction,
+      SmallVectorImpl<Identifier> &spiGroups,
+      llvm::function_ref<bool(Parser &)> parseSILTargetName,
+      llvm::function_ref<bool(Parser &)> parseSILSIPModule);
 
   /// Parse the @_implements attribute.
   /// \p Attr is where to store the parsed attribute