Merge pull request swiftlang#22753 from nkcsgexi/request-type

SourceKit: Add a new request to collect expression types in a source file.

Author: nkcsgexi

test/SourceKit/ExpressionType/basic.swift

@@ -0,0 +1,4 @@
+func foo() { return 1 }
+
+// RUN: %sourcekitd-test -req=collect-type %s -- %s | %FileCheck %s
+// CHECK: (20, 21): Int

tools/SourceKit/docs/Protocol.md

@@ -737,7 +737,46 @@ Welcome to SourceKit.  Type ':help' for assistance.
 }
 ```
 
+## Expression Type
+This request collects the types of all expressions in a source file after type checking.
+To fulfill this task, the client must provide the path to the Swift source file under
+type checking and the necessary compiler arguments to help resolve all dependencies.
 
+### Request
+
+```
+{
+    <key.request>:            (UID)     <source.request.expression.type>,
+    <key.sourcefile>:         (string)  // Absolute path to the file.
+    <key.compilerargs>:       [string*] // Array of zero or more strings for the compiler arguments,
+                                        // e.g ["-sdk", "/path/to/sdk"]. If key.sourcefile is provided,
+                                        // these must include the path to that file.
+}
+```
+
+### Response
+```
+{
+    <key.printedtypebuffer>:          (string)                    // A text buffer where all expression types are printed to.
+    <key.expression_type_list>:       (array) [expr-type-info*]   // A list of expression and type
+}
+```
+
+```
+expr-type-info ::=
+{
+  <key.expression_offset>:    (int64)    // Offset of an expression in the source file
+  <key.expression_length>:    (int64)    // Length of an expression in the source file
+  <key.type_offset>:          (int64)    // Offset of the printed type of the expression in the printed type buffer
+  <key.type_length>:          (int64)    // Length of the printed type of the expression in the printed type buffer
+}
+```
+
+### Testing
+
+```
+$ sourcekitd-test -req=collect-type /path/to/file.swift -- /path/to/file.swift
+```
 
 # UIDs
 

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

@@ -123,6 +123,18 @@ struct CodeCompletionInfo {
   Optional<ArrayRef<ParameterStructure>> parametersStructure;
 };
 
+struct ExpressionType {
+  unsigned ExprOffset;
+  unsigned ExprLength;
+  unsigned TypeOffset;
+  unsigned TypeLength;
+};
+
+struct ExpressionTypesInFile {
+  std::vector<ExpressionType> Results;
+  StringRef TypeBuffer;
+};
+
 class CodeCompletionConsumer {
   virtual void anchor();
 
@@ -708,6 +720,10 @@ class LangSupport {
                                    ArrayRef<const char*> Args,
                                    CategorizedEditsReceiver Receiver) = 0;
 
+  virtual void collectExpressionTypes(StringRef FileName,
+                                      ArrayRef<const char *> Args,
+                                      std::function<void(const ExpressionTypesInFile&)> Receiver) = 0;
+
   virtual void getDocInfo(llvm::MemoryBuffer *InputBuf,
                           StringRef ModuleName,
                           ArrayRef<const char *> Args,
@@ -726,7 +742,6 @@ class LangSupport {
 
   virtual void getStatistics(StatisticsReceiver) = 0;
 };
-
 } // namespace SourceKit
 
 #endif

tools/SourceKit/lib/SwiftLang/SwiftLangSupport.h

@@ -521,6 +521,9 @@ class SwiftLangSupport : public LangSupport {
                              unsigned Length, ArrayRef<const char *> Args,
                              CategorizedRenameRangesReceiver Receiver) override;
 
+  void collectExpressionTypes(StringRef FileName, ArrayRef<const char *> Args,
+                              std::function<void(const ExpressionTypesInFile&)> Receiver) override;
+
   void semanticRefactoring(StringRef Filename, SemanticRefactoringInfo Info,
                            ArrayRef<const char*> Args,
                            CategorizedEditsReceiver Receiver) override;

tools/SourceKit/lib/SwiftLang/SwiftSourceDocInfo.cpp

@@ -2032,3 +2032,49 @@ semanticRefactoring(StringRef Filename, SemanticRefactoringInfo Info,
   static const char OncePerASTToken = 0;
   getASTManager()->processASTAsync(Invok, std::move(Consumer), &OncePerASTToken);
 }
+
+void SwiftLangSupport::collectExpressionTypes(StringRef FileName,
+                                              ArrayRef<const char *> Args,
+                  std::function<void(const ExpressionTypesInFile&)> Receiver) {
+  std::string Error;
+  SwiftInvocationRef Invok = ASTMgr->getInvocation(Args, FileName, Error);
+  if (!Invok) {
+    // FIXME: Report it as failed request.
+    LOG_WARN_FUNC("failed to create an ASTInvocation: " << Error);
+    Receiver({});
+    return;
+  }
+  assert(Invok);
+  class ExpressionTypeCollector: public SwiftASTConsumer {
+    std::function<void(const ExpressionTypesInFile&)> Receiver;
+  public:
+    ExpressionTypeCollector(std::function<void(const ExpressionTypesInFile&)> Receiver):
+      Receiver(std::move(Receiver)) {}
+    void handlePrimaryAST(ASTUnitRef AstUnit) override {
+      auto *SF = AstUnit->getCompilerInstance().getPrimarySourceFile();
+      std::vector<ExpressionTypeInfo> Scratch;
+      llvm::SmallString<256> TypeBuffer;
+      llvm::raw_svector_ostream OS(TypeBuffer);
+      ExpressionTypesInFile Result;
+      for (auto Item: collectExpressionType(*SF, Scratch, OS)) {
+        Result.Results.push_back({Item.offset, Item.length, Item.typeOffset,
+          Item.typeLength});
+      }
+      Result.TypeBuffer = OS.str();
+      Receiver(Result);
+    }
+
+    void cancelled() override {
+      Receiver({});
+    }
+
+    void failed(StringRef Error) override {
+      Receiver({});
+    }
+  };
+  auto Collector = std::make_shared<ExpressionTypeCollector>(Receiver);
+  /// FIXME: When request cancellation is implemented and Xcode adopts it,
+  /// don't use 'OncePerASTToken'.
+  static const char OncePerASTToken = 0;
+  getASTManager()->processASTAsync(Invok, std::move(Collector), &OncePerASTToken);
+}

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

@@ -152,6 +152,7 @@ bool TestOptions::parseArgs(llvm::ArrayRef<const char *> Args) {
         .Case("markup-xml", SourceKitRequest::MarkupToXML)
         .Case("stats", SourceKitRequest::Statistics)
         .Case("track-compiles", SourceKitRequest::EnableCompileNotifications)
+        .Case("collect-type", SourceKitRequest::CollectExpresstionType)
         .Default(SourceKitRequest::None);
 
       if (Request == SourceKitRequest::None) {
@@ -163,7 +164,7 @@ bool TestOptions::parseArgs(llvm::ArrayRef<const char *> Args) {
                "doc-info/sema/interface-gen/interface-gen-openfind-usr/find-interface/"
                "open/close/edit/print-annotations/print-diags/extract-comment/module-groups/"
                "range/syntactic-rename/find-rename-ranges/translate/markup-xml/stats/"
-               "track-compiles\n";
+               "track-compiles/collect-type\n";
         return true;
       }
       break;

tools/SourceKit/tools/sourcekitd-test/TestOptions.h

@@ -62,6 +62,7 @@ enum class SourceKitRequest {
   Statistics,
   SyntaxTree,
   EnableCompileNotifications,
+  CollectExpresstionType,
 #define SEMANTIC_REFACTORING(KIND, NAME, ID) KIND,
 #include "swift/IDE/RefactoringKinds.def"
 };

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

@@ -68,6 +68,7 @@ static void printCursorInfo(sourcekitd_variant_t Info, StringRef Filename,
 static void printNameTranslationInfo(sourcekitd_variant_t Info, llvm::raw_ostream &OS);
 static void printRangeInfo(sourcekitd_variant_t Info, StringRef Filename,
                            llvm::raw_ostream &OS);
+static void printExpressionType(sourcekitd_variant_t Info, llvm::raw_ostream &OS);
 static void printDocInfo(sourcekitd_variant_t Info, StringRef Filename);
 static void printInterfaceGen(sourcekitd_variant_t Info, bool CheckASCII);
 static void printSemanticInfo();
@@ -614,6 +615,11 @@ static int handleTestInvocation(TestOptions Opts, TestOptions &InitOpts) {
     break;
   }
 
+  case SourceKitRequest::CollectExpresstionType: {
+    sourcekitd_request_dictionary_set_uid(Req, KeyRequest, RequestCollectExpressionType);
+    break;
+  }
+
 #define SEMANTIC_REFACTORING(KIND, NAME, ID) case SourceKitRequest::KIND:                 \
     {                                                                                     \
       sourcekitd_request_dictionary_set_uid(Req, KeyRequest, RequestSemanticRefactoring); \
@@ -1057,6 +1063,10 @@ static bool handleResponse(sourcekitd_response_t Resp, const TestOptions &Opts,
       printRangeInfo(Info, SourceFile, llvm::outs());
       break;
 
+    case SourceKitRequest::CollectExpresstionType:
+      printExpressionType(Info, llvm::outs());
+      break;
+
     case SourceKitRequest::DocInfo:
       printDocInfo(Info, SourceFile);
       break;
@@ -1525,6 +1535,25 @@ static void printRangeInfo(sourcekitd_variant_t Info, StringRef FilenameIn,
   if (Typename)
     OS << "<type>" <<Typename << "</type>\n";
 }
+
+static void printExpressionType(sourcekitd_variant_t Info, llvm::raw_ostream &OS) {
+  auto *TypeBuffer = sourcekitd_variant_dictionary_get_string(Info, KeyTypeBuffer);
+  sourcekitd_variant_t ExprList = sourcekitd_variant_dictionary_get_value(Info,
+    KeyExpressionTypeList);
+  unsigned Count = sourcekitd_variant_array_get_count(ExprList);
+  for (unsigned i = 0; i != Count; ++i) {
+    sourcekitd_variant_t Item = sourcekitd_variant_array_get_value(ExprList, i);
+    unsigned Offset = sourcekitd_variant_dictionary_get_int64(Item, KeyExpressionOffset);
+    unsigned Length = sourcekitd_variant_dictionary_get_int64(Item, KeyExpressionLength);
+    StringRef PrintedType(TypeBuffer + sourcekitd_variant_dictionary_get_int64(Item,
+      KeyTypeOffset), sourcekitd_variant_dictionary_get_int64(Item, KeyTypeLength));
+    OS << "(" << Offset << ", " << Offset + Length << "): " << PrintedType << "\n";
+  }
+  if (!Count) {
+    OS << "cannot find expression types in the file\n";
+  }
+}
+
 static void printFoundInterface(sourcekitd_variant_t Info,
                                 llvm::raw_ostream &OS) {
   const char *Name = sourcekitd_variant_dictionary_get_string(Info,

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

@@ -141,6 +141,8 @@ static sourcekitd_response_t reportDocInfo(llvm::MemoryBuffer *InputBuf,
 
 static void reportCursorInfo(const CursorInfoData &Info, ResponseReceiver Rec);
 
+static void reportExpressionTypeInfo(const ExpressionTypesInFile &Info, ResponseReceiver Rec);
+
 static void reportRangeInfo(const RangeInfo &Info, ResponseReceiver Rec);
 
 static void reportNameInfo(const NameTranslatingInfo &Info, ResponseReceiver Rec);
@@ -1000,6 +1002,14 @@ handleSemanticRequest(RequestDict Req,
     return Rec(createErrorRequestInvalid("'key.line' or 'key.column' are required"));
   }
 
+  if (ReqUID == RequestCollectExpressionType) {
+    LangSupport &Lang = getGlobalContext().getSwiftLangSupport();
+    return Lang.collectExpressionTypes(*SourceFile, Args,
+      [Rec](const ExpressionTypesInFile &Info) {
+        reportExpressionTypeInfo(Info, Rec);
+      });
+  }
+
   if (ReqUID == RequestFindLocalRenameRanges) {
     int64_t Line = 0, Column = 0, Length = 0;
     if (Req.getInt64(KeyLine, Line, /*isOptional=*/false))
@@ -1723,6 +1733,25 @@ static void reportNameInfo(const NameTranslatingInfo &Info, ResponseReceiver Rec
   Rec(RespBuilder.createResponse());
 }
 
+//===----------------------------------------------------------------------===//
+// ReportExpressionTypeInfo
+//===----------------------------------------------------------------------===//
+static void reportExpressionTypeInfo(const ExpressionTypesInFile &Info,
+                                     ResponseReceiver Rec) {
+  ResponseBuilder Builder;
+  auto Dict = Builder.getDictionary();
+  Dict.set(KeyTypeBuffer, Info.TypeBuffer);
+  ResponseBuilder::Array Arr = Dict.setArray(KeyExpressionTypeList);
+  for (auto Result: Info.Results) {
+    auto Elem = Arr.appendDictionary();
+    Elem.set(KeyExpressionOffset, Result.ExprOffset);
+    Elem.set(KeyExpressionLength, Result.ExprLength);
+    Elem.set(KeyTypeOffset, Result.TypeOffset);
+    Elem.set(KeyTypeLength, Result.TypeLength);
+  }
+  Rec(Builder.createResponse());
+}
+
 //===----------------------------------------------------------------------===//
 // FindRelatedIdents
 //===----------------------------------------------------------------------===//

utils/gyb_sourcekit_support/UIDs.py

@@ -163,6 +163,12 @@ def __init__(self, internal_name, external_name):
     KEY('ImplicitMembers', 'key.implicitmembers'),
     KEY('ExpectedTypes', 'key.expectedtypes'),
     KEY('Members', 'key.members'),
+    KEY('TypeBuffer', 'key.printedtypebuffer'),
+    KEY('ExpressionTypeList', 'key.expression_type_list'),
+    KEY('ExpressionOffset', 'key.expression_offset'),
+    KEY('ExpressionLength', 'key.expression_length'),
+    KEY('TypeOffset', 'key.type_offset'),
+    KEY('TypeLength', 'key.type_length'),
 ]
 
 
@@ -219,6 +225,7 @@ def __init__(self, internal_name, external_name):
     REQUEST('EnableCompileNotifications',
             'source.request.enable-compile-notifications'),
     REQUEST('TestNotification', 'source.request.test_notification'),
+    REQUEST('CollectExpressionType', 'source.request.expression.type'),
 ]