Add source.request.index_to_store to sourcekitd

DavidGoldman · DavidGoldman · commit c0402fa9897b · 2023-07-13T12:29:33.000-04:00
This requests performs an index store index of the given file using the
given index store path and index unit output path. All other options are
derived from the index store related compiler flags.

This will allow IDEs like Xcode to index the file directly inside of
sourcekitd and potentially reuse an already built AST.
diff --git a/test/SourceKit/Indexing/Inputs/indexstore_multifile_other.swift b/test/SourceKit/Indexing/Inputs/indexstore_multifile_other.swift
@@ -0,0 +1,4 @@
+struct OtherStruct {
+  let x: Int
+  let y: Int
+}
diff --git a/test/SourceKit/Indexing/indexstore_multifile.swift b/test/SourceKit/Indexing/indexstore_multifile.swift
@@ -0,0 +1,23 @@
+// RUN: %empty-directory(%t)
+// RUN: %sourcekitd-test -req=index-to-store %s -index-store-path %t/idx -index-unit-output-path %t/indexstore_multifile.o -- -c -module-name indexstoremodule -o args_output_path.o %s %S/Inputs/indexstore_multifile_other.swift
+// RUN: c-index-test core -print-unit %t/idx | %FileCheck  --dump-input=always --dump-input-filter=all %s
+
+struct Foo {
+  let bar: Int
+  let counter: Int
+  let other: OtherStruct
+
+  init(bar: Int, counter: Int = 0, other: OtherStruct) {
+    self.bar = bar
+    self.counter = counter
+    self.other = other
+  }
+}
+
+// CHECK: indexstore_multifile.o-{{.*}}
+// CHECK: module-name: indexstoremodule
+// CHECK: main-path: SOURCE_DIR{{/|\\}}test{{/|\\}}SourceKit{{/|\\}}Indexing{{/|\\}}indexstore_multifile.swift
+// CHECK: out-file: BUILD_DIR{{.*}}indexstore_multifile.o
+// CHECK: is-debug: 1
+// CHECK: Unit | system | Swift | {{.*}}{{/|\\}}Swift.swiftmodule
+// CHECK: Record | user | SOURCE_DIR{{/|\\}}test{{/|\\}}SourceKit{{/|\\}}Indexing{{/|\\}}indexstore_multifile.swift | indexstore_multifile.swift-{{.*}}
diff --git a/test/SourceKit/Indexing/indexstore_with_remappings.swift b/test/SourceKit/Indexing/indexstore_with_remappings.swift
@@ -0,0 +1,17 @@
+// RUN: %empty-directory(%t)
+// RUN: %sourcekitd-test -req=index-to-store %s -index-store-path %t/idx -index-unit-output-path %t/indexstore_with_remappings.o -- -c -module-name indexstoremodule -file-prefix-map %S=REMAPPED_SRC_ROOT -file-prefix-map %t=REMAPPED_OUT_DIR %s
+// RUN: c-index-test core -print-unit %t/idx | %FileCheck %s
+
+class Foo {
+  let member: Int
+  init(member: Int) {
+    self.member = member
+  }
+}
+
+// CHECK: indexstore_with_remappings.o-{{.*}}
+// CHECK: module-name: indexstoremodule
+// CHECK: main-path: REMAPPED_SRC_ROOT{{/|\\}}indexstore_with_remappings.swift
+// CHECK: out-file: REMAPPED_OUT_DIR{{/|\\}}indexstore_with_remappings.o
+// CHECK: is-debug: 1
+// CHECK: Unit | system | Swift | {{.*}}{{/|\\}}Swift.swiftmodule
diff --git a/tools/SourceKit/include/SourceKit/Core/LangSupport.h b/tools/SourceKit/include/SourceKit/Core/LangSupport.h
@@ -873,10 +873,18 @@ struct RenameLocations {
   std::vector<RenameLocation> LineColumnLocs;
 };
 
+struct IndexStoreOptions {
+  std::string IndexStorePath;
+  std::string IndexUnitOutputPath;
+};
+
+struct IndexStoreInfo{};
+
 typedef std::function<void(RequestResult<ArrayRef<CategorizedEdits>> Result)>
     CategorizedEditsReceiver;
 typedef std::function<void(RequestResult<ArrayRef<CategorizedRenameRanges>> Result)>
     CategorizedRenameRangesReceiver;
