[SourceKit] Add optional declarations array to interface gen request (#75802)

Introduces the new DeclarationsArrayBuilder and adds it to the
EditorConsumer. Declaration info always includes a kind, offset, and
length, and includes a USR where applicable.
As the USR is already available for editor.open.interface type requests,
this doesn't compute any new information, it just exposes more of what's
there already.

Author: J-MR-T

test/SourceKit/InterfaceGen/gen_swift_module.swift

@@ -33,3 +33,67 @@ func f(s : inout [Int]) {
 // RUN: %swift -emit-module -o /dev/null -emit-module-interface-path %t.mod/swift_mod.swiftinterface -O %S/Inputs/swift_mod.swift -parse-as-library -disable-implicit-concurrency-module-import -disable-implicit-string-processing-module-import
 // RUN: %sourcekitd-test -req=interface-gen -module swift_mod -- -Xfrontend -disable-implicit-concurrency-module-import -Xfrontend -disable-implicit-string-processing-module-import  -I %t.mod -module-cache-path %t.mod/mcp > %t.response
 // RUN: %diff -u %s.from_swiftinterface.response %t.response
+
+// Separately test whether the optional declarations array contains the expected information from Inputs/swift_mod.swift
+// RUN: %sourcekitd-test -req=interface-gen -req-opts=enabledeclarations=true -module swift_mod -- -Xfrontend -disable-implicit-concurrency-module-import -Xfrontend -disable-implicit-string-processing-module-import  -I %t.mod -module-cache-path %t.mod/mcp | %FileCheck -check-prefix=CHECK-DECLARATIONS-ARRAY %s
+// check that the other arrays are still there as normal, by checking for a unique identifying member in each
+// CHECK-DECLARATIONS-ARRAY-DAG: key.kind:{{.*}}ref.associatedtype
+// CHECK-DECLARATIONS-ARRAY-DAG: key.kind:{{.*}}syntaxtype
+// CHECK-DECLARATIONS-ARRAY-DAG: key.substructure
+
+// CHECK-DECLARATIONS-ARRAY: {{^}}]{{$}}
+
+// public class MyClass
+// CHECK-DECLARATIONS-ARRAY: key.kind:
+// CHECK-DECLARATIONS-ARRAY-SAME: decl
+// CHECK-DECLARATIONS-ARRAY-SAME: class
+// CHECK-DECLARATIONS-ARRAY-NEXT: key.usr:
+// could check for parts of the usr, but as the .response files always check for exact matches, do the same here
+// CHECK-DECLARATIONS-ARRAY-SAME: s:9swift_mod7MyClassC
+// check that offsets are there, but not zero, as that might mean offsets are broken (none of the declaration offsets in this file should be zero)
+// CHECK-DECLARATIONS-ARRAY-NOT: key.offset:{{.*}}{{^[0-9]}}0{{^[0-9]}}
+// CHECK-DECLARATIONS-ARRAY: key.offset
+
+// public func pub_method
+// CHECK-DECLARATIONS-ARRAY: key.kind:
+// CHECK-DECLARATIONS-ARRAY-SAME: decl
+// CHECK-DECLARATIONS-ARRAY-SAME: method.instance
+// CHECK-DECLARATIONS-ARRAY-NEXT: key.usr:
+// CHECK-DECLARATIONS-ARRAY-SAME: s:9swift_mod7MyClassC10pub_methodyyF
+// CHECK-DECLARATIONS-ARRAY-NOT: key.offset:{{.*}}{{^[0-9]}}0{{^[0-9]}}
+// CHECK-DECLARATIONS-ARRAY: key.offset
+
+// CHECK-DECLARATIONS-ARRAY-NOT: int_method
+// CHECK-DECLARATIONS-ARRAY-NOT: fp_method
+// CHECK-DECLARATIONS-ARRAY-NOT: priv_method
+
+// public protocol MyProto
+// CHECK-DECLARATIONS-ARRAY: key.kind:
+// CHECK-DECLARATIONS-ARRAY-SAME: decl
+// CHECK-DECLARATIONS-ARRAY-SAME: protocol
+// CHECK-DECLARATIONS-ARRAY-NEXT: key.usr
+// CHECK-DECLARATIONS-ARRAY-SAME: s:9swift_mod7MyProtoP
+// CHECK-DECLARATIONS-ARRAY-NOT: key.offset:{{.*}}{{^[0-9]}}0{{^[0-9]}}
+// CHECK-DECLARATIONS-ARRAY: key.offset
+
+// associatedtype Assoc
+// CHECK-DECLARATIONS-ARRAY: key.kind:
+// CHECK-DECLARATIONS-ARRAY-SAME: decl
+// CHECK-DECLARATIONS-ARRAY-SAME: associatedtype
+// CHECK-DECLARATIONS-ARRAY-NEXT: key.usr
+// CHECK-DECLARATIONS-ARRAY-SAME: s:9swift_mod7MyProtoP5AssocQa
+// CHECK-DECLARATIONS-ARRAY-NOT: key.offset:{{.*}}{{^[0-9]}}0{{^[0-9]}}
+// CHECK-DECLARATIONS-ARRAY: key.offset
+
+// public func pub_function
+// CHECK-DECLARATIONS-ARRAY: key.kind:
+// CHECK-DECLARATIONS-ARRAY-SAME: decl
+// CHECK-DECLARATIONS-ARRAY-SAME: function.free
+// CHECK-DECLARATIONS-ARRAY-NEXT: key.usr
+// CHECK-DECLARATIONS-ARRAY-SAME: s:9swift_mod12pub_functionSiyF
+// CHECK-DECLARATIONS-ARRAY-NOT: key.offset:{{.*}}{{^[0-9]}}0{{^[0-9]}}
+// CHECK-DECLARATIONS-ARRAY: key.offset
+
+// CHECK-DECLARATIONS-ARRAY-NOT: int_function
+// CHECK-DECLARATIONS-ARRAY-NOT: fp_function
+// CHECK-DECLARATIONS-ARRAY-NOT: priv_function

tools/SourceKit/docs/Protocol.md

@@ -362,16 +362,19 @@ Welcome to SourceKit.  Type ':help' for assistance.
     <key.modulename>:       (string) // Full module name, e.g. "Foundation.NSArray"
     [opt] <key.compilerargs> [string*] // array of zero or more strings for the compiler arguments
                                        // e.g ["-sdk", "/path/to/sdk"]
+    [opt] <key.enabledeclarations> (int) // 0 by default, 1 to enable the declarations array in the output
 }
 ```
 
 ### Response
 
 This will return the Swift interface of the specified module.
 
-- `key.sourcetext`: The pretty-printed module interface in swift source code
+- `key.sourcetext`: The pretty-printed module interface in swift source code.
 - `key.syntaxmap`: An array of syntactic annotations, same as the one returned for the source.request.editor.open request.
 - `key.annotations`: An array of semantic annotations, same as the one returned for the source.request.editor.open request.
+- `key.substructure`: An array of dictionaries representing ranges of structural elements in the result description, such as the parameters of a function.
+- (optional, only if `key.enabledeclarations: 1`) `key.declarations`: An array of declarations, containing the kind, USR (if available), offset, and length of the declaration.
 
 All SourceKit requests that don't modify the source buffer should work on the
 opened document, by passing the associated 'name' for the document.

tools/SourceKit/include/SourceKit/Core/LangSupport.h

@@ -391,6 +391,9 @@ class EditorConsumer {
   virtual void handleSemanticAnnotation(unsigned Offset, unsigned Length,
                                         UIdent Kind, bool isSystem) = 0;
 
+  virtual void handleDeclaration(unsigned Offset, unsigned Length, UIdent Kind,
+                                 StringRef USR) = 0;
+
   virtual void beginDocumentSubStructure(unsigned Offset, unsigned Length,
                                          UIdent Kind, UIdent AccessLevel,
                                          UIdent SetterAccessLevel,

tools/SourceKit/lib/SwiftLang/SwiftEditorInterfaceGen.cpp

@@ -56,6 +56,8 @@ class SwiftInterfaceGenContext::Implementation {
     const Decl *Dcl = nullptr;
     /// The range in the interface source.
     TextRange Range;
+    /// The USR, if the declaration has one
+    StringRef USR;
 
     TextDecl(const Decl *D, TextRange Range)
       : Dcl(D), Range(Range) {}
@@ -66,6 +68,8 @@ class SwiftInterfaceGenContext::Implementation {
     std::string Text;
     std::vector<TextReference> References;
     std::vector<TextDecl> Decls;
+    /// Do not clear the map or erase elements,
+    /// without keeping the Decls vector in sync
     llvm::StringMap<TextDecl> USRMap;
   };
 
@@ -155,8 +159,17 @@ class AnnotatingPrinter : public StreamPrinter {
           OS << TargetUSR;
         }
         StringRef USR = OS.str();
-        Info.USRMap[USR] = Entry;
         DeclUSRs.emplace_back(VD, USR);
+        auto iterator = Info.USRMap.insert_or_assign(USR, Entry).first;
+        // Set the USR in the declarations to the key in the USRMap, because the
+        // lifetime of that matches/exceeds the lifetime of Decls. String keys
+        // in the StringMap are heap allocated and only get destroyed on
+        // explicit erase() or clear() calls, or on destructor calls (the
+        // Programmer's Manual description itself also states that StringMap
+        // "only ever copies a string if a value is inserted").
+        // Thus this never results in a dangling reference, as the USRMap is
+        // never cleared and no elements are erased in its lifetime.
+        Info.Decls.back().USR = iterator->getKey();
       }
     }
   }
@@ -265,6 +278,21 @@ static void reportSemanticAnnotations(const SourceTextInfo &IFaceInfo,
   }
 }
 
+/// Create the declarations array (sourcekitd::DeclarationsArrayBuilder) from
+/// the SourceTextInfo about declarations
+static void reportDeclarations(const SourceTextInfo &IFaceInfo,
+                               EditorConsumer &Consumer) {
+  for (auto &Dcl : IFaceInfo.Decls) {
+    if (!Dcl.Dcl)
+      continue;
+    UIdent Kind = SwiftLangSupport::getUIDForDecl(Dcl.Dcl);
+    if (Kind.isInvalid())
+      continue;
+    Consumer.handleDeclaration(Dcl.Range.Offset, Dcl.Range.Length, Kind,
+                               Dcl.USR);
+  }
+}
+
 namespace {
 /// A diagnostic consumer that picks up module loading errors.
 class ModuleLoadingErrorConsumer final : public DiagnosticConsumer {
@@ -593,6 +621,8 @@ void SwiftInterfaceGenContext::reportEditorInfo(EditorConsumer &Consumer) const
   reportSyntacticAnnotations(Impl.TextCI, Consumer);
   reportDocumentStructure(Impl.TextCI, Consumer);
   reportSemanticAnnotations(Impl.Info, Consumer);
+
+  reportDeclarations(Impl.Info, Consumer);
   Consumer.finished();
 }
 

tools/SourceKit/tools/sourcekitd-test/sourcekitd-test.cpp

@@ -1069,6 +1069,9 @@ static int handleTestInvocation(TestOptions Opts, TestOptions &InitOpts) {
                                                Opts.InterestedUSR.c_str());
     if (!Opts.USR.empty())
       sourcekitd_request_dictionary_set_string(Req, KeyUSR, Opts.USR.c_str());
+
+    // add possible EnableDeclarations in particular
+    addRequestOptionsDirect(Req, Opts);
     break;
 
   case SourceKitRequest::FindInterfaceDoc:
@@ -2453,6 +2456,12 @@ static void printInterfaceGen(sourcekitd_variant_t Info, bool CheckASCII) {
   sourcekitd_variant_t structure =
       sourcekitd_variant_dictionary_get_value(Info, KeySubStructure);
   printRawVariant(structure);
+  sourcekitd_variant_t declarations =
+      sourcekitd_variant_dictionary_get_value(Info, KeyDeclarations);
+  // only output declarations if there are any (because this might have been
+  // disabled in the request itself)
+  if (sourcekitd_variant_get_type(declarations) != SOURCEKITD_VARIANT_TYPE_NULL)
+    printRawVariant(declarations);
 }
 
 static void printRelatedIdents(sourcekitd_variant_t Info, StringRef Filename,

tools/SourceKit/tools/sourcekitd/include/sourcekitd/DeclarationsArray.h

@@ -0,0 +1,45 @@
+//===--- DeclarationsArray.h - ----------------------------------*- C++ -*-===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2024 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+// This is an array used in the response to editor.open.interface requests.
+// It contains all declarations, identified by their Kind, Offset, and Length,
+// and optionally includes a USR, if the declaration has one.
+//===----------------------------------------------------------------------===//
+#ifndef LLVM_SOURCEKITD_DECLARATIONS_ARRAY_H
+#define LLVM_SOURCEKITD_DECLARATIONS_ARRAY_H
+
+#include "sourcekitd/Internal.h"
+
+namespace sourcekitd {
+
+VariantFunctions *getVariantFunctionsForDeclarationsArray();
+
+/// Builds an array for declarations by kind, offset, length, and optionally USR
+class DeclarationsArrayBuilder {
+public:
+  DeclarationsArrayBuilder();
+  ~DeclarationsArrayBuilder();
+
+  void add(SourceKit::UIdent Kind, unsigned Offset, unsigned Length,
+           llvm::StringRef USR);
+
+  bool empty() const;
+
+  std::unique_ptr<llvm::MemoryBuffer> createBuffer();
+
+private:
+  struct Implementation;
+  Implementation &Impl;
+};
+
+} // namespace sourcekitd
+
+#endif

tools/SourceKit/tools/sourcekitd/include/sourcekitd/Internal.h

@@ -49,6 +49,7 @@ static const unsigned ProtocolMinorVersion = 0;
 
 enum class CustomBufferKind {
   TokenAnnotationsArray,
+  DeclarationsArray,
   DocSupportAnnotationArray,
   CodeCompletionResultsArray,
   DocStructureArray,

tools/SourceKit/tools/sourcekitd/lib/API/CMakeLists.txt

@@ -9,6 +9,7 @@ add_sourcekit_library(sourcekitdAPI
   DocSupportAnnotationArray.cpp
   RawData.cpp
   sourcekitdAPI-Common.cpp
+  DeclarationsArray.cpp
   TokenAnnotationsArray.cpp
   ExpressionTypeArray.cpp
   VariableTypeArray.cpp