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

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
}

rdar:://35199889

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/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'),
 ]