Refactors willEditDocument in favor of SourceEdits

Removes willEditDocument completion block and moves that code into the DocumentManager using SourceEdit

Author: mustiikhalil

Sources/SourceKitLSP/DocumentManager.swift

@@ -14,6 +14,7 @@ import Dispatch
 import LSPLogging
 import LanguageServerProtocol
 import SKSupport
+import SwiftSyntax
 
 /// An immutable snapshot of a document at a given time.
 ///
@@ -140,26 +141,27 @@ public final class DocumentManager {
   ///   - newVersion: The new version of the document. Must be greater than the
   ///     latest version of the document.
   ///   - edits: The edits to apply to the document
-  ///   - willEditDocument: Optional closure to call before each edit. Will be
-  ///     called multiple times if there are multiple edits.
+  ///   - requiresUpcomingEdits: Flag to return upcoming editst. Will be
+  ///     calling ``willEditDocument`` for each edit made
   /// - Returns: The snapshot of the document before the edit and the snapshot
-  ///   of the document after the edit.
+  ///   of the document after the edit, including all edits that were made.
   @discardableResult
   public func edit(
     _ uri: DocumentURI,
     newVersion: Int,
     edits: [TextDocumentContentChangeEvent],
-    willEditDocument: ((_ before: LineTable, TextDocumentContentChangeEvent) -> Void)? = nil
-  ) throws -> (preEditSnapshot: DocumentSnapshot, postEditSnapshot: DocumentSnapshot) {
+    requiresUpcomingEdits: Bool = false
+  ) throws -> (preEditSnapshot: DocumentSnapshot, postEditSnapshot: DocumentSnapshot, edits: [SourceEdit]) {
     return try queue.sync {
       guard let document = documents[uri] else {
         throw Error.missingDocument(uri)
       }
       let preEditSnapshot = document.latestSnapshot
 
+      var upcomingEdits: [SourceEdit] = []
       for edit in edits {
-        if let willEditDocument {
-          willEditDocument(document.latestLineTable, edit)
+        if requiresUpcomingEdits {
+          upcomingEdits.append(willEditDocument(document.latestLineTable, edit: edit))
         }
 
         if let range = edit.range {
@@ -180,7 +182,7 @@ public final class DocumentManager {
         logger.error("Document version did not increase on edit from \(document.latestVersion) to \(newVersion)")
       }
       document.latestVersion = newVersion
-      return (preEditSnapshot, document.latestSnapshot)
+      return (preEditSnapshot, document.latestSnapshot, upcomingEdits)
     }
   }
 
@@ -192,6 +194,35 @@ public final class DocumentManager {
       return document.latestSnapshot
     }
   }
+
+  private func willEditDocument(_ before: LineTable, edit: TextDocumentContentChangeEvent) -> SourceEdit {
+    guard let range = edit.range else {
+      return SourceEdit(
+        range: Range(
+          uncheckedBounds: (
+            AbsolutePosition(utf8Offset: 0), 
+            upper: AbsolutePosition(utf8Offset: before.content.utf8.count)
+          )
+        ),
+        replacement: edit.text
+      )
+    }
+    guard let offset = before.utf8OffsetOf(line: range.lowerBound.line, utf16Column: range.lowerBound.utf16index),
+            let end = before.utf8OffsetOf(line: range.upperBound.line, utf16Column: range.upperBound.utf16index)
+    else {
+      fatalError("invalid edit \(range)")
+    }
+    return SourceEdit(
+      range: Range(
+        uncheckedBounds: (
+          AbsolutePosition(utf8Offset: offset),
+          upper: AbsolutePosition(utf8Offset: end)
+        )
+      ),
+      replacement: edit.text
+    )
+  }
+
 }
 
 extension DocumentManager {
@@ -214,19 +245,19 @@ extension DocumentManager {
     }
   }
 
-  /// Convenience wrapper for `edit(_:newVersion:edits:willEditDocument:updateDocumentTokens:)`
+  /// Convenience wrapper for `edit(_:newVersion:edits:requiresUpcomingEdits:updateDocumentTokens:)`
   /// that logs on failure.
   @discardableResult
   func edit(
     _ note: DidChangeTextDocumentNotification,
-    willEditDocument: ((_ before: LineTable, TextDocumentContentChangeEvent) -> Void)? = nil
-  ) -> (preEditSnapshot: DocumentSnapshot, postEditSnapshot: DocumentSnapshot)? {
+    requiresUpcomingEdits: Bool = false
+  ) -> (preEditSnapshot: DocumentSnapshot, postEditSnapshot: DocumentSnapshot, edits: [SourceEdit])? {
     return orLog("failed to edit document", level: .error) {
       return try edit(
         note.textDocument.uri,
         newVersion: note.textDocument.version,
         edits: note.contentChanges,
-        willEditDocument: willEditDocument
+        requiresUpcomingEdits: requiresUpcomingEdits
       )
     }
   }

Sources/SourceKitLSP/Swift/SwiftLanguageServer.swift

@@ -497,33 +497,12 @@ extension SwiftLanguageServer {
       let replacement: String
     }
 
-    var edits: [Edit] = []
+    let editResult = self.documentManager.edit(
+      note,
+      requiresUpcomingEdits: true
+    )
 
-    let editResult = self.documentManager.edit(note) {
-      (before: LineTable, edit: TextDocumentContentChangeEvent) in
-      if let range = edit.range {
-        guard let offset = before.utf8OffsetOf(line: range.lowerBound.line, utf16Column: range.lowerBound.utf16index),
-          let end = before.utf8OffsetOf(line: range.upperBound.line, utf16Column: range.upperBound.utf16index)
-        else {
-          fatalError("invalid edit \(range)")
-        }
-        edits.append(
-          Edit(
-            offset: offset,
-            length: end - offset,
-            replacement: edit.text
-          )
-        )
-      } else {
-        edits.append(
-          Edit(
-            offset: 0,
-            length: before.content.utf8.count,
-            replacement: edit.text
-          )
-        )
-      }
-    }
+    let edits = editResult?.edits ?? []
     for edit in edits {
       let req = sourcekitd.dictionary([
         keys.request: self.requests.editorReplaceText,
@@ -532,8 +511,8 @@ extension SwiftLanguageServer {
         keys.enableStructure: 0,
         keys.enableDiagnostics: 0,
         keys.syntacticOnly: 1,
-        keys.offset: edit.offset,
-        keys.length: edit.length,
+        keys.offset: edit.range.lowerBound.utf8Offset,
+        keys.length: edit.positiveLength,
         keys.sourceText: edit.replacement,
       ])
       do {
@@ -543,12 +522,16 @@ extension SwiftLanguageServer {
       }
     }
 
-    guard let (preEditSnapshot, postEditSnapshot) = editResult else {
+    guard let (preEditSnapshot, postEditSnapshot, _) = editResult else {
       return
     }
     let concurrentEdits = ConcurrentEdits(
       fromSequential: edits.map {
-        IncrementalEdit(offset: $0.offset, length: $0.length, replacementLength: $0.replacement.utf8.count)
+        IncrementalEdit(
+          offset: $0.range.lowerBound.utf8Offset,
+          length: $0.positiveLength,
+          replacementLength: $0.replacement.utf8.count
+        )
       }
     )
     await syntaxTreeManager.registerEdit(
@@ -1129,3 +1112,10 @@ extension sourcekitd_api_uid_t {
     }
   }
 }
+
+extension SourceEdit {
+  var positiveLength: Int {
+    // Required since SourceEdit returns a negative length
+    -length.utf8Length
+  }
+}