Add support for _specialize(exported: true, ...)

This attribute allows to define a pre-specialized entry point of a
generic function in a library.

The following definition provides a pre-specialized entry point for
`genericFunc(_:)` for the parameter type `Int` that clients of the
library can call.

```
@_specialize(exported: true, where T == Int)
public func genericFunc<T>(_ t: T) { ... }
```

Pre-specializations of internal `@inlinable` functions are allowed.

```
@usableFromInline
internal struct GenericThing<T> {
  @_specialize(exported: true, where T == Int)
  @inlinable
  internal func genericMethod(_ t: T) {
  }
}
```

There is syntax to pre-specialize a method from a different module.

```
import ModuleDefiningGenericFunc

@_specialize(exported: true, target: genericFunc(_:), where T == Double)
func prespecialize_genericFunc(_ t: T) { fatalError("dont call") }

```

Specially marked extensions allow for pre-specialization of internal
methods accross module boundries (respecting `@inlinable` and
`@usableFromInline`).

```
import ModuleDefiningGenericThing
public struct Something {}

@_specializeExtension
extension GenericThing {
  @_specialize(exported: true, target: genericMethod(_:), where T == Something)
  func prespecialize_genericMethod(_ t: T) { fatalError("dont call") }
}
```

rdar://64993425

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

@@ -1409,6 +1409,8 @@ class SynthesizedProtocolAttr : public DeclAttribute {
 /// The @_specialize attribute, which forces specialization on the specified
 /// type list.
 class SpecializeAttr : public DeclAttribute {
+  friend class SpecializeAttrTargetDeclRequest;
+
 public:
   // NOTE: When adding new kinds, you must update the inline bitfield macro.
   enum class SpecializationKind {
@@ -1420,18 +1422,35 @@ class SpecializeAttr : public DeclAttribute {
   TrailingWhereClause *trailingWhereClause;
   GenericSignature specializedSignature;
 
+  DeclNameRef targetFunctionName;
+  LazyMemberLoader *resolver = nullptr;
+  uint64_t resolverContextData;
+
   SpecializeAttr(SourceLoc atLoc, SourceRange Range,
                  TrailingWhereClause *clause, bool exported,
                  SpecializationKind kind,
-                 GenericSignature specializedSignature);
+                 GenericSignature specializedSignature,
+                 DeclNameRef targetFunctionName);
 
 public:
   static SpecializeAttr *create(ASTContext &Ctx, SourceLoc atLoc,
                                 SourceRange Range, TrailingWhereClause *clause,
                                 bool exported, SpecializationKind kind,
+                                DeclNameRef targetFunctionName,
                                 GenericSignature specializedSignature
                                     = nullptr);
 
+  static SpecializeAttr *create(ASTContext &ctx, bool exported,
+                                SpecializationKind kind,
+                                GenericSignature specializedSignature,
+                                DeclNameRef replacedFunction);
+
+  static SpecializeAttr *create(ASTContext &ctx, bool exported,
+                                SpecializationKind kind,
+                                GenericSignature specializedSignature,
+                                DeclNameRef replacedFunction,
+                                LazyMemberLoader *resolver, uint64_t data);
+
   TrailingWhereClause *getTrailingWhereClause() const;
 
   GenericSignature getSpecializedSignature() const {
@@ -1458,6 +1477,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,8 @@ 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", ())
 
 // _implements
 ERROR(attr_implements_expected_member_name,PointsToFirstBadToken,

include/swift/AST/DiagnosticsSema.def

@@ -5113,6 +5113,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/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.