Port incremental parse ability to sourcekit-lsp

This feature will be used when we call `changeDocument` in SwiftLanguageServer

Author: StevenWong12

Sources/SourceKitLSP/Swift/SwiftLanguageServer.swift

@@ -119,6 +119,11 @@ public final class SwiftLanguageServer: ToolchainLanguageServer {
 
   var commandsByFile: [DocumentURI: SwiftCompileCommand] = [:]
 
+  var documentLookaheadRanges: [DocumentURI: LookaheadRanges] = [:]
+  
+  /// *For Testing*
+  public var documentReusedNodeCollector: [DocumentURI: IncrementalParseReusedNodeCollector] = [:]
+
   var keys: sourcekitd_keys { return sourcekitd.keys }
   var requests: sourcekitd_requests { return sourcekitd.requests }
   var values: sourcekitd_values { return sourcekitd.values }
@@ -197,14 +202,36 @@ public final class SwiftLanguageServer: ToolchainLanguageServer {
   }
 
   /// Returns the updated lexical tokens for the given `snapshot`.
+  ///
+  /// - Parameters:
+  ///   - edits: If we are in the context of editing the contents of a file, i.e. calling ``SwiftLanguageServer/changeDocument(_:)``, we should pass `edits` to enable incremental parse. Otherwise, `edits` should be `nil`.
   private func updateSyntaxTree(
-    for snapshot: DocumentSnapshot
+    for snapshot: DocumentSnapshot,
+    with edits: ConcurrentEdits? = nil
   ) -> DocumentTokens {
     logExecutionTime(level: .debug) {
       var docTokens = snapshot.tokens
+      let documentURI = snapshot.document.uri
+      
+      var reusedNodeCollector = documentReusedNodeCollector[documentURI]
+      if reusedNodeCollector == nil {
+        documentReusedNodeCollector[documentURI] = IncrementalParseReusedNodeCollector()
+        reusedNodeCollector = documentReusedNodeCollector[documentURI]
+      }
+      
+      var parseTransition: IncrementalParseTransition? = nil
+      if let previousTree = snapshot.tokens.syntaxTree,
+         let lookaheadRanges = documentLookaheadRanges[documentURI],
+         let edits {
+        parseTransition = IncrementalParseTransition(previousTree: previousTree, edits: edits, lookaheadRanges: lookaheadRanges, reusedNodeDelegate: reusedNodeCollector)
+      }
+      let (tree, nextLookaheadRanges) = Parser.parseIncrementally(
+        source: snapshot.text, parseTransition: parseTransition)
 
-      docTokens.syntaxTree = Parser.parse(source: snapshot.text)
+      docTokens.syntaxTree = tree
 
+      documentLookaheadRanges[documentURI] = nextLookaheadRanges
+      
       return docTokens
     }
   }
@@ -520,34 +547,45 @@ extension SwiftLanguageServer {
       // Clear settings that should not be cached for closed documents.
       self.commandsByFile[uri] = nil
       self.currentDiagnostics[uri] = nil
+      self.documentLookaheadRanges[uri] = nil
+      self.documentReusedNodeCollector[uri] = nil
 
       _ = try? self.sourcekitd.sendSync(req)
     }
   }
 
   public func changeDocument(_ note: DidChangeTextDocumentNotification) {
     let keys = self.keys
+    var edits: [IncrementalEdit] = []
 
     self.queue.async {
       var lastResponse: SKDResponseDictionary? = nil
 
-      let snapshot = self.documentManager.edit(note) { (before: DocumentSnapshot, edit: TextDocumentContentChangeEvent) in
+      let snapshot = self.documentManager.edit(note) {
+        (before: DocumentSnapshot, edit: TextDocumentContentChangeEvent) in
         let req = SKDRequestDictionary(sourcekitd: self.sourcekitd)
         req[keys.request] = self.requests.editor_replacetext
         req[keys.name] = note.textDocument.uri.pseudoPath
 
         if let range = edit.range {
-          guard let offset = before.utf8Offset(of: range.lowerBound), let end = before.utf8Offset(of: range.upperBound) else {
+          guard let offset = before.utf8Offset(of: range.lowerBound),
+            let end = before.utf8Offset(of: range.upperBound)
+          else {
             fatalError("invalid edit \(range)")
           }
 
+          let length = end - offset
           req[keys.offset] = offset
-          req[keys.length] = end - offset
+          req[keys.length] = length
 
+          edits.append(IncrementalEdit(offset: offset, length: length, replacementLength: edit.text.utf8.count))
         } else {
           // Full text
+          let length = before.text.utf8.count
           req[keys.offset] = 0
-          req[keys.length] = before.text.utf8.count
+          req[keys.length] = length
+
+          edits.append(IncrementalEdit(offset: 0, length: length, replacementLength: edit.text.utf8.count))
         }
 
         req[keys.sourcetext] = edit.text
@@ -556,7 +594,7 @@ extension SwiftLanguageServer {
         self.adjustDiagnosticRanges(of: note.textDocument.uri, for: edit)
       } updateDocumentTokens: { (after: DocumentSnapshot) in
         if lastResponse != nil {
-          return self.updateSyntaxTree(for: after)
+          return self.updateSyntaxTree(for: after, with: ConcurrentEdits(fromSequential: edits))
         } else {
           return DocumentTokens()
         }

Tests/SourceKitLSPTests/LocalSwiftTests.swift

@@ -1476,4 +1476,49 @@ final class LocalSwiftTests: XCTestCase {
     data = EditorPlaceholder(text)
     XCTAssertNil(data)
   }
+  
+  func testIncrementalParse() throws {
+    let url = URL(fileURLWithPath: "/\(UUID())/a.swift")
+    let uri = DocumentURI(url)
+
+    sk.allowUnexpectedNotification = false
+    
+    sk.send(DidOpenTextDocumentNotification(textDocument: TextDocumentItem(
+      uri: uri,
+      language: .swift,
+      version: 0,
+      text: """
+      func foo() {
+      }
+      class bar {
+      }
+      """
+    )))
+    
+    let didChangeTextDocumentExpectation = self.expectation(description: "didChangeTextDocument")
+    sk.sendNoteSync(DidChangeTextDocumentNotification(textDocument: .init(uri, version: 1), contentChanges: [
+      .init(range: Range(Position(line: 2, utf16index: 7)), text: "a"),
+    ]), { (note: LanguageServerProtocol.Notification<PublishDiagnosticsNotification>) -> Void in
+      log("Received diagnostics for text edit - syntactic")
+      didChangeTextDocumentExpectation.fulfill()
+    }, { (note: LanguageServerProtocol.Notification<PublishDiagnosticsNotification>) -> Void in
+      log("Received diagnostics for text edit - semantic")
+    })
+    
+    self.wait(for: [didChangeTextDocumentExpectation], timeout: defaultTimeout)
+    
+    let swiftLanguageServer = connection.server!._languageService(for: uri, .swift, in: connection.server!.workspaceForDocumentOnQueue(uri: uri)!) as! SwiftLanguageServer
+
+    let nodes = try XCTUnwrap(swiftLanguageServer.documentReusedNodeCollector[uri]?.nodes)
+    XCTAssertEqual(nodes.count, 1)
+    
+    let firstNode = try XCTUnwrap(nodes.first)
+    XCTAssertEqual(firstNode.description,
+      """
+      func foo() {
+      }
+      """
+    )
+    XCTAssertEqual(firstNode.kind, .codeBlockItem)
+  }
 }