Merge pull request #66958 from DavidGoldman/main

bnbarham · web-flow · commit 6e93ac16276b · 2023-07-12T13:53:33.000-07:00
Add source.request.index_to_store to sourcekitd
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'),
 ]
 
 

Original file line number	Diff line number	Diff line change
`@@ -212,6 +212,8 @@ def __init__(self, internal_name, external_name):`
`212`	`212`	`KEY('Expansions', 'key.expansions'),`
`213`	`213`	`KEY('MacroRoles', 'key.macro_roles'),`
`214`	`214`	`KEY('ExpandedMacroReplacements', 'key.expanded_macro_replacements'),`
	`215`	`+ KEY('IndexStorePath', 'key.index_store_path'),`
	`216`	`+ KEY('IndexUnitOutputPath', 'key.index_unit_output_path'),`
`215`	`217`	`]`
`216`	`218`
`217`	`219`
`@@ -279,6 +281,7 @@ def __init__(self, internal_name, external_name):`
`279`	`281`	`REQUEST('EnableRequestBarriers', 'source.request.enable_request_barriers'),`
`280`	`282`	`REQUEST('SyntacticMacroExpansion',`
`281`	`283`	`'source.request.syntactic_macro_expansion'),`
	`284`	`+ REQUEST('IndexToStore', 'source.request.index_to_store'),`
`282`	`285`	`]`
`283`	`286`
`284`	`287`