AST: Requestify generic signature building for @_specialized attributes.

tshortli · tshortli · commit 497034ef6425 · 2023-09-16T00:02:33.000-07:00
In order to support lazy typechecking during module emission for modules
containing specialized functions, the computation of generic signatures for
`@_specialized` attributes must be requestified.

Resolves rdar://115569606
diff --git a/include/swift/AST/Attr.h b/include/swift/AST/Attr.h
@@ -1427,6 +1427,7 @@ class SpecializeAttr final
       private llvm::TrailingObjects<SpecializeAttr, Identifier,
                                     AvailableAttr *, Type> {
   friend class SpecializeAttrTargetDeclRequest;
+  friend class SerializeAttrGenericSignatureRequest;
   friend TrailingObjects;
 
 public:
@@ -1519,14 +1520,6 @@ class SpecializeAttr final
 
   TrailingWhereClause *getTrailingWhereClause() const;
 
-  GenericSignature getSpecializedSignature() const {
-    return specializedSignature;
-  }
-
-  void setSpecializedSignature(GenericSignature newSig) {
-    specializedSignature = newSig;
-  }
-
   bool isExported() const {
     return Bits.SpecializeAttr.exported;
   }
@@ -1550,6 +1543,10 @@ class SpecializeAttr final
   /// \p forDecl is the value decl that the attribute belongs to.
   ValueDecl *getTargetFunctionDecl(const ValueDecl *forDecl) const;
 
+  /// \p forDecl is the value decl that the attribute belongs to.
+  GenericSignature
+  getSpecializedSignature(const AbstractFunctionDecl *forDecl) const;
+
   static bool classof(const DeclAttribute *DA) {
     return DA->getKind() == DAK_Specialize;
   }
diff --git a/include/swift/AST/TypeCheckRequests.h b/include/swift/AST/TypeCheckRequests.h
@@ -4399,6 +4399,28 @@ class ExpandChildTypeRefinementContextsRequest
   bool isCached() const { return true; }
 };
 
+class SerializeAttrGenericSignatureRequest
+    : public SimpleRequest<SerializeAttrGenericSignatureRequest,
+                           GenericSignature(const AbstractFunctionDecl *,
+                                            SpecializeAttr *),
+                           RequestFlags::Cached> {
+public:
+  using SimpleRequest::SimpleRequest;
+
+private:
+  friend SimpleRequest;
+
+  GenericSignature evaluate(Evaluator &evaluator,
+                            const AbstractFunctionDecl *decl,
+                            SpecializeAttr *attr) const;
+
+public:
+  // Separate caching.
+  bool isCached() const { return true; }
+  llvm::Optional<GenericSignature> getCachedResult() const;
+  void cacheResult(GenericSignature signature) const;
+};
+
 #define SWIFT_TYPEID_ZONE TypeChecker
 #define SWIFT_TYPEID_HEADER "swift/AST/TypeCheckerTypeIDZone.def"
 #include "swift/Basic/DefineTypeIDZone.h"
diff --git a/include/swift/AST/TypeCheckerTypeIDZone.def b/include/swift/AST/TypeCheckerTypeIDZone.def
@@ -500,3 +500,6 @@ SWIFT_REQUEST(TypeChecker, InitAccessorReferencedVariablesRequest,
 SWIFT_REQUEST(TypeChecker, ExpandChildTypeRefinementContextsRequest,
               bool(Decl *, TypeRefinementContext *),
               Cached, NoLocationInfo)
+SWIFT_REQUEST(TypeChecker, SerializeAttrGenericSignatureRequest,
+              bool(Decl *, SpecializeAttr *),
+              Cached, NoLocationInfo)
diff --git a/lib/AST/Attr.cpp b/lib/AST/Attr.cpp
@@ -1198,12 +1198,13 @@ bool DeclAttribute::printImpl(ASTPrinter &Printer, const PrintOptions &Options,
       }
     }
     SmallVector<Requirement, 4> requirementsScratch;
-    auto requirements = attr->getSpecializedSignature().getRequirements();
     auto *FnDecl = dyn_cast_or_null<AbstractFunctionDecl>(D);
