Avoid crashing on invalid range in editDocument

Rename `LineTable.replace(utf8Offset:length:with)` to `tryReplace`
and bail if the provided range is out of bounds of the buffer. This
ensures we match the behavior of SourceKit when handling an
`editor.replacetext` request.

rdar://90385969

Author: hamishknight

Sources/SKTestSupport/Utils.swift

@@ -13,6 +13,7 @@
 package import Foundation
 package import LanguageServerProtocol
 import SwiftExtensions
+import XCTest
 
 import struct TSCBasic.AbsolutePath
 
@@ -124,3 +125,15 @@ var globalModuleCache: URL? {
       .appendingPathComponent("shared-module-cache")
   }
 }
+
+/// Similar to XCTAssertThrowsError, but accepts an async function.
+package func assertAsyncThrowsError<T>(
+  _ fn: @autoclosure () async throws -> T,
+  file: StaticString = #filePath,
+  line: UInt = #line
+) async {
+  do {
+    _ = try await fn()
+    XCTFail("Expected thrown error", file: file, line: line)
+  } catch {}
+}

Sources/SKUtilities/LineTable.swift

@@ -116,20 +116,26 @@ extension LineTable {
   }
 
   /// Replace the line table's `content` in the given range and update the line data.
+  /// If the given range is out-of-bounds, logs an error and returns without mutating.
   ///
-  /// - parameter fromLine: Starting line number (zero-based).
-  /// - parameter fromOff: Starting UTF-8 column offset (zero-based).
-  /// - parameter toLine: Ending line number (zero-based).
-  /// - parameter toOff: Ending UTF-8 column offset (zero-based).
+  /// - parameter utf8Offset: Starting UTF-8 offset (zero-based).
+  /// - parameter length: UTF-8 length.
   /// - parameter replacement: The new text for the given range.
   @inlinable
-  mutating package func replace(
+  mutating package func tryReplace(
     utf8Offset fromOff: Int,
     length: Int,
     with replacement: String
   ) {
-    let start = content.utf8.index(content.startIndex, offsetBy: fromOff)
-    let end = content.utf8.index(content.startIndex, offsetBy: fromOff + length)
+    let utf8 = self.content.utf8
+    guard
+      fromOff >= 0, length >= 0,
+      let start = utf8.index(utf8.startIndex, offsetBy: fromOff, limitedBy: utf8.endIndex),
+      let end = utf8.index(start, offsetBy: length, limitedBy: utf8.endIndex)
+    else {
+      logger.error("Range end is out of bounds \(fromOff) + \(length) > \(utf8.count)")
+      return
+    }
 
     var newText = self.content
     newText.replaceSubrange(start..<end, with: replacement)

Sources/SwiftSourceKitPlugin/CodeCompletion/Connection.swift

@@ -140,7 +140,9 @@ final class Connection {
       return
     }
 
-    document.lineTable.replace(utf8Offset: offset, length: length, with: newText)
+    // Try replace the range, ignoring an invalid input. This matches SourceKit's
+    // behavior.
+    document.lineTable.tryReplace(utf8Offset: offset, length: length, with: newText)
 
     sourcekitd.ideApi.set_file_contents(impl, path, document.lineTable.content)
   }

Tests/SwiftSourceKitPluginTests/SwiftSourceKitPluginTests.swift

@@ -401,6 +401,80 @@ final class SwiftSourceKitPluginTests: XCTestCase {
     XCTAssertEqual(result3.items.count, 1)
   }
 
+  func testEditBounds() async throws {
+    try await SkipUnless.sourcekitdSupportsPlugin()
+    let sourcekitd = try await getSourceKitD()
+    let path = scratchFilePath()
+    _ = try await sourcekitd.openDocument(
+      path,
+      contents: "",
+      compilerArguments: [path]
+    )
+
+    let typeWithMethod = """
+      struct S {
+        static func foo() -> Int {}
+      }
+
+      """
+    var fullText = typeWithMethod
+
+    try await sourcekitd.editDocument(path, fromOffset: 0, length: 0, newContents: typeWithMethod)
+
+    let completion = """
+      S.
+      """
+    fullText += completion
+
+    try await sourcekitd.editDocument(path, fromOffset: typeWithMethod.utf8.count, length: 0, newContents: completion)
+
+    func testCompletion(file: StaticString = #filePath, line: UInt = #line) async throws {
+      let result = try await sourcekitd.completeOpen(
+        path: path,
+        position: Position(line: 3, utf16index: 2),
+        filter: "foo",
+        flags: []
+      )
+      XCTAssertGreaterThan(result.unfilteredResultCount, 1, file: file, line: line)
+      XCTAssertEqual(result.items.count, 1, file: file, line: line)
+    }
+    try await testCompletion()
+
+    // Bogus edits are ignored (negative offsets crash SourceKit itself so we don't test them here).
+    await assertAsyncThrowsError(
+      try await sourcekitd.editDocument(path, fromOffset: 0, length: 99999, newContents: "")
+    )
+    await assertAsyncThrowsError(
+      try await sourcekitd.editDocument(path, fromOffset: 99999, length: 1, newContents: "")
+    )
+    await assertAsyncThrowsError(
+      try await sourcekitd.editDocument(path, fromOffset: 99999, length: 0, newContents: "unrelated")
+    )
+    // SourceKit doesn't throw an error for a no-op edit.
+    try await sourcekitd.editDocument(path, fromOffset: 99999, length: 0, newContents: "")
+
+    try await sourcekitd.editDocument(path, fromOffset: 0, length: 0, newContents: "")
+    try await sourcekitd.editDocument(path, fromOffset: fullText.utf8.count, length: 0, newContents: "")
+
+    try await testCompletion()
+
+    let badCompletion = """
+      X.
+      """
+    fullText = fullText.dropLast(2) + badCompletion
+
+    try await sourcekitd.editDocument(path, fromOffset: fullText.utf8.count - 2, length: 2, newContents: badCompletion)
+
+    let result = try await sourcekitd.completeOpen(
+      path: path,
+      position: Position(line: 3, utf16index: 2),
+      filter: "foo",
+      flags: []
+    )
+    XCTAssertEqual(result.unfilteredResultCount, 0)
+    XCTAssertEqual(result.items.count, 0)
+  }
+
   func testDocumentation() async throws {
     try await SkipUnless.sourcekitdSupportsPlugin()
     let sourcekitd = try await getSourceKitD()