Fix offset conversions at the end of the document

src/main/java/com/sourcegraph/cody/vscode/IntelliJTextDocument.java

@@ -6,11 +6,8 @@
 import com.intellij.openapi.editor.Editor;
 import com.intellij.openapi.fileEditor.FileDocumentManager;
 import com.intellij.openapi.project.Project;
-import com.intellij.openapi.util.TextRange;
 import com.intellij.openapi.vfs.VirtualFile;
-import com.sourcegraph.cody.agent.protocol_extensions.PositionKt;
 import com.sourcegraph.cody.agent.protocol_generated.Position;
-import com.sourcegraph.cody.agent.protocol_generated.Range;
 import java.net.URI;
 import java.util.Optional;
 import org.jetbrains.annotations.NotNull;
@@ -39,28 +36,17 @@ public String fileName() {
     return file.getName();
   }
 
-  @Override
-  public int offsetAt(Position position) {
-    return PositionKt.toOffset(position, this.editor.getDocument());
-  }
-
   @Override
   public String getText() {
     return this.editor.getDocument().getText();
   }
 
-  @Override
-  public String getText(Range range) {
-    return this.editor
-        .getDocument()
-        .getText(TextRange.create(offsetAt(range.getStart()), offsetAt(range.getEnd())));
-  }
-
   @Override
   public Position positionAt(int offset) {
-    int line = this.editor.getDocument().getLineNumber(offset);
-    int lineStartOffset = offsetAt(new Position(line, 0));
-    return new Position(line, offset - lineStartOffset);
+    Document document = this.editor.getDocument();
+    int line = document.getLineNumber(offset);
+    int character = offset - document.getLineStartOffset(line);
+    return new Position(line, character);
   }
 
   @Override

src/main/java/com/sourcegraph/cody/vscode/TextDocument.java

@@ -1,7 +1,6 @@
 package com.sourcegraph.cody.vscode;
 
 import com.sourcegraph.cody.agent.protocol_generated.Position;
-import com.sourcegraph.cody.agent.protocol_generated.Range;
 import java.net.URI;
 import java.util.Optional;
 import org.jetbrains.annotations.NotNull;