+    auto specializedSig = attr->getSpecializedSignature(FnDecl);
+    auto requirements = specializedSig.getRequirements();
     if (FnDecl && FnDecl->getGenericSignature()) {
       auto genericSig = FnDecl->getGenericSignature();
 
-      if (auto sig = attr->getSpecializedSignature()) {
+      if (auto sig = specializedSig) {
         requirementsScratch = sig.requirementsNotSatisfiedBy(genericSig);
         requirements = requirementsScratch;
       }
@@ -2204,6 +2205,14 @@ ValueDecl * SpecializeAttr::getTargetFunctionDecl(const ValueDecl *onDecl) const
                            nullptr);
 }
 
+GenericSignature SpecializeAttr::getSpecializedSignature(
+    const AbstractFunctionDecl *onDecl) const {
+  return evaluateOrDefault(onDecl->getASTContext().evaluator,
+                           SerializeAttrGenericSignatureRequest{
+                               onDecl, const_cast<SpecializeAttr *>(this)},
+                           nullptr);
+}
+
 SPIAccessControlAttr::SPIAccessControlAttr(SourceLoc atLoc, SourceRange range,
                                            ArrayRef<Identifier> spiGroups)
       : DeclAttribute(DAK_SPIAccessControl, atLoc, range,
diff --git a/lib/SIL/IR/SILFunctionBuilder.cpp b/lib/SIL/IR/SILFunctionBuilder.cpp
@@ -72,7 +72,7 @@ void SILFunctionBuilder::addFunctionAttributes(
             : SILSpecializeAttr::SpecializationKind::Partial;
     assert(!constant.isNull());
     SILFunction *targetFunction = nullptr;
-    auto *attributedFuncDecl = constant.getDecl();
+    auto *attributedFuncDecl = constant.getAbstractFunctionDecl();
     auto *targetFunctionDecl = SA->getTargetFunctionDecl(attributedFuncDecl);
     // Filter out _spi.
     auto spiGroups = SA->getSPIGroups();
@@ -91,16 +91,17 @@ void SILFunctionBuilder::addFunctionAttributes(
     auto availability =
       AvailabilityInference::annotatedAvailableRangeForAttr(SA,
          M.getSwiftModule()->getASTContext());
+    auto specializedSignature = SA->getSpecializedSignature(attributedFuncDecl);
     if (targetFunctionDecl) {
       SILDeclRef declRef(targetFunctionDecl, constant.kind, false);
       targetFunction = getOrCreateDeclaration(targetFunctionDecl, declRef);
       F->addSpecializeAttr(SILSpecializeAttr::create(
-          M, SA->getSpecializedSignature(), SA->getTypeErasedParams(),
+          M, specializedSignature, SA->getTypeErasedParams(),
           SA->isExported(), kind, targetFunction, spiGroupIdent,
           attributedFuncDecl->getModuleContext(), availability));
     } else {
       F->addSpecializeAttr(SILSpecializeAttr::create(
-          M, SA->getSpecializedSignature(), SA->getTypeErasedParams(),
+          M, specializedSignature, SA->getTypeErasedParams(),
           SA->isExported(), kind, nullptr, spiGroupIdent,
           attributedFuncDecl->getModuleContext(), availability));
     }
diff --git a/lib/SIL/IR/SILSymbolVisitor.cpp b/lib/SIL/IR/SILSymbolVisitor.cpp
@@ -442,8 +442,9 @@ class SILSymbolVisitorImpl : public ASTVisitor<SILSymbolVisitorImpl> {
       if (!attr->isExported())
         continue;
 
-      auto erasedSignature = attr->getSpecializedSignature().typeErased(
-          attr->getTypeErasedParams());
+      auto specializedSignature = attr->getSpecializedSignature(AFD);
+      auto erasedSignature =
+          specializedSignature.typeErased(attr->getTypeErasedParams());
 
       if (auto *targetFun = attr->getTargetFunctionDecl(AFD)) {
         addFunction(SILDeclRef(targetFun, erasedSignature),
diff --git a/lib/SILOptimizer/Transforms/GenericSpecializer.cpp b/lib/SILOptimizer/Transforms/GenericSpecializer.cpp
@@ -68,7 +68,7 @@ static void transferSpecializeAttributeTargets(SILModule &M,
           SA, M.getSwiftModule()->getASTContext());
 
       targetSILFunction->addSpecializeAttr(SILSpecializeAttr::create(
-          M, SA->getSpecializedSignature(), SA->getTypeErasedParams(),
+          M, SA->getSpecializedSignature(vd), SA->getTypeErasedParams(),
           SA->isExported(), kind, nullptr,
           spiGroupIdent, vd->getModuleContext(), availability));
     }
diff --git a/lib/Sema/TypeCheckAttr.cpp b/lib/Sema/TypeCheckAttr.cpp
@@ -2834,10 +2834,25 @@ void AttributeChecker::visitSpecializeAttr(SpecializeAttr *attr) {
     return;
   }
 
+  (void)attr->getSpecializedSignature(FD);
+}
+
+GenericSignature
+SerializeAttrGenericSignatureRequest::evaluate(Evaluator &evaluator,
+                                               const AbstractFunctionDecl *FD,
+                                               SpecializeAttr *attr) const {
+  if (attr->specializedSignature)
+    return attr->specializedSignature;
+
+  auto &Ctx = FD->getASTContext();
+  auto genericSig = FD->getGenericSignature();
+  if (!genericSig)
+    return nullptr;
+
   InferredGenericSignatureRequest request{
       genericSig.getPointer(),
       /*genericParams=*/nullptr,
-      WhereClauseOwner(FD, attr),
+      WhereClauseOwner(const_cast<AbstractFunctionDecl *>(FD), attr),
       /*addedRequirements=*/{},
       /*inferenceSources=*/{},
       /*allowConcreteGenericParams=*/true};
@@ -2865,10 +2880,26 @@ void AttributeChecker::visitSpecializeAttr(SpecializeAttr *attr) {
     attr->setTypeErasedParams(typeErasedParams);
   }
 
-  attr->setSpecializedSignature(specializedSig);
-
   // Check the target function if there is one.
   attr->getTargetFunctionDecl(FD);
+
+  return specializedSig;
+}
+
+llvm::Optional<GenericSignature>
+SerializeAttrGenericSignatureRequest::getCachedResult() const {
+  const auto &storage = getStorage();
+  SpecializeAttr *attr = std::get<1>(storage);
+  if (auto signature = attr->specializedSignature)
+    return signature;
+  return llvm::None;
+}
+
+void SerializeAttrGenericSignatureRequest::cacheResult(
+    GenericSignature signature) const {
+  const auto &storage = getStorage();
+  SpecializeAttr *attr = std::get<1>(storage);
+  attr->specializedSignature = signature;
 }
 
 void AttributeChecker::visitFixedLayoutAttr(FixedLayoutAttr *attr) {
diff --git a/lib/Serialization/Serialization.cpp b/lib/Serialization/Serialization.cpp
@@ -2883,7 +2883,8 @@ class Serializer::DeclSerializer : public DeclVisitor<DeclSerializer> {
       auto abbrCode = S.DeclTypeAbbrCodes[SpecializeDeclAttrLayout::Code];
       auto attr = cast<SpecializeAttr>(DA);
       auto targetFun = attr->getTargetFunctionName();
-      auto *targetFunDecl = attr->getTargetFunctionDecl(cast<ValueDecl>(D));
+      auto *afd = cast<AbstractFunctionDecl>(D);
+      auto *targetFunDecl = attr->getTargetFunctionDecl(afd);
 
       SmallVector<IdentifierID, 4> pieces;
 
@@ -2917,7 +2918,7 @@ class Serializer::DeclSerializer : public DeclVisitor<DeclSerializer> {
       SpecializeDeclAttrLayout::emitRecord(
           S.Out, S.ScratchRecord, abbrCode, (unsigned)attr->isExported(),
           (unsigned)attr->getSpecializationKind(),
-          S.addGenericSignatureRef(attr->getSpecializedSignature()),
+          S.addGenericSignatureRef(attr->getSpecializedSignature(afd)),
           S.addDeclRef(targetFunDecl), numArgs, numSPIGroups,
           numAvailabilityAttrs, numTypeErasedParams,
           pieces);
diff --git a/test/Inputs/lazy_typecheck.swift b/test/Inputs/lazy_typecheck.swift
@@ -33,6 +33,11 @@ public func constrainedGenericPublicFunction<T>(_ t: T) where T: PublicProto {
   doesNotExist() // expected-error {{cannot find 'doesNotExist' in scope}}
 }
 
+@_specialize(exported: true, where T == PublicProto)
+public func publicSpecializedFunc<T>(_ t: T) -> T {
+  return doesNotExist() // expected-error {{cannot find 'doesNotExist' in scope}}
+}
+
 @available(SwiftStdlib 5.1, *)
 public func publicFuncWithOpaqueReturnType() -> some PublicProto { // expected-note {{opaque return type declared here}}
   return 1 // expected-error {{return type of global function 'publicFuncWithOpaqueReturnType()' requires that 'Int' conform to 'PublicProto'}}
diff --git a/test/Inputs/lazy_typecheck_client.swift b/test/Inputs/lazy_typecheck_client.swift
@@ -14,6 +14,8 @@ func testGlobalFunctions() {
   _ = packageFunc()
   #endif
   constrainedGenericPublicFunction(ConformsToPublicProto())
+  _ = publicSpecializedFunc(4)
+  _ = publicSpecializedFunc(ConformsToPublicProto())
   if #available(SwiftStdlib 5.1, *) {
     _ = publicFuncWithOpaqueReturnType()
     _ = publicAEICFuncWithOpaqueReturnType()
diff --git a/test/ModuleInterface/lazy-typecheck.swift b/test/ModuleInterface/lazy-typecheck.swift
@@ -14,6 +14,8 @@
 // CHECK-NEXT:    return 1
 // CHECK-NEXT:  }
 // CHECK:       public func constrainedGenericPublicFunction<T>(_ t: T) where T : lazy_typecheck.PublicProto
+// CHECK:       @_specialize(exported: true, kind: full, where T == any lazy_typecheck.PublicProto)
+// CHECK-NEXT:  public func publicSpecializedFunc<T>(_ t: T) -> T
 // CHECK:       @available(macOS 10.15, iOS 13.0, watchOS 6.0, tvOS 13.0, *)
 // CHECK-NEXT:  public func publicFuncWithOpaqueReturnType() -> some lazy_typecheck.PublicProto
 
diff --git a/test/Serialization/serialize-external-decls-only-lazy-typecheck.swift b/test/Serialization/serialize-external-decls-only-lazy-typecheck.swift
@@ -3,3 +3,6 @@
 // RUN: %target-swift-frontend -swift-version 5 %S/../Inputs/lazy_typecheck.swift -module-name lazy_typecheck -emit-module -emit-module-path %t/lazy_typecheck.swiftmodule -enable-library-evolution -parse-as-library -package-name Package -DFLAG -experimental-lazy-typecheck -experimental-skip-all-function-bodies -experimental-serialize-external-decls-only
 
 // RUN: %target-swift-frontend -package-name Package -typecheck -verify %S/../Inputs/lazy_typecheck_client.swift -DFLAG -I %t
+
+// FIXME: Re-run the test with -experimental-skip-non-inlinable-function-bodies
+
diff --git a/test/TBD/lazy-typecheck.swift b/test/TBD/lazy-typecheck.swift
@@ -123,6 +123,7 @@ exports:
                        '_$s14lazy_typecheck19PublicGenericStructVyxGAA05EmptyC5ProtoA2A08InternalE13ForConstraintVRszlWP',
                        '_$s14lazy_typecheck19PublicRethrowsProtoMp', '_$s14lazy_typecheck19PublicRethrowsProtoP3reqSiyKFTj',
                        '_$s14lazy_typecheck19PublicRethrowsProtoP3reqSiyKFTq', '_$s14lazy_typecheck19PublicRethrowsProtoTL',
+                       '_$s14lazy_typecheck21publicSpecializedFuncyxxlF', '_$s14lazy_typecheck21publicSpecializedFuncyxxlFAA11PublicProto_p_Ts5',
                        '_$s14lazy_typecheck24publicFuncWithDefaultArgyS2iF', '_$s14lazy_typecheck27publicGlobalVarInferredTypeSSvM',
                        '_$s14lazy_typecheck27publicGlobalVarInferredTypeSSvg', '_$s14lazy_typecheck27publicGlobalVarInferredTypeSSvs',
                        '_$s14lazy_typecheck30publicFuncWithOpaqueReturnTypeQryF',

Original file line number	Diff line number	Diff line change
`@@ -68,7 +68,7 @@ static void transferSpecializeAttributeTargets(SILModule &M,`
`68`	`68`	`SA, M.getSwiftModule()->getASTContext());`
`69`	`69`
`70`	`70`	`targetSILFunction->addSpecializeAttr(SILSpecializeAttr::create(`
`71`		`- M, SA->getSpecializedSignature(), SA->getTypeErasedParams(),`
	`71`	`+ M, SA->getSpecializedSignature(vd), SA->getTypeErasedParams(),`
`72`	`72`	`SA->isExported(), kind, nullptr,`
`73`	`73`	`spiGroupIdent, vd->getModuleContext(), availability));`
`74`	`74`	`}`