+typedef std::function<void(RequestResult<IndexStoreInfo> Result)> IndexToStoreReceiver;
 
 class DocInfoConsumer {
   virtual void anchor();
@@ -987,6 +995,12 @@ class LangSupport {
                            IndexingConsumer &Consumer,
                            ArrayRef<const char *> Args) = 0;
 
+  virtual void indexToStore(StringRef InputFile,
+                            ArrayRef<const char *> Args,
+                            IndexStoreOptions Opts,
+                            SourceKitCancellationToken CancellationToken,
+                            IndexToStoreReceiver Receiver) = 0;
+
   virtual void codeComplete(llvm::MemoryBuffer *InputBuf, unsigned Offset,
                             OptionsDictionary *options,
                             CodeCompletionConsumer &Consumer,
diff --git a/tools/SourceKit/lib/SwiftLang/SwiftIndexing.cpp b/tools/SourceKit/lib/SwiftLang/SwiftIndexing.cpp
@@ -19,6 +19,7 @@
 #include "swift/Frontend/Frontend.h"
 #include "swift/Frontend/PrintingDiagnosticConsumer.h"
 #include "swift/Index/Index.h"
+#include "swift/Index/IndexRecord.h"
 #include "swift/Serialization/SerializedModuleLoader.h"
 // This is included only for createLazyResolver(). Move to different header ?
 #include "swift/Sema/IDETypeChecking.h"
@@ -367,3 +368,79 @@ void SwiftLangSupport::indexSource(StringRef InputFile,
   SKIndexDataConsumer IdxDataConsumer(IdxConsumer);
   index::indexSourceFile(CI.getPrimarySourceFile(), IdxDataConsumer);
 }
+
+static void emitIndexDataForSourceFile(SourceFile &PrimarySourceFile,
+                                       StringRef IndexStorePath,
+                                       StringRef IndexUnitOutputPath,
+                                       const CompilerInstance &Instance) {
+  const auto &Invocation = Instance.getInvocation();
+  const auto &Opts = Invocation.getFrontendOptions();
+
+  bool isDebugCompilation;
+  switch (Invocation.getSILOptions().OptMode) {
+    case OptimizationMode::NotSet:
+    case OptimizationMode::NoOptimization:
+      isDebugCompilation = true;
+      break;
+    case OptimizationMode::ForSpeed:
+    case OptimizationMode::ForSize:
+      isDebugCompilation = false;
+      break;
+  }
+
+  (void) index::indexAndRecord(&PrimarySourceFile, IndexUnitOutputPath,
+                               IndexStorePath,
+                               !Opts.IndexIgnoreClangModules,
+                               Opts.IndexSystemModules,
+                               Opts.IndexIgnoreStdlib,
+                               Opts.IndexIncludeLocals,
+                               isDebugCompilation,
+                               Invocation.getTargetTriple(),
+                               *Instance.getDependencyTracker(),
+                               Invocation.getIRGenOptions().FilePrefixMap);
+}
+
+void SwiftLangSupport::indexToStore(
+    StringRef PrimaryFilePath, ArrayRef<const char *> Args,
+    IndexStoreOptions Opts,
+    SourceKitCancellationToken CancellationToken,
+    IndexToStoreReceiver Receiver) {
+  std::string Error;
+  SwiftInvocationRef Invok =
+      ASTMgr->getTypecheckInvocation(Args, PrimaryFilePath, Error);
+  if (!Invok) {
+    LOG_WARN_FUNC("failed to create an ASTInvocation: " << Error);
+    Receiver(RequestResult<IndexStoreInfo>::fromError(Error));
+    return;
+  }
+
+  struct IndexStoreASTConsumer : public SwiftASTConsumer {
+    IndexToStoreReceiver Receiver;
+    IndexStoreOptions Opts;
+
+    IndexStoreASTConsumer(IndexToStoreReceiver Receiver, IndexStoreOptions Opts)
+        : Receiver(std::move(Receiver)), Opts(std::move(Opts)) {}
+
+    void handlePrimaryAST(ASTUnitRef AstUnit) override {
+      auto &SF = AstUnit->getPrimarySourceFile();
+      auto &CI = AstUnit->getCompilerInstance();
+      emitIndexDataForSourceFile(
+          SF, Opts.IndexStorePath, Opts.IndexUnitOutputPath, CI);
+      Receiver(RequestResult<IndexStoreInfo>::fromResult(IndexStoreInfo{}));
+    }
+
+    void cancelled() override {
+      Receiver(RequestResult<IndexStoreInfo>::cancelled());
+    }
+
+    void failed(StringRef Error) override {
+      Receiver(RequestResult<IndexStoreInfo>::fromError(Error));
+    }
+  };
+
+  auto ASTConsumer = std::make_shared<IndexStoreASTConsumer>(std::move(Receiver), std::move(Opts));
+  getASTManager()->processASTAsync(Invok, ASTConsumer,
+                                   /*OncePerASTToken=*/nullptr,
+                                   CancellationToken,
+                                   llvm::vfs::getRealFileSystem());
+}
diff --git a/tools/SourceKit/lib/SwiftLang/SwiftLangSupport.h b/tools/SourceKit/lib/SwiftLang/SwiftLangSupport.h
@@ -563,6 +563,12 @@ class SwiftLangSupport : public LangSupport {
   void indexSource(StringRef Filename, IndexingConsumer &Consumer,
                    ArrayRef<const char *> Args) override;
 
+  void indexToStore(StringRef InputFile,
+                    ArrayRef<const char *> Args,
+                    IndexStoreOptions Opts,
+                    SourceKitCancellationToken CancellationToken,
+                    IndexToStoreReceiver Receiver) override;
+
   void codeComplete(llvm::MemoryBuffer *InputBuf, unsigned Offset,
                     OptionsDictionary *options,
                     SourceKit::CodeCompletionConsumer &Consumer,
diff --git a/tools/SourceKit/tools/sourcekitd-test/Options.td b/tools/SourceKit/tools/sourcekitd-test/Options.td
@@ -166,6 +166,12 @@ def vfs_files : CommaJoined<["-"], "vfs-files=">,
 def vfs_name : Separate<["-"], "vfs-name">,
   HelpText<"Specify a virtual filesystem name">;
 
+def index_store_path : Separate<["-"], "index-store-path">,
+  HelpText<"Index Store path to use for indexing">;
+
+def index_unit_output_path : Separate<["-"], "index-unit-output-path">,
+  HelpText<"Index Store unit output path to use for indexing">;
+
 def suppress_config_request : Flag<["-"], "suppress-config-request">,
   HelpText<"Suppress the default global configuration request, that is otherwise sent before any other request (except for the global-config request itself)">;
 
diff --git a/tools/SourceKit/tools/sourcekitd-test/TestOptions.cpp b/tools/SourceKit/tools/sourcekitd-test/TestOptions.cpp
@@ -154,6 +154,7 @@ bool TestOptions::parseArgs(llvm::ArrayRef<const char *> Args) {
         .Case("compile", SourceKitRequest::Compile)
         .Case("compile.close", SourceKitRequest::CompileClose)
         .Case("syntactic-expandmacro", SourceKitRequest::SyntacticMacroExpansion)
+        .Case("index-to-store", SourceKitRequest::IndexToStore)
 #define SEMANTIC_REFACTORING(KIND, NAME, ID) .Case("refactoring." #ID, SourceKitRequest::KIND)
 #include "swift/Refactoring/RefactoringKinds.def"
         .Default(SourceKitRequest::None);
@@ -205,6 +206,7 @@ bool TestOptions::parseArgs(llvm::ArrayRef<const char *> Args) {
                      << "- global-config\n"
                      << "- dependency-updated\n"
                      << "- syntactic-expandmacro\n"
+                     << "- index-to-store\n"
 #define SEMANTIC_REFACTORING(KIND, NAME, ID) << "- refactoring." #ID "\n"
 #include "swift/Refactoring/RefactoringKinds.def"
                         "\n";
@@ -428,6 +430,14 @@ bool TestOptions::parseArgs(llvm::ArrayRef<const char *> Args) {
       VFSName = InputArg->getValue();
       break;
 
+    case OPT_index_store_path:
+      IndexStorePath = InputArg->getValue();
+      break;
+
+    case OPT_index_unit_output_path:
+      IndexUnitOutputPath = InputArg->getValue();
+      break;
+
     case OPT_module_cache_path:
       ModuleCachePath = InputArg->getValue();
       break;
diff --git a/tools/SourceKit/tools/sourcekitd-test/TestOptions.h b/tools/SourceKit/tools/sourcekitd-test/TestOptions.h
@@ -71,6 +71,7 @@ enum class SourceKitRequest {
   Compile,
   CompileClose,
   SyntacticMacroExpansion,
+  IndexToStore,
 #define SEMANTIC_REFACTORING(KIND, NAME, ID) KIND,
 #include "swift/Refactoring/RefactoringKinds.def"
 };
@@ -148,6 +149,8 @@ struct TestOptions {
   llvm::Optional<std::string> VFSName;
   llvm::Optional<bool> CancelOnSubsequentRequest;
   bool ShellExecution = false;
+  std::string IndexStorePath;
+  std::string IndexUnitOutputPath;
   bool parseArgs(llvm::ArrayRef<const char *> Args);
   void printHelp(bool ShowHidden) const;
 };
diff --git a/tools/SourceKit/tools/sourcekitd-test/sourcekitd-test.cpp b/tools/SourceKit/tools/sourcekitd-test/sourcekitd-test.cpp
@@ -1153,6 +1153,13 @@ static int handleTestInvocation(TestOptions Opts, TestOptions &InitOpts) {
                                           RequestSyntacticMacroExpansion);
     setSyntacticMacroExpansions(Req, Opts, SourceBuf.get());
     break;
+
+  case SourceKitRequest::IndexToStore:
+    sourcekitd_request_dictionary_set_string(Req, KeyName, SemaName.c_str());
+    sourcekitd_request_dictionary_set_string(Req, KeyIndexStorePath, Opts.IndexStorePath.c_str());
+    sourcekitd_request_dictionary_set_string(Req, KeyIndexUnitOutputPath, Opts.IndexUnitOutputPath.c_str());
+    sourcekitd_request_dictionary_set_uid(Req, KeyRequest, RequestIndexToStore);
+    break;
   }
 
   if (!Opts.SourceFile.empty()) {
@@ -1585,6 +1592,9 @@ static bool handleResponse(sourcekitd_response_t Resp, const TestOptions &Opts,
       case SourceKitRequest::Statistics:
         printStatistics(Info, llvm::outs());
         break;
+      case SourceKitRequest::IndexToStore:
+        printRawResponse(Resp);
+        break;
     }
   }
 
diff --git a/tools/SourceKit/tools/sourcekitd/lib/Service/Requests.cpp b/tools/SourceKit/tools/sourcekitd/lib/Service/Requests.cpp
@@ -1421,6 +1421,48 @@ static void handleRequestIndex(const RequestDict &Req,
   });
 }
 
+static void handleRequestIndexToStore(
+    const RequestDict &Req, SourceKitCancellationToken CancellationToken,
+    ResponseReceiver Rec) {
+  if (checkVFSNotSupported(Req, Rec))
+    return;
+
+  handleSemanticRequest(Req, Rec, [Req, Rec, CancellationToken]() {
+    Optional<StringRef> PrimaryFilePath =
+        getPrimaryFilePathForRequestOrEmitError(Req, Rec);
+    if (!PrimaryFilePath)
+      return;
+
+    IndexStoreOptions Opts;
+    if (auto IndexStorePath = Req.getString(KeyIndexStorePath))
+      Opts.IndexStorePath = IndexStorePath->str();
+    else
+      return Rec(createErrorRequestInvalid("'key.index_store_path' is required"));
+
+    if (auto IndexUnitOutputPath = Req.getString(KeyIndexUnitOutputPath))
+      Opts.IndexUnitOutputPath = IndexUnitOutputPath->str();
+    else
+      return Rec(createErrorRequestInvalid("'key.index_unit_output_path' is required"));
+
+    SmallVector<const char *, 0> Args;
+    if (getCompilerArgumentsForRequestOrEmitError(Req, Args, Rec))
+      return;
+    LangSupport &Lang = getGlobalContext().getSwiftLangSupport();
+    Lang.indexToStore(*PrimaryFilePath, Args, std::move(Opts),
+      CancellationToken,
+      [Rec](const RequestResult<IndexStoreInfo> &Result) {
+        if (Result.isCancelled())
+          return Rec(createErrorRequestCancelled());
+        if (Result.isError())
+          return Rec(createErrorRequestFailed(Result.getError()));
+
+        ResponseBuilder RespBuilder;
+        return Rec(RespBuilder.createResponse());
+      });
+  });
+}
+
+
 static void
 handleRequestCursorInfo(const RequestDict &Req,
                         SourceKitCancellationToken CancellationToken,
@@ -2044,6 +2086,7 @@ void handleRequestImpl(sourcekitd_object_t ReqObj,
   HANDLE_REQUEST(RequestConformingMethodList, handleRequestConformingMethodList)
 
   HANDLE_REQUEST(RequestIndex, handleRequestIndex)
+  HANDLE_REQUEST(RequestIndexToStore, handleRequestIndexToStore)
   HANDLE_REQUEST(RequestCursorInfo, handleRequestCursorInfo)
   HANDLE_REQUEST(RequestRangeInfo, handleRequestRangeInfo)
   HANDLE_REQUEST(RequestSemanticRefactoring, handleRequestSemanticRefactoring)
diff --git a/utils/gyb_sourcekit_support/UIDs.py b/utils/gyb_sourcekit_support/UIDs.py
@@ -212,6 +212,8 @@ def __init__(self, internal_name, external_name):
     KEY('Expansions', 'key.expansions'),
     KEY('MacroRoles', 'key.macro_roles'),
     KEY('ExpandedMacroReplacements', 'key.expanded_macro_replacements'),
+    KEY('IndexStorePath', 'key.index_store_path'),
+    KEY('IndexUnitOutputPath', 'key.index_unit_output_path'),
 ]
 
 
@@ -279,6 +281,7 @@ def __init__(self, internal_name, external_name):
     REQUEST('EnableRequestBarriers', 'source.request.enable_request_barriers'),
     REQUEST('SyntacticMacroExpansion',
             'source.request.syntactic_macro_expansion'),
+    REQUEST('IndexToStore', 'source.request.index_to_store'),
 ]
 
 
