Add CollectVariableType request to SourceKit

- Add VariableTypeCollector
  This new SourceEntityWalker collects types from variable declarations.
- Add SwiftLangSupport::collectVariableTypes
- Implement CollectVariableType request
- Provide information about explicit types in CollectVarType
- Fix HasExplicitType in VariableTypeArray
- Fix typo
- Implement ranged CollectVariableTypes requests
- Use offset/length params for CollectVariableType in sourcekitd-test
- Address a bunch of PR suggestions
- Remove CanonicalType from VariableTypeCollector
  This turned out not to be needed (for now).
- Improve doc comment on VariableTypeCollector::getTypeOffsets
- Remove unused CanonicalTy variable
- Remove out-of-date comment
- Run clang-format on the CollectVariableType implementation
- Fix some minor style issues
- Use emplace_back while collecting variable infos
- Pass CollectVariableType range to VariableTypeCollector
- Use capitalized variable names in VariableTypeArray
  ...as recommended by the LLVM coding standards
- Use PrintOptions for type printing in VariableTypeCollector
- Return void for collectVariableType
  This seems to be cleaner stylistically.
- Avoid visiting subranges of invalid range in VariableTypeCollector
- Use std::string for type buffer in VariableTypeCollectorASTConsumer
- Use plural for PrintedType in VariableTypeArray
- Remove unused fields in VariableTypeArrayBuilder
- Add suggested doc comments to VariableTypeArray
- Remove unused VariableTypeInfo.TypeLength
- Fix typo of ostream in VariableTypeCollectorASTConsumer
- Fix typo
- Document Offset and Length semantics in collectVariableTypes

Author: fwcd

include/swift/Sema/IDETypeChecking.h

