Merge pull request #37867 from fwcd/sourcekit-var-types

[SourceKit] Add `CollectVariableType` request

include/swift/Basic/SourceLoc.h

@@ -123,6 +123,14 @@ class SourceRange {
   /// includes both this range and the other one.
   void widen(SourceRange Other);
 
+  /// Checks whether this range contains the given location. Note that the given
+  /// location should correspond to the start of a token, since locations inside
+  /// the last token may be considered outside the range by this function.
+  bool contains(SourceLoc Loc) const;
+
+  /// Checks whether this range overlaps with the given range.
+  bool overlaps(SourceRange Other) const;
+
   bool operator==(const SourceRange &other) const {
     return Start == other.Start && End == other.End;
   }

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 into \c VariableTypeInfos.
+  /// 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, SourceRange Range,
+                           std::vector<VariableTypeInfo> &VariableTypeInfos,
+                           llvm::raw_ostream &OS);
+
   /// FIXME: All of the below goes away once CallExpr directly stores its
   /// arguments.
 

lib/Basic/SourceLoc.cpp

@@ -270,6 +270,15 @@ void SourceRange::widen(SourceRange Other) {
     End = Other.End;
 }
 
+bool SourceRange::contains(SourceLoc Loc) const {
+  return Start.Value.getPointer() <= Loc.Value.getPointer() &&
+         Loc.Value.getPointer() <= End.Value.getPointer();
+}
+
+bool SourceRange::overlaps(SourceRange Other) const {
+  return contains(Other.Start) || Other.contains(Start);
+}
+
 void SourceLoc::printLineAndColumn(raw_ostream &OS, const SourceManager &SM,
                                    unsigned BufferID) const {
   if (isInvalid()) {

lib/IDE/IDETypeChecking.cpp

@@ -731,6 +731,117 @@ swift::collectExpressionType(SourceFile &SF,
   return Scratch;
 }
 
+/// This walker will traverse the AST and report types for every variable
+/// declaration.
+class VariableTypeCollector : public SourceEntityWalker {
+private:
+  const SourceManager &SM;
+  unsigned int BufferId;
+
+  /// The range in which variable types are to be collected.
+  SourceRange TotalRange;
+
+  /// 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. \c 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 \c OS.
+  llvm::raw_ostream &OS;
+
+  /// Map from a printed type to the offset in \c OS where the type starts.
+  llvm::StringMap<uint32_t> TypeOffsets;
+
+  /// Returns the start offset of this string in \c OS. If \c PrintedType
+  /// hasn't been printed to \c OS yet, this function will do so.
+  uint32_t getTypeOffset(StringRef PrintedType) {
+    auto It = TypeOffsets.find(PrintedType);
+    if (It == TypeOffsets.end()) {
+      TypeOffsets[PrintedType] = OS.tell();
+      OS << PrintedType << '\0';
+    }
+    return TypeOffsets[PrintedType];
+  }
+
+  /// Checks whether the given range overlaps the total range in which we
+  /// collect variable types.
+  bool overlapsTotalRange(SourceRange Range) {
+    return TotalRange.isInvalid() || Range.overlaps(TotalRange);
+  }
+
+public:
+  VariableTypeCollector(const SourceFile &SF, SourceRange Range,
+                        std::vector<VariableTypeInfo> &Results,
+                        llvm::raw_ostream &OS)
+      : SM(SF.getASTContext().SourceMgr), BufferId(*SF.getBufferID()),
+        TotalRange(Range), Results(Results), OS(OS) {}
+
+  bool walkToDeclPre(Decl *D, CharSourceRange DeclNameRange) override {
+    if (DeclNameRange.isInvalid()) {
+      return true;
+    }
+    // Skip this declaration and its subtree if outside the range
+    if (!overlapsTotalRange(D->getSourceRange())) {
+      return false;
+    }
+    if (auto VD = dyn_cast<VarDecl>(D)) {
+      unsigned VarOffset =
+          SM.getLocOffsetInBuffer(DeclNameRange.getStart(), BufferId);
+      unsigned VarLength = DeclNameRange.getByteLength();
+      // Print the type to a temporary buffer
+      SmallString<64> Buffer;
+      {
+        llvm::raw_svector_ostream OS(Buffer);
+        PrintOptions Options;
+        Options.SynthesizeSugarOnTypes = true;
+        auto Ty = VD->getType();
+        // Skip this declaration and its children if the type is an error type.
+        if (Ty->is<ErrorType>()) {
+          return false;
+        }
+        Ty->print(OS, Options);
+      }
+      // Transfer the type to `OS` if needed and get the offset of this string
+      // in `OS`.
+      auto TyOffset = getTypeOffset(Buffer.str());
+      bool HasExplicitType =
+          VD->getTypeReprOrParentPatternTypeRepr() != nullptr;
+      // Add the type information to the result list.
+      Results.emplace_back(VarOffset, VarLength, HasExplicitType, TyOffset);
+    }
+    return true;
+  }
+
+  bool walkToStmtPre(Stmt *S) override {
+    // Skip this statement and its subtree if outside the range
+    return overlapsTotalRange(S->getSourceRange());
+  }
+
+  bool walkToExprPre(Expr *E) override {
+    // Skip this expression and its subtree if outside the range
+    return overlapsTotalRange(E->getSourceRange());
+  }
+
+  bool walkToPatternPre(Pattern *P) override {
+    // Skip this pattern and its subtree if outside the range
+    return overlapsTotalRange(P->getSourceRange());
+  }
+};
+
+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, SourceRange Range,
+    std::vector<VariableTypeInfo> &VariableTypeInfos, llvm::raw_ostream &OS) {
+  VariableTypeCollector Walker(SF, Range, VariableTypeInfos, OS);
+  Walker.walk(SF);
+}
+
 ArrayRef<ValueDecl*> swift::
 canDeclProvideDefaultImplementationFor(ValueDecl* VD) {
   return evaluateOrDefault(VD->getASTContext().evaluator,

test/SourceKit/VariableType/basic.swift

@@ -0,0 +1,32 @@
+let x: Int = 3
+let y = "abc"
+
+var foo = ["abc" + "def"]
+
+struct A {
+  let x: String = ""
+  let y = ""
+}
+
+class B {
+  var x = 4.0
+  var y = [A]()
+  var z: [Int: Int] = [:]
+  var w: (Int) -> Int { { $0 * 2 } }
+}
+
+func foo() {
+  var local = 5
+}
+
+// RUN: %sourcekitd-test -req=collect-var-type %s -- %s | %FileCheck %s
+// CHECK: (1:5, 1:6): Int (explicit type: 1)
+// CHECK: (2:5, 2:6): String (explicit type: 0)
+// CHECK: (4:5, 4:8): [String] (explicit type: 0)
+// CHECK: (7:7, 7:8): String (explicit type: 1)
+// CHECK: (8:7, 8:8): String (explicit type: 0)
+// CHECK: (12:7, 12:8): Double (explicit type: 0)
+// CHECK: (13:7, 13:8): [A] (explicit type: 0)
+// CHECK: (14:7, 14:8): [Int : Int] (explicit type: 1)
+// CHECK: (15:7, 15:8): (Int) -> Int (explicit type: 1)
+// CHECK: (19:7, 19:12): Int (explicit type: 0)

test/SourceKit/VariableType/error.swift

@@ -0,0 +1,6 @@
+let x = "an error type" + 3
+let y = "not an error type"
+
+// RUN: %sourcekitd-test -req=collect-var-type %s -- %s | %FileCheck %s
+// CHECK-NOT: (1:5, 1:6)
+// CHECK: (2:5, 2:6): String (explicit type: 0)

test/SourceKit/VariableType/inout.swift

@@ -0,0 +1,10 @@
+func x(_ param: inout Int) -> Int {
+  param = 4
+}
+
+let z = { (param: inout String) in }
+
+// RUN: %sourcekitd-test -req=collect-var-type %s -- %s | %FileCheck %s
+// CHECK: (1:10, 1:15): Int (explicit type: 1)
+// CHECK: (5:5, 5:6): (inout String) -> () (explicit type: 0)
+// CHECK: (5:12, 5:17): String (explicit type: 1)

test/SourceKit/VariableType/params.swift

@@ -0,0 +1,16 @@
+func x(_ param: Int) -> Int {
+  param
+}
+
+let y: (String) -> Void = { param in }
+
+let z = { (param: String) in
+  param.count
+}
+
+// RUN: %sourcekitd-test -req=collect-var-type %s -- %s | %FileCheck %s
+// CHECK: (1:10, 1:15): Int (explicit type: 1)
+// CHECK: (5:5, 5:6): (String) -> Void (explicit type: 1)
+// CHECK: (5:29, 5:34): String (explicit type: 0)
+// CHECK: (7:5, 7:6): (String) -> Int (explicit type: 0)
+// CHECK: (7:12, 7:17): String (explicit type: 1)

test/SourceKit/VariableType/ranged.swift

@@ -0,0 +1,10 @@
+let x = 1
+let y = "abc"
+let z: String = "def"
+var w = 4
+
+// RUN: %sourcekitd-test -req=collect-var-type -pos=2:1 -end-pos=4:1 %s -- %s | %FileCheck %s
+// CHECK-NOT: (1:5, 1:6)
+// CHECK: (2:5, 2:6): String (explicit type: 0)
+// CHECK: (3:5, 3:6): String (explicit type: 1)
+// CHECK-NOT: (4:5, 4:6)

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;