[SourceKit] Add option to index local symbols to IndexSource request

Adds optional `key.should_index_locals` flag to the SourceKit `source.request.indexsource` request.
If true, the response includes symbol info for both locals and globals.

Author: Simon Barinka

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

@@ -878,6 +878,10 @@ struct RenameLocations {
   std::vector<RenameLocation> LineColumnLocs;
 };
 
+struct IndexSourceOptions {
+  bool IndexLocals = false;
+};
+
 struct IndexStoreOptions {
   std::string IndexStorePath;
   std::string IndexUnitOutputPath;
@@ -998,7 +1002,8 @@ class LangSupport {
 
   virtual void indexSource(StringRef Filename,
                            IndexingConsumer &Consumer,
-                           ArrayRef<const char *> Args) = 0;
+                           ArrayRef<const char *> Args,
+                           IndexSourceOptions Opts) = 0;
 
   virtual void indexToStore(StringRef InputFile,
                             ArrayRef<const char *> Args,

tools/SourceKit/lib/SwiftLang/SwiftIndexing.cpp

@@ -38,7 +38,7 @@ static UIdent getUIDForDependencyKind(bool isClangModule) {
 
 class SKIndexDataConsumer : public IndexDataConsumer {
 public:
-  SKIndexDataConsumer(IndexingConsumer &C) : impl(C) {}
+  SKIndexDataConsumer(IndexingConsumer &C, IndexSourceOptions Opts) : impl(C), Opts(Opts) {}
 
 private:
   void failed(StringRef error) override { impl.failed(error); }
@@ -51,6 +51,8 @@ class SKIndexDataConsumer : public IndexDataConsumer {
     return Logger::isLoggingEnabledForLevel(Logger::Level::Warning);
   }
 
+  virtual bool indexLocals() override { return Opts.IndexLocals; }
+
   bool startDependency(StringRef name, StringRef path, bool isClangModule, bool isSystem) override {
     auto kindUID = getUIDForDependencyKind(isClangModule);
     return impl.startDependency(kindUID, name, path, isSystem);
@@ -61,7 +63,7 @@ class SKIndexDataConsumer : public IndexDataConsumer {
   }
 
   Action startSourceEntity(const IndexSymbol &symbol) override {
-    if (symbol.symInfo.Kind == SymbolKind::Parameter)
+    if (!indexLocals() && symbol.symInfo.Kind == SymbolKind::Parameter)
       return Skip;
 
     // report any parent relations to this reference
@@ -204,13 +206,15 @@ class SKIndexDataConsumer : public IndexDataConsumer {
 
 private:
   IndexingConsumer &impl;
+  IndexSourceOptions Opts;
 };
 
 static void indexModule(llvm::MemoryBuffer *Input,
                         StringRef ModuleName,
                         IndexingConsumer &IdxConsumer,
                         CompilerInstance &CI,
-                        ArrayRef<const char *> Args) {
+                        ArrayRef<const char *> Args,
+                        IndexSourceOptions Opts) {
   ASTContext &Ctx = CI.getASTContext();
   std::unique_ptr<ImplicitSerializedModuleLoader> Loader;
   ModuleDecl *Mod = nullptr;
@@ -245,7 +249,7 @@ static void indexModule(llvm::MemoryBuffer *Input,
     Mod->setHasResolvedImports();
   }
 
-  SKIndexDataConsumer IdxDataConsumer(IdxConsumer);
+  SKIndexDataConsumer IdxDataConsumer(IdxConsumer, Opts);
   index::indexModule(Mod, IdxDataConsumer);
 }
 
@@ -277,7 +281,8 @@ void trace::initTraceInfo(trace::SwiftInvocation &SwiftArgs,
 
 void SwiftLangSupport::indexSource(StringRef InputFile,
                                    IndexingConsumer &IdxConsumer,
-                                   ArrayRef<const char *> OrigArgs) {
+                                   ArrayRef<const char *> OrigArgs,
+                                   IndexSourceOptions Opts) {
   std::string Error;
   auto InputBuf =
       ASTMgr->getMemoryBuffer(InputFile, llvm::vfs::getRealFileSystem(), Error);
@@ -333,7 +338,7 @@ void SwiftLangSupport::indexSource(StringRef InputFile,
     }
 
     indexModule(InputBuf.get(), llvm::sys::path::stem(Filename),
-                IdxConsumer, CI, Args);
+                IdxConsumer, CI, Args, Opts);
     return;
   }
 
@@ -365,7 +370,7 @@ void SwiftLangSupport::indexSource(StringRef InputFile,
     return;
   }
 
-  SKIndexDataConsumer IdxDataConsumer(IdxConsumer);
+  SKIndexDataConsumer IdxDataConsumer(IdxConsumer, Opts);
   index::indexSourceFile(CI.getPrimarySourceFile(), IdxDataConsumer);
 }
 

tools/SourceKit/lib/SwiftLang/SwiftLangSupport.cpp

@@ -33,6 +33,7 @@
 
 #include "clang/Lex/HeaderSearch.h"
 #include "clang/Lex/Preprocessor.h"
+#include "clang/Index/IndexSymbol.h"
 #include "llvm/ADT/APInt.h"
 #include "llvm/ADT/Hashing.h"
 #include "llvm/Support/ConvertUTF.h"
@@ -58,6 +59,7 @@ using swift::index::SymbolPropertySet;
 using swift::index::SymbolInfo;
 using swift::index::SymbolRole;
 using swift::index::SymbolRoleSet;
+using swift::index::operator&;
 
 #define KIND(NAME, CONTENT) static UIdent Kind##NAME(CONTENT);
 #include "SourceKit/Core/ProtocolUIDs.def"
@@ -711,8 +713,19 @@ UIdent SwiftLangSupport::getUIDForSymbol(SymbolInfo sym, bool isRef) {
     } else {
       return UID_FOR(FunctionFree);
     }
-  case SymbolKind::Variable:
+  case SymbolKind::Parameter:
+    if (isRef) {
+      return KindRefVarLocal;
+    } else {
+      return KindDeclVarParam;
+    }
+  case SymbolKind::Variable: {
+    bool IsLocalSymbol = sym.Properties & SymbolProperty::Local;
+    if (IsLocalSymbol) {
+      return UID_FOR(VarLocal);
+    }
     return UID_FOR(VarGlobal);
+  }
   case SymbolKind::InstanceMethod:
     return UID_FOR(MethodInstance);
   case SymbolKind::ClassMethod:

tools/SourceKit/lib/SwiftLang/SwiftLangSupport.h

@@ -561,7 +561,7 @@ class SwiftLangSupport : public LangSupport {
           Receiver) override;
 
   void indexSource(StringRef Filename, IndexingConsumer &Consumer,
-                   ArrayRef<const char *> Args) override;
+                   ArrayRef<const char *> Args, IndexSourceOptions Opts) override;
 
   void indexToStore(StringRef InputFile,
                     ArrayRef<const char *> Args,

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

@@ -725,6 +725,7 @@ static int handleTestInvocation(TestOptions Opts, TestOptions &InitOpts) {
 
   case SourceKitRequest::Index:
     sourcekitd_request_dictionary_set_uid(Req, KeyRequest, RequestIndex);
+    addRequestOptionsDirect(Req, Opts);
     break;
 
   case SourceKitRequest::CodeComplete:

tools/SourceKit/tools/sourcekitd/lib/Service/Requests.cpp

@@ -183,7 +183,7 @@ static SourceKit::Context &getGlobalContext() {
 }
 
 static sourcekitd_response_t indexSource(StringRef Filename,
-                                         ArrayRef<const char *> Args);
+                                         ArrayRef<const char *> Args, IndexSourceOptions Opts);
 
 static sourcekitd_response_t reportDocInfo(llvm::MemoryBuffer *InputBuf,
                                            StringRef ModuleName,
@@ -1417,7 +1417,12 @@ static void handleRequestIndex(const RequestDict &Req,
     SmallVector<const char *, 8> Args;
     if (getCompilerArgumentsForRequestOrEmitError(Req, Args, Rec))
       return;
-    return Rec(indexSource(*PrimaryFilePath, Args));
+
+    int64_t ShouldIndexLocals = false;
+    Req.getInt64(KeyShouldIndexLocals, ShouldIndexLocals, /*isOptional*/true);
+    IndexSourceOptions Opts{.IndexLocals = ShouldIndexLocals > 0};
+
+    return Rec(indexSource(*PrimaryFilePath, Args, Opts));
   });
 }
 
@@ -2174,11 +2179,11 @@ class SKIndexingConsumer : public IndexingConsumer {
 } // end anonymous namespace
 
 static sourcekitd_response_t indexSource(StringRef Filename,
-                                         ArrayRef<const char *> Args) {
+                                         ArrayRef<const char *> Args, IndexSourceOptions Opts) {
   ResponseBuilder RespBuilder;
   SKIndexingConsumer IdxConsumer(RespBuilder);
   LangSupport &Lang = getGlobalContext().getSwiftLangSupport();
-  Lang.indexSource(Filename, IdxConsumer, Args);
+  Lang.indexSource(Filename, IdxConsumer, Args, Opts);
 
   if (!IdxConsumer.ErrorDescription.empty())
     return createErrorRequestFailed(IdxConsumer.ErrorDescription.c_str());

utils/gyb_sourcekit_support/UIDs.py

@@ -214,6 +214,7 @@ def __init__(self, internal_name, external_name):
     KEY('ExpandedMacroReplacements', 'key.expanded_macro_replacements'),
     KEY('IndexStorePath', 'key.index_store_path'),
     KEY('IndexUnitOutputPath', 'key.index_unit_output_path'),
+    KEY('ShouldIndexLocals', 'key.should_index_locals'),
 ]