@@ -206,6 +206,33 @@ namespace swift {
   /// the decl context.
   ProtocolDecl *resolveProtocolName(DeclContext *dc, StringRef Name);
 
+  /// Reported type of a variable declaration.
+  struct VariableTypeInfo {
+    /// The start of the variable identifier.
+    uint32_t Offset;
+
+    /// The length of the variable identifier.
+    uint32_t Length;
+
+    /// Whether the variable has an explicit type annotation.
+    bool HasExplicitType;
+
+    /// The start of the printed type in a separate string buffer.
+    uint32_t TypeOffset;
+
+    VariableTypeInfo(uint32_t Offset, uint32_t Length, bool HasExplicitType,
+                     uint32_t TypeOffset);
+  };
+
+  /// Collect type information for every variable declaration in \c SF
+  /// within the given range, defined by \c Offset and \c Length.
+  /// All types will be printed to \c OS and the type offsets of the
+  /// \c VariableTypeInfos will index into the string that backs this
+  /// stream.
+  void collectVariableType(
+      SourceFile &SF, Optional<unsigned> Offset, Optional<unsigned> Length,
+      std::vector<VariableTypeInfo> &Scratch, llvm::raw_ostream &OS);
+
   /// FIXME: All of the below goes away once CallExpr directly stores its
   /// arguments.
 

lib/IDE/IDETypeChecking.cpp

@@ -731,6 +731,95 @@ swift::collectExpressionType(SourceFile &SF,
   return Scratch;
 }
 
+/// This walker will traverse the AST and report types for every variable
+/// declaration.
+class VariableTypeCollector : public SourceEntityWalker {
+private:
+  SourceManager &SM;
+  unsigned int BufferId;
+
+  /// The start offset of the range in which variable types are to be collected.
+  Optional<unsigned> TotalOffset;
+
+  /// The length of the range in which variable types are to be collected.
+  Optional<unsigned> TotalLength;
+
+  /// The output vector for VariableTypeInfos emitted during traversal.
+  std::vector<VariableTypeInfo> &Results;
+
+  /// We print all types into a single output stream (e.g. into a string buffer)
+  /// and provide offsets into this string buffer to describe individual types,
+  /// i.e. `OS` builds a string that contains all null-terminated printed type
+  /// strings. When referring to one of these types, we can use the offsets at
+  /// which it starts in the `OS`.
+  llvm::raw_ostream &OS;
+
+  /// Map from a printed type to the offset in OS where the type starts.
+  llvm::StringMap<uint32_t> TypeOffsets;
+
+  /// Returns the start and end offset of this string in `OS`. If `PrintedType`
+  /// hasn't been printed to `OS` yet, this function will do so.
+  std::pair<uint32_t, uint32_t> getTypeOffsets(StringRef PrintedType) {
+    auto It = TypeOffsets.find(PrintedType);
+    if (It == TypeOffsets.end()) {
+      TypeOffsets[PrintedType] = OS.tell();
+      OS << PrintedType << '\0';
+    }
+    return {TypeOffsets[PrintedType], PrintedType.size()};
+  }
+
+public:
+  VariableTypeCollector(SourceFile &SF, Optional<unsigned> Offset,
+                        Optional<unsigned> Length,
+                        std::vector<VariableTypeInfo> &Results,
+                        llvm::raw_ostream &OS)
+      : SM(SF.getASTContext().SourceMgr), BufferId(*SF.getBufferID()),
+        TotalOffset(Offset), TotalLength(Length), Results(Results), OS(OS) {}
+
+  bool walkToDeclPre(Decl *D, CharSourceRange Range) override {
+    if (Range.isInvalid()) {
+      return false;
+    }
+    unsigned Offset = SM.getLocOffsetInBuffer(Range.getStart(), BufferId);
+    unsigned Length = Range.getByteLength();
+    // Skip this declaration and its subtree if outside the range
+    if (TotalOffset.hasValue() && TotalLength.hasValue() &&
+        (Offset < *TotalOffset || Offset >= (*TotalOffset + *TotalLength))) {
+      return false;
+    }
+    if (auto VD = dyn_cast<VarDecl>(D)) {
+      // Print the type to a temporary buffer
+      SmallString<64> Buffer;
+      {
+        llvm::raw_svector_ostream OS(Buffer);
+        PrintOptions Options;
+        Options.SynthesizeSugarOnTypes = true;
+        VD->getType()->print(OS, Options);
+      }
+      // Transfer the type to `OS` if needed and get the offsets of this string
+      // in `OS`.
+      auto Ty = getTypeOffsets(Buffer.str());
+      bool HasExplicitType =
+          VD->getTypeReprOrParentPatternTypeRepr() != nullptr;
+      // Add the type information to the result list.
+      Results.emplace_back(Offset, Length, HasExplicitType, Ty.first);
+    }
+    return true;
+  }
+};
+
+VariableTypeInfo::VariableTypeInfo(uint32_t Offset, uint32_t Length,
+                                   bool HasExplicitType, uint32_t TypeOffset)
+    : Offset(Offset), Length(Length), HasExplicitType(HasExplicitType),
+      TypeOffset(TypeOffset) {}
+
+void swift::collectVariableType(
+    SourceFile &SF, Optional<unsigned> Offset, Optional<unsigned> Length,
+    std::vector<VariableTypeInfo> &Scratch, llvm::raw_ostream &OS) {
+  VariableTypeCollector Walker(SF, Offset, Length, Scratch, OS);
+  Walker.walk(SF);
+}
+
 ArrayRef<ValueDecl*> swift::
 canDeclProvideDefaultImplementationFor(ValueDecl* VD) {
   return evaluateOrDefault(VD->getASTContext().evaluator,

tools/SourceKit/docs/Protocol.md

@@ -785,6 +785,49 @@ expr-type-info ::=
 $ sourcekitd-test -req=collect-type /path/to/file.swift -- /path/to/file.swift
 ```
 
+## Variable Type
+
+This request collects the types of all variable declarations 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.variable.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.
+    [opt] <key.offset>:       (int64)   // Offset of the requested range. Defaults to zero.
+    [opt] <key.length>:       (int64)   // Length of the requested range. Defaults to the entire file.
+}
+```
+
+### Response
+```
+{
+    <key.variable_type_list>: (array) [var-type-info*]   // A list of variable declarations and types
+}
+```
+
+```
+var-type-info ::=
+{
+    <key.variable_offset>:       (int64)    // Offset of a variable identifier in the source file
+    <key.variable_length>:       (int64)    // Length of a variable identifier an expression in the source file
+    <key.variable_type>:         (string)   // Printed type of the variable declaration
+    <key.variable_type_explicit> (bool)     // Whether the declaration has an explicit type annotation
+}
+```
+
+### Testing
+
+```
+$ sourcekitd-test -req=collect-var-type /path/to/file.swift -- /path/to/file.swift
+```
+
 # UIDs
 
 ## Keys

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

@@ -138,6 +138,27 @@ struct ExpressionTypesInFile {
   StringRef TypeBuffer;
 };
 
+struct VariableType {
+  /// The variable identifier's offset in the file.
+  unsigned VarOffset;
+  /// The variable identifier's length.
+  unsigned VarLength;
+  /// The offset of the type's string representation inside
+  /// `VariableTypesInFile.TypeBuffer`.
+  unsigned TypeOffset;
+  /// Whether the variable declaration has an explicit type annotation.
+  bool HasExplicitType;
+};
+
+struct VariableTypesInFile {
+  /// The typed variable declarations in the file.
+  std::vector<VariableType> Results;
+  /// A String containing the printed representation of all types in
+  /// `Results`. Entries in `Results` refer to their types by using
+  /// an offset into this string.
+  StringRef TypeBuffer;
+};
+
 class CodeCompletionConsumer {
   virtual void anchor();
 
@@ -864,6 +885,15 @@ class LangSupport {
                                       std::function<void(const
                                           RequestResult<ExpressionTypesInFile> &)> Receiver) = 0;
 
+  /// Collects variable types for a range defined by `Offset` and `Length` in
+  /// the source file. If `Offset` or `Length` are empty, variable types for
+  /// the entire document are collected.
+  virtual void collectVariableTypes(
+      StringRef FileName, ArrayRef<const char *> Args,
+      Optional<unsigned> Offset, Optional<unsigned> Length,
+      std::function<void(const RequestResult<VariableTypesInFile> &)>
+          Receiver) = 0;
+
   virtual void getDocInfo(llvm::MemoryBuffer *InputBuf,
                           StringRef ModuleName,
                           ArrayRef<const char *> Args,

tools/SourceKit/lib/SwiftLang/SwiftLangSupport.h

@@ -603,6 +603,12 @@ class SwiftLangSupport : public LangSupport {
                               bool CanonicalType,
                               std::function<void(const RequestResult<ExpressionTypesInFile> &)> Receiver) override;
 
+  void collectVariableTypes(
+      StringRef FileName, ArrayRef<const char *> Args,
+      Optional<unsigned> Offset, Optional<unsigned> Length,
+      std::function<void(const RequestResult<VariableTypesInFile> &)> Receiver)
+      override;
+
   void semanticRefactoring(StringRef Filename, SemanticRefactoringInfo Info,
                            ArrayRef<const char*> Args,
                            CategorizedEditsReceiver Receiver) override;

tools/SourceKit/lib/SwiftLang/SwiftSourceDocInfo.cpp

@@ -2368,3 +2368,62 @@ void SwiftLangSupport::collectExpressionTypes(StringRef FileName,
                                    &OncePerASTToken,
                                    llvm::vfs::getRealFileSystem());
 }
+
+void SwiftLangSupport::collectVariableTypes(
+    StringRef FileName, ArrayRef<const char *> Args, Optional<unsigned> Offset,
+    Optional<unsigned> Length,
+    std::function<void(const RequestResult<VariableTypesInFile> &)> Receiver) {
+  std::string Error;
+  SwiftInvocationRef Invok = ASTMgr->getInvocation(Args, FileName, Error);
+  if (!Invok) {
+    LOG_WARN_FUNC("failed to create an ASTInvocation: " << Error);
+    Receiver(RequestResult<VariableTypesInFile>::fromError(Error));
+    return;
+  }
+  assert(Invok);
+
+  class VariableTypeCollectorASTConsumer : public SwiftASTConsumer {
+  private:
+    std::function<void(const RequestResult<VariableTypesInFile> &)> Receiver;
+    Optional<unsigned> Offset;
+    Optional<unsigned> Length;
+
+  public:
+    VariableTypeCollectorASTConsumer(
+        std::function<void(const RequestResult<VariableTypesInFile> &)>
+            Receiver,
+        Optional<unsigned> Offset, Optional<unsigned> Length)
+        : Receiver(std::move(Receiver)), Offset(Offset), Length(Length) {}
+
+    void handlePrimaryAST(ASTUnitRef AstUnit) override {
+      auto *SF = AstUnit->getCompilerInstance().getPrimarySourceFile();
+      std::vector<VariableTypeInfo> Infos;
+      std::string TypeBuffer;
+      llvm::raw_string_ostream OS(TypeBuffer);
+      VariableTypesInFile Result;
+      collectVariableType(*SF, Offset, Length, Infos, OS);
+      for (auto Info : Infos) {
+        Result.Results.push_back({Info.Offset, Info.Length, Info.TypeOffset, Info.HasExplicitType});
+      }
+      Result.TypeBuffer = OS.str();
+      Receiver(RequestResult<VariableTypesInFile>::fromResult(Result));
+    }
+
+    void cancelled() override {
+      Receiver(RequestResult<VariableTypesInFile>::cancelled());
+    }
+
+    void failed(StringRef Error) override {
+      Receiver(RequestResult<VariableTypesInFile>::fromError(Error));
+    }
+  };
+
+  auto Collector = std::make_shared<VariableTypeCollectorASTConsumer>(
+      Receiver, Offset, Length);
+  /// 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,
+                                   llvm::vfs::getRealFileSystem());
+}

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

@@ -147,6 +147,7 @@ bool TestOptions::parseArgs(llvm::ArrayRef<const char *> Args) {
         .Case("stats", SourceKitRequest::Statistics)
         .Case("track-compiles", SourceKitRequest::EnableCompileNotifications)
         .Case("collect-type", SourceKitRequest::CollectExpresstionType)
+        .Case("collect-var-type", SourceKitRequest::CollectVariableType)
         .Case("global-config", SourceKitRequest::GlobalConfiguration)
         .Case("dependency-updated", SourceKitRequest::DependencyUpdated)
 #define SEMANTIC_REFACTORING(KIND, NAME, ID) .Case("refactoring." #ID, SourceKitRequest::KIND)

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

@@ -64,6 +64,7 @@ enum class SourceKitRequest {
   SyntaxTree,
   EnableCompileNotifications,
   CollectExpresstionType,
+  CollectVariableType,
   GlobalConfiguration,
   DependencyUpdated,
 #define SEMANTIC_REFACTORING(KIND, NAME, ID) KIND,

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

@@ -67,6 +67,7 @@ static void printNameTranslationInfo(sourcekitd_variant_t Info, llvm::raw_ostrea
 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 printVariableType(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();
@@ -738,6 +739,17 @@ static int handleTestInvocation(TestOptions Opts, TestOptions &InitOpts) {
     break;
   }
 
+  case SourceKitRequest::CollectVariableType: {
+    sourcekitd_request_dictionary_set_uid(Req, KeyRequest,
+                                          RequestCollectVariableType);
+    if (Opts.Length) {
+      sourcekitd_request_dictionary_set_int64(Req, KeyOffset, ByteOffset);
+      sourcekitd_request_dictionary_set_int64(Req, KeyLength, Opts.Length);
+    }
+    addRequestOptionsDirect(Req, Opts);
+    break;
+  }
+
 #define SEMANTIC_REFACTORING(KIND, NAME, ID)                                   \
   case SourceKitRequest::KIND:                                                 \
     setRefactoringFields(Req, Opts, KindRefactoring##KIND, SourceBuf.get());   \
@@ -1254,6 +1266,10 @@ static bool handleResponse(sourcekitd_response_t Resp, const TestOptions &Opts,
       printExpressionType(Info, llvm::outs());
       break;
 
+    case SourceKitRequest::CollectVariableType:
+      printVariableType(Info, llvm::outs());
+      break;
+
     case SourceKitRequest::DocInfo:
       printDocInfo(Info, SourceFile);
       break;
@@ -1945,6 +1961,27 @@ static void printExpressionType(sourcekitd_variant_t Info, llvm::raw_ostream &OS
   OS << "</ExpressionTypes>\n";
 }
 
+static void printVariableType(sourcekitd_variant_t Info, llvm::raw_ostream &OS) {
+  auto TypeBuffer =
+      sourcekitd_variant_dictionary_get_value(Info, KeyVariableTypeList);
+  unsigned Count = sourcekitd_variant_array_get_count(TypeBuffer);
+  if (!Count) {
+    OS << "cannot find variable types in the file\n";
+    return;
+  }
+  OS << "<VariableTypes>\n";
+  for (unsigned i = 0; i != Count; ++i) {
+    sourcekitd_variant_t Item = sourcekitd_variant_array_get_value(TypeBuffer, i);
+    unsigned Offset = sourcekitd_variant_dictionary_get_int64(Item, KeyVariableOffset);
+    unsigned Length = sourcekitd_variant_dictionary_get_int64(Item, KeyVariableLength);
+    bool HasExplicitType = sourcekitd_variant_dictionary_get_bool(Item, KeyVariableTypeExplicit);
+    OS << "(" << Offset << ", " << Offset + Length << "): "
+       << sourcekitd_variant_dictionary_get_string(Item, KeyVariableType)
+       << " (explicit type: " << HasExplicitType << ")\n";
+  }
+  OS << "</VariableTypes>\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/include/sourcekitd/Internal.h

@@ -60,6 +60,7 @@ enum class CustomBufferKind {
   DocStructureElementArray,
   AttributesArray,
   ExpressionTypeArray,
+  VariableTypeArray,
   RawData
 };