Implement server-side support for inlay hints using CollectVariableType

Editors/vscode/src/inlayHints.ts

@@ -53,9 +53,9 @@ const hintStyle: InlayHintStyle = {
 
     makeDecoration: (hint, converter) => ({
         range: converter.asRange({
-            start: hint.position,
-            end: { ...hint.position, character: hint.position.character + 1 } }
-        ),
+            start: { ...hint.position, character: hint.position.character - 1 },
+            end: hint.position
+        }),
         renderOptions: {
             after: {
                 // U+200C is a zero-width non-joiner to prevent the editor from

Sources/SourceKitD/SKDResponseDictionary.swift

@@ -32,6 +32,9 @@ public final class SKDResponseDictionary {
   public subscript(key: sourcekitd_uid_t?) -> Int? {
     return Int(sourcekitd.api.variant_dictionary_get_int64(dict, key))
   }
+  public subscript(key: sourcekitd_uid_t?) -> Bool? {
+    return sourcekitd.api.variant_dictionary_get_bool(dict, key)
+  }
   public subscript(key: sourcekitd_uid_t?) -> sourcekitd_uid_t? {
     return sourcekitd.api.variant_dictionary_get_uid(dict, key)
   }

Sources/SourceKitD/sourcekitd_uids.swift

@@ -64,6 +64,11 @@ public struct sourcekitd_keys {
   public let text: sourcekitd_uid_t
   public let typename: sourcekitd_uid_t
   public let usr: sourcekitd_uid_t
+  public let variable_offset: sourcekitd_uid_t
+  public let variable_length: sourcekitd_uid_t
+  public let variable_type: sourcekitd_uid_t
+  public let variable_type_explicit: sourcekitd_uid_t
+  public let variable_type_list: sourcekitd_uid_t
 
   // Code Completion options.
   public let codecomplete_options: sourcekitd_uid_t
@@ -129,6 +134,11 @@ public struct sourcekitd_keys {
     text = api.uid_get_from_cstr("key.text")!
     typename = api.uid_get_from_cstr("key.typename")!
     usr = api.uid_get_from_cstr("key.usr")!
+    variable_offset = api.uid_get_from_cstr("key.variable_offset")!
+    variable_length = api.uid_get_from_cstr("key.variable_length")!
+    variable_type = api.uid_get_from_cstr("key.variable_type")!
+    variable_type_explicit = api.uid_get_from_cstr("key.variable_type_explicit")!
+    variable_type_list = api.uid_get_from_cstr("key.variable_type_list")!
 
     // Code Completion options
     codecomplete_options = api.uid_get_from_cstr("key.codecomplete.options")!
@@ -155,6 +165,7 @@ public struct sourcekitd_requests {
   public let codecomplete_close: sourcekitd_uid_t
   public let cursorinfo: sourcekitd_uid_t
   public let expression_type: sourcekitd_uid_t
+  public let variable_type: sourcekitd_uid_t
   public let relatedidents: sourcekitd_uid_t
   public let semantic_refactoring: sourcekitd_uid_t
 
@@ -169,6 +180,7 @@ public struct sourcekitd_requests {
     codecomplete_close = api.uid_get_from_cstr("source.request.codecomplete.close")!
     cursorinfo = api.uid_get_from_cstr("source.request.cursorinfo")!
     expression_type = api.uid_get_from_cstr("source.request.expression.type")!
+    variable_type = api.uid_get_from_cstr("source.request.variable.type")!
     relatedidents = api.uid_get_from_cstr("source.request.relatedidents")!
     semantic_refactoring = api.uid_get_from_cstr("source.request.semantic.refactoring")!
   }

Sources/SourceKitLSP/CMakeLists.txt

@@ -30,6 +30,7 @@ target_sources(SourceKitLSP PRIVATE
   Swift/SourceKitD+ResponseError.swift
   Swift/SwiftCommand.swift
   Swift/SwiftLanguageServer.swift)
+  Swift/VariableTypeInfo.swift
 set_target_properties(SourceKitLSP PROPERTIES
   INTERFACE_INCLUDE_DIRECTORIES ${CMAKE_Swift_MODULE_DIRECTORY})
 # TODO(compnerd) reduce the exposure here, why is everything PUBLIC-ly linked?

Sources/SourceKitLSP/Swift/SwiftLanguageServer.swift

@@ -1055,77 +1055,29 @@ extension SwiftLanguageServer {
   }
 
   public func inlayHints(_ req: Request<InlayHintsRequest>) {
-    // TODO: Introduce a new SourceKit request for inlay hints
-    // instead of computing them from document symbols here.
-
     guard req.params.only?.contains(.type) ?? true else {
       req.reply([])
       return
     }
 
     let uri = req.params.textDocument.uri
-    documentSymbols(uri) { symbolsResult in
+    variableTypeInfos(uri, req.params.range) { infosResult in
       do {
-        /// Filters all the document symbols for which inlay type hints
-        /// should be displayed, i.e. variable bindings, fields and properties.
-        func bindings(_ symbols: [DocumentSymbol]) -> [DocumentSymbol] {
-          symbols
-            .flatMap { bindings($0.children ?? []) + ([.variable, .field, .property].contains($0.kind) ? [$0] : []) }
-        }
-
-        let symbols = try symbolsResult.get()
-        let bindingPositions = Set(bindings(symbols).map { $0.range.upperBound })
-
-        self.expressionTypeInfos(uri) { infosResult in
-          do {
-            let infos = try infosResult.get()
-
-            // The unfiltered infos may contain multiple infos for a single position
-            // as the ending position does not necessarily identify an expression uniquely.
-            // Consider the following example:
-            //
-            //    var x = "abc" + "def"
-            //
-            // Both `"abc" + "def"` and `"def"` are matching expressions. Since we are only
-            // interested in the first expression, i.e. the one that corresponds to the
-            // bound expression, we have to do some pre-processing here. Note that this
-            // mechanism currently relies on the outermost expression being reported first.
-
-            var visitedPositions: Set<Position> = []
-            var processedInfos: [ExpressionTypeInfo] = []
-
-            // TODO: Compute inlay hints only for the requested range/categories
-            // instead of filtering them afterwards.
-
-            for info in infos {
-              let pos = info.range.upperBound
-              if (req.params.range?.contains(pos) ?? true)
-                && bindingPositions.contains(pos)
-                && !visitedPositions.contains(pos) {
-                processedInfos.append(info)
-                visitedPositions.insert(pos)
-              }
-            }
-
-            let hints = processedInfos
-              .lazy
-              .map { info in
-                InlayHint(
-                  position: info.range.upperBound,
-                  category: .type,
-                  label: info.printedType
-                )
-              }
-
-            req.reply(.success(Array(hints)))
-          } catch {
-            let message = "expression types for inlay hints failed for \(uri): \(error)"
-            log(message, level: .warning)
-            req.reply(.failure(.unknown(message)))
+        let infos = try infosResult.get()
+        let hints = infos
+          .lazy
+          .filter { !$0.hasExplicitType }
+          .map { info in
+            InlayHint(
+              position: info.range.upperBound,
+              category: .type,
+              label: info.printedType
+            )
           }
-        }
+
+        req.reply(.success(Array(hints)))
       } catch {
-        let message = "document symbols for inlay hints failed for \(uri): \(error)"
+        let message = "variable types for inlay hints failed for \(uri): \(error)"
         log(message, level: .warning)
         req.reply(.failure(.unknown(message)))
       }

Sources/SourceKitLSP/Swift/VariableTypeInfo.swift

@@ -0,0 +1,125 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2014 - 2021 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+
+import Dispatch
+import LanguageServerProtocol
+import SourceKitD
+
+/// A typed variable as returned by sourcekitd's CollectVariableType.
+struct VariableTypeInfo {
+  /// Range of the variable identifier in the source file.
+  var range: Range<Position>
+  /// The printed type of the variable.
+  var printedType: String
+  /// Whether the variable has an explicit type annotation in the source file.
+  var hasExplicitType: Bool
+
+  init?(_ dict: SKDResponseDictionary, in snapshot: DocumentSnapshot) {
+    let keys = dict.sourcekitd.keys
+
+    guard let offset: Int = dict[keys.variable_offset],
+          let length: Int = dict[keys.variable_length],
+          let startIndex = snapshot.positionOf(utf8Offset: offset),
+          let endIndex = snapshot.positionOf(utf8Offset: offset + length),
+          let printedType: String = dict[keys.variable_type],
+          let hasExplicitType: Bool = dict[keys.variable_type_explicit] else {
+      return nil
+    }
+
+    self.range = startIndex..<endIndex
+    self.printedType = printedType
+    self.hasExplicitType = hasExplicitType
+  }
+}
+
+enum VariableTypeInfoError: Error, Equatable {
+  /// The given URL is not a known document.
+  case unknownDocument(DocumentURI)
+
+  /// The underlying sourcekitd request failed with the given error.
+  case responseError(ResponseError)
+}
+
+extension SwiftLanguageServer {
+  /// Must be called on self.queue.
+  private func _variableTypeInfos(
+    _ uri: DocumentURI,
+    _ range: Range<Position>? = nil,
+    _ completion: @escaping (Swift.Result<[VariableTypeInfo], VariableTypeInfoError>) -> Void
+  ) {
+    dispatchPrecondition(condition: .onQueue(queue))
+
+    guard let snapshot = documentManager.latestSnapshot(uri) else {
+      return completion(.failure(.unknownDocument(uri)))
+    }
+
+    let keys = self.keys
+
+    let skreq = SKDRequestDictionary(sourcekitd: sourcekitd)
+    skreq[keys.request] = requests.variable_type
+    skreq[keys.sourcefile] = snapshot.document.uri.pseudoPath
+
+    if let range = range,
+       let start = snapshot.utf8Offset(of: range.lowerBound),
+       let end = snapshot.utf8Offset(of: range.upperBound) {
+      skreq[keys.offset] = start
+      skreq[keys.length] = end - start
+    }
+
+    // FIXME: SourceKit should probably cache this for us
+    if let compileCommand = self.commandsByFile[uri] {
+      skreq[keys.compilerargs] = compileCommand.compilerArgs
+    }
+
+    let handle = self.sourcekitd.send(skreq, self.queue) { result in
+      guard let dict = result.success else {
+        return completion(.failure(.responseError(ResponseError(result.failure!))))
+      }
+
+      guard let skVariableTypeInfos: SKDResponseArray = dict[keys.variable_type_list] else {
+        return completion(.success([]))
+      }
+
+      var variableTypeInfos: [VariableTypeInfo] = []
+      variableTypeInfos.reserveCapacity(skVariableTypeInfos.count)
+
+      skVariableTypeInfos.forEach { (_, skVariableTypeInfo) -> Bool in
+        guard let info = VariableTypeInfo(skVariableTypeInfo, in: snapshot) else {
+          assertionFailure("VariableTypeInfo failed to deserialize")
+          return true
+        }
+        variableTypeInfos.append(info)
+        return true
+      }
+
+      completion(.success(variableTypeInfos))
+    }
+
+    // FIXME: cancellation
+    _ = handle
+  }
+
+  /// Provides typed variable declarations in a document.
+  ///
+  /// - Parameters:
+  ///   - url: Document URL in which to perform the request. Must be an open document.
+  ///   - completion: Completion block to asynchronously receive the VariableTypeInfos, or error.
+  func variableTypeInfos(
+    _ uri: DocumentURI,
+    _ range: Range<Position>? = nil,
+    _ completion: @escaping (Swift.Result<[VariableTypeInfo], VariableTypeInfoError>) -> Void
+  ) {
+    queue.async {
+      self._variableTypeInfos(uri, range, completion)
+    }
+  }
+}