@@ -12,12 +11,8 @@ public interface TextDocument {
   @NotNull
   String fileName();
 
-  int offsetAt(Position position);
-
   String getText();
 
-  String getText(Range range);
-
   Position positionAt(int offset);
 
   @NotNull

src/main/kotlin/com/sourcegraph/cody/agent/protocol_extensions/Position.kt

@@ -12,15 +12,7 @@ fun Position(line: Int, character: Int): Position {
   return Position(line.toLong(), character.toLong())
 }
 
-object PositionFactory {
-  fun fromOffset(document: Document, offset: Int): Position {
-    val line = document.getLineNumber(offset)
-    val lineStartOffset = document.getLineStartOffset(line)
-    return Position(line.toLong(), (offset - lineStartOffset).toLong())
-  }
-}
-
-fun Position.isStartOrEndOfDocumentMarker(document: Document): Boolean {
+fun Position.isOutsideOfDocument(document: Document): Boolean {
   return line < 0 || line > document.lineCount
 }
 
@@ -38,8 +30,7 @@ fun Position.toLogicalPosition(document: Document): LogicalPosition {
   return LogicalPosition(getRealLine(document), getRealColumn(document))
 }
 
-/** Return zero-based offset of this position in the document. */
-fun Position.toOffset(document: Document): Int {
+fun Position.toOffsetOrZero(document: Document): Int {
   val lineStartOffset = document.getLineStartOffset(getRealLine(document))
   return lineStartOffset + getRealColumn(document)
 }

src/main/kotlin/com/sourcegraph/cody/agent/protocol_extensions/Range.kt

@@ -1,23 +1,11 @@
 package com.sourcegraph.cody.agent.protocol_extensions
 
-import com.intellij.openapi.application.ReadAction
 import com.intellij.openapi.editor.Document
-import com.intellij.openapi.editor.RangeMarker
 import com.sourcegraph.cody.agent.protocol_generated.Range
 
 typealias RangePair = Pair<Long, Long>
 
-object RangeFactory {
-  fun fromRangeMarker(rm: RangeMarker): Range =
-      ReadAction.compute<Range, RuntimeException> {
-        Range(
-            PositionFactory.fromOffset(rm.document, rm.startOffset),
-            PositionFactory.fromOffset(rm.document, rm.endOffset))
-      }
-}
-
-public fun Range.toIntellijRange(): RangePair =
-    RangePair(this.start.line.plus(1L), this.end.line.plus(1L))
+typealias RangeOffset = Pair<Int, Int>
 
 // We need to .plus(1) since the ranges use 0-based indexing
 // but IntelliJ presents it as 1-based indexing.
@@ -26,17 +14,23 @@ public fun Range.intellijRange(): RangePair = RangePair(start.line.plus(1), end.
 // The link to Sourcegraph Search on the other hand looks like this:
 fun Range.toSearchRange(): RangePair = RangePair(start.line.plus(1), end.line)
 
-fun Range.toRangeMarker(document: Document, surviveOnExternalChange: Boolean = false): RangeMarker =
-    ReadAction.compute<RangeMarker, RuntimeException> {
-      document
-          .createRangeMarker(
-              start.toOffset(document), end.toOffset(document), surviveOnExternalChange)
-          .also {
-            it.isGreedyToLeft = true
-            it.isGreedyToRight = true
-          }
-    }
+/**
+ * Converts the range represented by this [Range] object to a pair of offsets within the given
+ * [Document].
+ *
+ * If the start or end position of the range is outside the document, the corresponding offset will
+ * be set to 0 or the document's text length, respectively.
+ *
+ * @param document The [Document] to use for converting the range to offsets.
+ * @return A [RangeOffset] pair containing the start and end offsets of the range within the
+ *   document.
+ */
+fun Range.toOffsetRange(document: Document): RangeOffset {
+  val startOffset = if (start.isOutsideOfDocument(document)) 0 else start.toOffsetOrZero(document)
+  val endOffset =
+      if (end.isOutsideOfDocument(document)) document.textLength else end.toOffsetOrZero(document)
+
+  return RangeOffset(startOffset, endOffset)
+}
 
 fun Range.length() = end.line - start.line + 1
-
-fun Range.toOffset() = 0

src/main/kotlin/com/sourcegraph/cody/edit/EditService.kt

@@ -6,7 +6,8 @@ import com.intellij.openapi.components.service
 import com.intellij.openapi.diagnostic.Logger
 import com.intellij.openapi.fileEditor.FileDocumentManager
 import com.intellij.openapi.project.Project
-import com.sourcegraph.cody.agent.protocol_extensions.toOffset
+import com.sourcegraph.cody.agent.protocol_extensions.toOffsetOrZero
+import com.sourcegraph.cody.agent.protocol_extensions.toOffsetRange
 import com.sourcegraph.cody.agent.protocol_generated.CreateFileOperation
 import com.sourcegraph.cody.agent.protocol_generated.DeleteFileOperation
 import com.sourcegraph.cody.agent.protocol_generated.DeleteTextEdit
@@ -47,17 +48,17 @@ class EditService(val project: Project) {
       edits.reversed().all { edit ->
         when (edit) {
           is ReplaceTextEdit -> {
-            document.replaceString(
-                edit.range.start.toOffset(document), edit.range.end.toOffset(document), edit.value)
+            val (startOffset, endOffset) = edit.range.toOffsetRange(document)
+            document.replaceString(startOffset, endOffset, edit.value)
             true
           }
           is DeleteTextEdit -> {
-            document.deleteString(
-                edit.range.start.toOffset(document), edit.range.end.toOffset(document))
+            val (startOffset, endOffset) = edit.range.toOffsetRange(document)
+            document.deleteString(startOffset, endOffset)
             true
           }
           is InsertTextEdit -> {
-            document.insertString(edit.position.toOffset(document), edit.value)
+            document.insertString(edit.position.toOffsetOrZero(document), edit.value)
             true
           }
         }

src/main/kotlin/com/sourcegraph/cody/edit/actions/lenses/EditShowDiffAction.kt

@@ -3,7 +3,7 @@ package com.sourcegraph.cody.edit.actions.lenses
 import com.intellij.openapi.command.WriteCommandAction
 import com.intellij.openapi.editor.EditorFactory
 import com.sourcegraph.cody.agent.CodyAgentService
-import com.sourcegraph.cody.agent.protocol_extensions.toOffset
+import com.sourcegraph.cody.agent.protocol_extensions.toOffsetRange
 import com.sourcegraph.cody.agent.protocol_generated.EditTask_GetTaskDetailsParams
 import com.sourcegraph.common.ShowDocumentDiffAction
 
@@ -16,10 +16,8 @@ class EditShowDiffAction :
           if (editTask != null) {
             val documentAfter = editor.document
             val documentBefore = EditorFactory.getInstance().createDocument(documentAfter.text)
-            documentBefore.replaceString(
-                editTask.selectionRange.start.toOffset(documentBefore),
-                editTask.selectionRange.end.toOffset(documentBefore),
-                editTask.originalText ?: "")
+            val (startOffset, endOffset) = editTask.selectionRange.toOffsetRange(documentBefore)
+            documentBefore.replaceString(startOffset, endOffset, editTask.originalText ?: "")
             ShowDocumentDiffAction(documentBefore, documentAfter).actionPerformed(event)
           }
         }

src/main/kotlin/com/sourcegraph/cody/edit/widget/LensWidgetGroup.kt

@@ -24,7 +24,7 @@ import com.intellij.ui.JBColor
 import com.intellij.util.concurrency.annotations.RequiresEdt
 import com.intellij.util.ui.UIUtil
 import com.sourcegraph.cody.agent.protocol_extensions.toLogicalPosition
-import com.sourcegraph.cody.agent.protocol_extensions.toOffset
+import com.sourcegraph.cody.agent.protocol_extensions.toOffsetRange
 import com.sourcegraph.cody.agent.protocol_generated.Range
 import java.awt.Cursor
 import java.awt.Font
@@ -114,11 +114,11 @@ class LensWidgetGroup(parentComponent: Editor) : EditorCustomElementRenderer, Di
 
   @RequiresEdt
   fun show(range: Range, shouldScrollToLens: Boolean) {
-    val offset = range.start.toOffset(editor.document)
+    val (startOffset, _) = range.toOffsetRange(editor.document)
     if (isDisposed.get()) {
       throw IllegalStateException("Request to show disposed inlay: $this")
     }
-    inlay = editor.inlayModel.addBlockElement(offset, false, true, 0, this)
+    inlay = editor.inlayModel.addBlockElement(startOffset, false, true, 0, this)
     Disposer.register(this, inlay!!)
 
     if (shouldScrollToLens) {

src/main/kotlin/com/sourcegraph/utils/CodyEditorUtil.kt

@@ -28,7 +28,7 @@ import com.intellij.openapi.vfs.VirtualFile
 import com.intellij.openapi.vfs.ex.temp.TempFileSystem
 import com.intellij.util.concurrency.annotations.RequiresEdt
 import com.intellij.util.withScheme
-import com.sourcegraph.cody.agent.protocol_extensions.toOffset
+import com.sourcegraph.cody.agent.protocol_extensions.toOffsetRange
 import com.sourcegraph.cody.agent.protocol_generated.Range
 import com.sourcegraph.config.ConfigUtil
 import java.net.URI
@@ -60,8 +60,7 @@ object CodyEditorUtil {
 
   @JvmStatic
   fun getTextRange(document: Document, range: Range): TextRange {
-    val start = range.start.toOffset(document)
-    val end = range.end.toOffset(document)
+    val (start, end) = range.toOffsetRange(document)
     return TextRange.create(start, end)
   }
 

src/test/kotlin/com/sourcegraph/cody/agent/protocol/PositionTest.kt

@@ -1,7 +1,10 @@
 package com.sourcegraph.cody.agent.protocol
 
 import com.intellij.testFramework.fixtures.BasePlatformTestCase
-import com.sourcegraph.cody.agent.protocol_extensions.*
+import com.sourcegraph.cody.agent.protocol_extensions.getRealColumn
+import com.sourcegraph.cody.agent.protocol_extensions.getRealLine
+import com.sourcegraph.cody.agent.protocol_extensions.isOutsideOfDocument
+import com.sourcegraph.cody.agent.protocol_extensions.toOffsetOrZero
 import com.sourcegraph.cody.agent.protocol_generated.Position
 
 class PositionTest : BasePlatformTestCase() {
@@ -18,19 +21,19 @@ class PositionTest : BasePlatformTestCase() {
 
   fun test_isStartOrEndOfDocumentMarkerReturnsTrueWhenLineIsLessThanZero() {
     val position = Position(-1, 0)
-    val result = position.isStartOrEndOfDocumentMarker(document)
+    val result = position.isOutsideOfDocument(document)
     assertEquals(true, result)
   }
 
   fun test_isStartOrEndOfDocumentMarkerReturnsTrueWhenLineIsGreaterThanLineCount() {
     val position = Position(3, 0)
-    val result = position.isStartOrEndOfDocumentMarker(document)
+    val result = position.isOutsideOfDocument(document)
     assertEquals(true, result)
   }
 
   fun test_isStartOrEndOfDocumentMarkerReturnsFalseWhenLineIsWithinBounds() {
     val position = Position(1, 0)
-    val result = position.isStartOrEndOfDocumentMarker(document)
+    val result = position.isOutsideOfDocument(document)
     assertEquals(false, result)
   }
 
@@ -64,9 +67,9 @@ class PositionTest : BasePlatformTestCase() {
     myFixture.openFileInEditor(file)
     val document = myFixture.editor.document
 
-    assertEquals(0, Position(0, 0).toOffset(document))
-    assertEquals(0, Position(0, 3).toOffset(document))
-    assertEquals(0, Position(1, 1).toOffset(document))
+    assertEquals(0, Position(0, 0).toOffsetOrZero(document))
+    assertEquals(0, Position(0, 3).toOffsetOrZero(document))
+    assertEquals(0, Position(1, 1).toOffsetOrZero(document))
   }
 
   fun test_toOffsetReturnsCorrectOffsetOnOneNewlineFile() {
@@ -75,16 +78,16 @@ class PositionTest : BasePlatformTestCase() {
     myFixture.openFileInEditor(file)
     val document = myFixture.editor.document
 
-    assertEquals(0, Position(0, 0).toOffset(document))
-    assertEquals(0, Position(0, 3).toOffset(document))
-    assertEquals(1, Position(2, 2).toOffset(document))
+    assertEquals(0, Position(0, 0).toOffsetOrZero(document))
+    assertEquals(0, Position(0, 3).toOffsetOrZero(document))
+    assertEquals(1, Position(2, 2).toOffsetOrZero(document))
   }
 
   fun test_toOffsetReturnsCorrectOffset() {
-    assertEquals(0, Position(0, 0).toOffset(document))
-    assertEquals(2, Position(0, 2).toOffset(document))
-    assertEquals(5, Position(0, 5).toOffset(document))
-    assertEquals(6, Position(1, 0).toOffset(document))
-    assertEquals(content.length, Position(12, 12).toOffset(document))
+    assertEquals(0, Position(0, 0).toOffsetOrZero(document))
+    assertEquals(2, Position(0, 2).toOffsetOrZero(document))
+    assertEquals(5, Position(0, 5).toOffsetOrZero(document))
+    assertEquals(6, Position(1, 0).toOffsetOrZero(document))
+    assertEquals(content.length, Position(12, 12).toOffsetOrZero(document))
   }
 }

src/test/kotlin/com/sourcegraph/cody/agent/protocol/ProtocolTextDocumentTest.kt

@@ -1,11 +1,11 @@
 package com.sourcegraph.cody.agent.protocol
 
 import com.intellij.openapi.application.WriteAction
+import com.intellij.openapi.editor.Document
 import com.intellij.openapi.fileEditor.FileEditorProvider
 import com.intellij.openapi.fileEditor.impl.text.TextEditorProvider
 import com.intellij.openapi.vfs.VirtualFile
 import com.intellij.testFramework.fixtures.BasePlatformTestCase
-import com.sourcegraph.cody.agent.protocol_extensions.PositionFactory
 import com.sourcegraph.cody.agent.protocol_generated.Position
 import com.sourcegraph.cody.agent.protocol_generated.Range
 import com.sourcegraph.cody.listeners.EditorChangesBus
@@ -15,6 +15,12 @@ class ProtocolTextDocumentTest : BasePlatformTestCase() {
   private val filename = "test.txt"
   private val file: VirtualFile by lazy { myFixture.createFile(filename, content) }
 
+  fun fromOffset(document: Document, offset: Int): Position {
+    val line = document.getLineNumber(offset)
+    val lineStartOffset = document.getLineStartOffset(line)
+    return Position(line.toLong(), (offset - lineStartOffset).toLong())
+  }
+
   override fun setUp() {
     super.setUp()
     file.putUserData(FileEditorProvider.KEY, TextEditorProvider.getInstance())
@@ -120,8 +126,8 @@ class ProtocolTextDocumentTest : BasePlatformTestCase() {
       val offset = if (isAppend) currentContent.length else 0
       assertEquals(
           Range(
-              PositionFactory.fromOffset(myFixture.editor.document, offset),
-              PositionFactory.fromOffset(myFixture.editor.document, offset),
+              fromOffset(myFixture.editor.document, offset),
+              fromOffset(myFixture.editor.document, offset),
           ),
           lastTextDocument!!.contentChanges!!.first().range)
     }
@@ -141,10 +147,8 @@ class ProtocolTextDocumentTest : BasePlatformTestCase() {
 
       val removalStartOffset = myFixture.editor.document.text.indexOf(removedContent)
       val removalEndOffset = removalStartOffset + removedContent.length
-      val removalStartPosition =
-          PositionFactory.fromOffset(myFixture.editor.document, removalStartOffset)
-      val removalEndPosition =
-          PositionFactory.fromOffset(myFixture.editor.document, removalEndOffset)
+      val removalStartPosition = fromOffset(myFixture.editor.document, removalStartOffset)
+      val removalEndPosition = fromOffset(myFixture.editor.document, removalEndOffset)
 
       WriteAction.run<RuntimeException> { file.setBinaryContent(newContent.toByteArray()) }
 
@@ -173,8 +177,8 @@ class ProtocolTextDocumentTest : BasePlatformTestCase() {
 
       val startOffset = currentContent.indexOf(oldSubstring)
       val endOffset = startOffset + oldSubstring.length
-      val startPosition = PositionFactory.fromOffset(myFixture.editor.document, startOffset)
-      val endPosition = PositionFactory.fromOffset(myFixture.editor.document, endOffset)
+      val startPosition = fromOffset(myFixture.editor.document, startOffset)
+      val endPosition = fromOffset(myFixture.editor.document, endOffset)
 
       WriteAction.run<RuntimeException> { file.setBinaryContent(newContent.toByteArray()) }
 
@@ -199,8 +203,8 @@ class ProtocolTextDocumentTest : BasePlatformTestCase() {
     fun insert(afterSubstring: String, insertSubstring: String) {
       val currentContent = myFixture.editor.document.text
       val startOffset = currentContent.indexOf(afterSubstring)
-      val startPosition = PositionFactory.fromOffset(myFixture.editor.document, startOffset)
-      val endPosition = PositionFactory.fromOffset(myFixture.editor.document, startOffset)
+      val startPosition = fromOffset(myFixture.editor.document, startOffset)
+      val endPosition = fromOffset(myFixture.editor.document, startOffset)
 
       val newContent =
           StringBuilder(currentContent).apply { insert(startOffset, insertSubstring) }.toString()