[IDE] A couple of fixes for @abi

include/swift/AST/ASTMangler.h

@@ -106,6 +106,12 @@ class ASTMangler : public Mangler {
   /// defined in.
   bool RespectOriginallyDefinedIn = true;
 
+  /// Whether to always mangle using the declaration's API, ignoring e.g
+  /// attached `@abi` attributes. This is necessary for USR mangling since for
+  /// semantic functionality we're only concerned about the API entity, and need
+  /// to be able to do name lookups to find the original decl based on the USR.
+  bool UseAPIMangling = false;
+
 public:
   class SymbolicReferent {
   public:
@@ -187,17 +193,26 @@ class ASTMangler : public Mangler {
     HasSymbolQuery,
   };
 
-  /// lldb overrides the defaulted argument to 'true'.
-  ASTMangler(const ASTContext &Ctx, bool DWARFMangling = false) : Context(Ctx) {
+  /// lldb overrides \p DWARFMangling to 'true'.
+  ASTMangler(const ASTContext &Ctx, bool DWARFMangling = false,
+             bool UseAPIMangling = false)
+      : Context(Ctx) {
     if (DWARFMangling) {
-      DWARFMangling = true;
+      this->DWARFMangling = true;
       RespectOriginallyDefinedIn = false;
     }
+    this->UseAPIMangling = UseAPIMangling;
     Flavor = Ctx.LangOpts.hasFeature(Feature::Embedded)
                  ? ManglingFlavor::Embedded
                  : ManglingFlavor::Default;
   }
 
+  /// Create an ASTMangler suitable for mangling a USR for use in semantic
+  /// functionality.
+  static ASTMangler forUSR(const ASTContext &Ctx) {
+    return ASTMangler(Ctx, /*DWARFMangling*/ true, /*UseAPIMangling*/ true);
+  }
+
   const ASTContext &getASTContext() { return Context; }
 
   void addTypeSubstitution(Type type, GenericSignature sig) {
@@ -392,9 +407,8 @@ class ASTMangler : public Mangler {
   std::string mangleTypeAsContextUSR(const NominalTypeDecl *type);
 
   void appendAnyDecl(const ValueDecl *Decl);
-  std::string mangleAnyDecl(const ValueDecl *Decl, bool prefix,
-                            bool respectOriginallyDefinedIn = false);
-  std::string mangleDeclAsUSR(const ValueDecl *Decl, StringRef USRPrefix);
+  std::string mangleAnyDecl(const ValueDecl *decl, bool addPrefix);
+  std::string mangleDeclWithPrefix(const ValueDecl *decl, StringRef prefix);
 
   std::string mangleAccessorEntityAsUSR(AccessorKind kind,
                                         const AbstractStorageDecl *decl,
@@ -791,6 +805,28 @@ class ASTMangler : public Mangler {
                                     const ValueDecl *forDecl);
 
   void appendLifetimeDependence(LifetimeDependenceInfo info);
+
+  void gatherExistentialRequirements(SmallVectorImpl<Requirement> &reqs,
+                                     ParameterizedProtocolType *PPT) const;
+
+  /// Extracts a list of inverse requirements from a PCT serving as the
+  /// constraint type of an existential.
+  void extractExistentialInverseRequirements(
+      SmallVectorImpl<InverseRequirement> &inverses,
+      ProtocolCompositionType *PCT) const;
+
+  template <typename DeclType>
+  DeclType *getABIDecl(DeclType *D) const {
+    if (!D || UseAPIMangling)
+      return nullptr;
+
+    auto abiRole = ABIRoleInfo(D);
+    if (!abiRole.providesABI())
+      return abiRole.getCounterpart();
+    return nullptr;
+  }
+
+  std::optional<SymbolicReferent> getABIDecl(SymbolicReferent ref) const;
 };
 
 } // end namespace Mangle

lib/APIDigester/ModuleAnalyzerNodes.cpp

@@ -1120,8 +1120,7 @@ static StringRef calculateMangledName(SDKContext &Ctx, ValueDecl *VD) {
     return Ctx.buffer(attr->Name);
   }
   Mangle::ASTMangler NewMangler(VD->getASTContext());
-  return Ctx.buffer(NewMangler.mangleAnyDecl(VD, true,
-                                    /*bool respectOriginallyDefinedIn*/true));
+  return Ctx.buffer(NewMangler.mangleAnyDecl(VD, /*addPrefix*/ true));
 }
 
 static StringRef calculateLocation(SDKContext &SDKCtx, Decl *D) {

lib/AST/ASTMangler.cpp

@@ -68,33 +68,22 @@
 using namespace swift;
 using namespace swift::Mangle;
 
-template<typename DeclType>
-static DeclType *getABIDecl(DeclType *D) {
-  if (!D)
-    return nullptr;
-
-  auto abiRole = ABIRoleInfo(D);
-  if (!abiRole.providesABI())
-    return abiRole.getCounterpart();
-  return nullptr;
-}
-
-static std::optional<ASTMangler::SymbolicReferent>
-getABIDecl(ASTMangler::SymbolicReferent ref) {
+std::optional<ASTMangler::SymbolicReferent>
+ASTMangler::getABIDecl(SymbolicReferent ref) const {
   switch (ref.getKind()) {
-  case ASTMangler::SymbolicReferent::NominalType:
+  case SymbolicReferent::NominalType:
     if (auto abiTypeDecl = getABIDecl(ref.getNominalType())) {
-      return ASTMangler::SymbolicReferent(abiTypeDecl);
+      return SymbolicReferent(abiTypeDecl);
     }
     break;
 
-  case ASTMangler::SymbolicReferent::OpaqueType:
+  case SymbolicReferent::OpaqueType:
     if (auto abiTypeDecl = getABIDecl(ref.getOpaqueType())) {
-      return ASTMangler::SymbolicReferent(abiTypeDecl);
+      return SymbolicReferent(abiTypeDecl);
     }
     break;
 
-  case ASTMangler::SymbolicReferent::ExtendedExistentialTypeShape:
+  case SymbolicReferent::ExtendedExistentialTypeShape:
     // Do nothing; mangling will use the underlying ABI decls in the end.
     break;
   }
@@ -969,8 +958,6 @@ std::string ASTMangler::mangleTypeAsContextUSR(const NominalTypeDecl *type) {
 }
 
 std::string ASTMangler::mangleTypeAsUSR(Type Ty) {
-  DWARFMangling = true;
-  RespectOriginallyDefinedIn = false;
   beginMangling();
 
   Ty = getTypeForDWARFMangling(Ty);
@@ -1006,42 +993,34 @@ void ASTMangler::appendAnyDecl(const ValueDecl *Decl) {
   }
 }
 
-std::string
-ASTMangler::mangleAnyDecl(const ValueDecl *Decl,
-                          bool prefix,
-                          bool respectOriginallyDefinedIn) {
+std::string ASTMangler::mangleAnyDecl(const ValueDecl *decl, bool addPrefix) {
   DWARFMangling = true;
-  RespectOriginallyDefinedIn = respectOriginallyDefinedIn;
-  if (prefix) {
+  if (addPrefix) {
     beginMangling();
   } else {
     beginManglingWithoutPrefix();
   }
   llvm::SaveAndRestore<bool> allowUnnamedRAII(AllowNamelessEntities, true);
 
-  appendAnyDecl(Decl);
+  appendAnyDecl(decl);
 
   // We have a custom prefix, so finalize() won't verify for us. If we're not
   // in invalid code (coming from an IDE caller) verify manually.
-  if (CONDITIONAL_ASSERT_enabled() && !prefix && !Decl->isInvalid())
+  if (CONDITIONAL_ASSERT_enabled() && !addPrefix && !decl->isInvalid())
     verify(Storage.str(), Flavor);
   return finalize();
 }
 
-std::string ASTMangler::mangleDeclAsUSR(const ValueDecl *Decl,
-                                        StringRef USRPrefix) {
-  llvm::SaveAndRestore<bool> respectOriginallyDefinedInRAII(
-      RespectOriginallyDefinedIn, false);
-  return (llvm::Twine(USRPrefix) + mangleAnyDecl(Decl, false)).str();
+std::string ASTMangler::mangleDeclWithPrefix(const ValueDecl *decl,
+                                             StringRef prefix) {
+  return (llvm::Twine(prefix) + mangleAnyDecl(decl, /*addPrefix*/ false)).str();
 }
 
 std::string ASTMangler::mangleAccessorEntityAsUSR(AccessorKind kind,
                                                   const AbstractStorageDecl *decl,
                                                   StringRef USRPrefix,
                                                   bool isStatic) {
   beginManglingWithoutPrefix();
-  llvm::SaveAndRestore<bool> respectOriginallyDefinedInRAII(
-      RespectOriginallyDefinedIn, false);
   llvm::SaveAndRestore<bool> allowUnnamedRAII(AllowNamelessEntities, true);
   Buffer << USRPrefix;
   appendAccessorEntity(getCodeForAccessorKind(kind), decl, isStatic);
@@ -5329,18 +5308,18 @@ std::string ASTMangler::mangleAttachedMacroExpansion(
   return finalize();
 }
 
-static void gatherExistentialRequirements(SmallVectorImpl<Requirement> &reqs,
-                                          ParameterizedProtocolType *PPT) {
+void ASTMangler::gatherExistentialRequirements(
+    SmallVectorImpl<Requirement> &reqs, ParameterizedProtocolType *PPT) const {
   auto protoTy = PPT->getBaseType();
   ASSERT(!getABIDecl(protoTy->getDecl()) && "need to figure out behavior");
   PPT->getRequirements(protoTy->getDecl()->getSelfInterfaceType(), reqs);
 }
 
 /// Extracts a list of inverse requirements from a PCT serving as the constraint
 /// type of an existential.
-static void extractExistentialInverseRequirements(
-                                SmallVectorImpl<InverseRequirement> &inverses,
-                                ProtocolCompositionType *PCT) {
+void ASTMangler::extractExistentialInverseRequirements(
+    SmallVectorImpl<InverseRequirement> &inverses,
+    ProtocolCompositionType *PCT) const {
   if (!PCT->hasInverse())
     return;
 

lib/AST/ASTPrinter.cpp

@@ -1348,9 +1348,9 @@ void PrintAST::printAttributes(const Decl *D) {
   if (Options.PrintSyntheticSILGenName
         && !D->getAttrs().hasAttribute<SILGenNameAttr>()) {
     if (canPrintSyntheticSILGenName(D)) {
-      auto mangledName = Mangle::ASTMangler(D->getASTContext())
-                            .mangleAnyDecl(cast<ValueDecl>(D), /*prefix=*/true,
-                                           /*respectOriginallyDefinedIn=*/true);
+      auto mangledName =
+          Mangle::ASTMangler(D->getASTContext())
+              .mangleAnyDecl(cast<ValueDecl>(D), /*addPrefix*/ true);
       Printer.printAttrName("@_silgen_name");
       Printer << "(";
       Printer.printEscapedStringLiteral(mangledName);

lib/AST/USRGeneration.cpp

@@ -40,7 +40,8 @@ static inline StringRef getUSRSpacePrefix() {
 bool ide::printTypeUSR(Type Ty, raw_ostream &OS) {
   assert(!Ty->hasArchetype() && "cannot have contextless archetypes mangled.");
   Ty = Ty->getCanonicalType()->getRValueType();
-  Mangle::ASTMangler Mangler(Ty->getASTContext());
+  using namespace Mangle;
+  ASTMangler Mangler = ASTMangler::forUSR(Ty->getASTContext());
   OS << Mangler.mangleTypeAsUSR(Ty);
   return false;
 }
@@ -245,8 +246,9 @@ swift::SwiftUSRGenerationRequest::evaluate(Evaluator &evaluator,
   if (declIFaceTy.findIf([](Type t) -> bool { return t->is<ModuleType>(); }))
     return std::string();
 
-  Mangle::ASTMangler NewMangler(D->getASTContext());
-  return NewMangler.mangleDeclAsUSR(D, getUSRSpacePrefix());
+  using namespace Mangle;
+  ASTMangler NewMangler = ASTMangler::forUSR(D->getASTContext());
+  return NewMangler.mangleDeclWithPrefix(D, getUSRSpacePrefix());
 }
 
 std::string
@@ -375,7 +377,8 @@ bool ide::printAccessorUSR(const AbstractStorageDecl *D, AccessorKind AccKind,
     return printObjCUSRForAccessor(SD, AccKind, OS);
   }
 
-  Mangle::ASTMangler NewMangler(D->getASTContext());
+  using namespace Mangle;
+  ASTMangler NewMangler = ASTMangler::forUSR(D->getASTContext());
   std::string Mangled = NewMangler.mangleAccessorEntityAsUSR(AccKind,
                           SD, getUSRSpacePrefix(), SD->isStatic());
 

lib/ClangImporter/SwiftDeclSynthesizer.cpp

@@ -998,7 +998,7 @@ getAccessorDeclarationName(clang::ASTContext &Ctx, NominalTypeDecl *structDecl,
   std::string id;
   llvm::raw_string_ostream IdStream(id);
   Mangle::ASTMangler mangler(structDecl->getASTContext());
-  IdStream << "$" << mangler.mangleDeclAsUSR(structDecl, "") << "$"
+  IdStream << "$" << mangler.mangleDeclWithPrefix(structDecl, "") << "$"
            << fieldDecl->getName() << "$" << suffix;
 
   return clang::DeclarationName(&Ctx.Idents.get(IdStream.str()));

lib/PrintAsClang/ModuleContentsWriter.cpp

@@ -155,8 +155,7 @@ static std::string getMangledNameString(const Decl *D) {
   if (!VD)
     return std::string();
   Mangle::ASTMangler mangler(VD->getASTContext());
-  return mangler.mangleAnyDecl(VD, /*prefix=*/true,
-                               /*respectOriginallyDefinedIn=*/true);
+  return mangler.mangleAnyDecl(VD, /*addPrefix*/ true);
 }
 
 static std::string getTypeString(const ValueDecl *VD) {

lib/Serialization/ModuleFile.cpp

@@ -753,7 +753,7 @@ void ModuleFile::loadDerivativeFunctionConfigurations(
     return;
   auto &ctx = originalAFD->getASTContext();
   Mangle::ASTMangler Mangler(ctx);
-  auto mangledName = Mangler.mangleDeclAsUSR(originalAFD, "");
+  auto mangledName = Mangler.mangleDeclWithPrefix(originalAFD, "");
   auto configs = Core->DerivativeFunctionConfigurations->find(mangledName);
   if (configs == Core->DerivativeFunctionConfigurations->end())
     return;

lib/Serialization/Serialization.cpp

@@ -6897,14 +6897,15 @@ static void recordDerivativeFunctionConfig(
   auto &ctx = AFD->getASTContext();
   Mangle::ASTMangler Mangler(AFD->getASTContext());
   for (auto *attr : AFD->getAttrs().getAttributes<DifferentiableAttr>()) {
-    auto mangledName = ctx.getIdentifier(Mangler.mangleDeclAsUSR(AFD, ""));
+    auto mangledName = ctx.getIdentifier(Mangler.mangleDeclWithPrefix(AFD, ""));
     derivativeConfigs[mangledName].insert(
         {ctx.getIdentifier(attr->getParameterIndices()->getString()),
          attr->getDerivativeGenericSignature()});
   }
   for (auto *attr : AFD->getAttrs().getAttributes<DerivativeAttr>()) {
     auto *origAFD = attr->getOriginalFunction(ctx);
-    auto mangledName = ctx.getIdentifier(Mangler.mangleDeclAsUSR(origAFD, ""));
+    auto mangledName =
+        ctx.getIdentifier(Mangler.mangleDeclWithPrefix(origAFD, ""));
     derivativeConfigs[mangledName].insert(
         {ctx.getIdentifier(attr->getParameterIndices()->getString()),
          AFD->getGenericSignature()});

test/Index/index_abi_attr.swift

@@ -0,0 +1,6 @@
+// RUN: %target-swift-ide-test -print-indexed-symbols -include-locals -source-filename %s | %FileCheck %s
+
+// Make sure we use the API decl for mangling the USR.
+@abi(func bar())
+public func foo() {}
+// CHECK: [[@LINE-1]]:13 | function/Swift | foo() | s:14swift_ide_test3fooyyF | Def | rel: 0

test/SourceKit/DocSupport/doc_abi.swift

@@ -0,0 +1,44 @@
+// RUN: %empty-directory(%t)
+// RUN: split-file %s %t
+// RUN: %target-swift-frontend -emit-module -module-name Mod %t/mod.swift -o %t/Mod.swiftmodule
+// RUN: %sourcekitd-test -req=doc-info -print-raw-response -module Mod -- -I %t -target %target-triple > %t/output.response
+// RUN: %diff -u %t/expected.response %t/output.response
+
+// Make sure we don't crash here, we should exclude the @abi decl from the
+// list of entities.
+
+//--- mod.swift
+@abi(func bar())
+public func foo() {}
+
+//--- expected.response
+{
+  key.annotations: [
+    {
+      key.kind: source.lang.swift.syntaxtype.attribute.builtin,
+      key.offset: 0,
+      key.length: 16
+    },
+    {
+      key.kind: source.lang.swift.syntaxtype.keyword,
+      key.offset: 17,
+      key.length: 4
+    },
+    {
+      key.kind: source.lang.swift.syntaxtype.identifier,
+      key.offset: 22,
+      key.length: 3
+    }
+  ],
+  key.entities: [
+    {
+      key.kind: source.lang.swift.decl.function.free,
+      key.name: "foo()",
+      key.usr: "s:3Mod3fooyyF",
+      key.offset: 0,
+      key.length: 27,
+      key.fully_annotated_decl: "<decl.function.free><syntaxtype.attribute.builtin>@abi(<decl.function.free><syntaxtype.keyword>func</syntaxtype.keyword> <decl.name>bar</decl.name>()</decl.function.free>)</syntaxtype.attribute.builtin> <syntaxtype.keyword>func</syntaxtype.keyword> <decl.name>foo</decl.name>()</decl.function.free>"
+    }
+  ],
+  key.sourcetext: "@abi(func bar())\nfunc foo()\n\n"
+}