Fix offset conversions at the end of the document (#2203)

Fixes #2196

## Changes

Conversion of the end of range to offset was buggy.
VSC sometimes sends `-1` as line index. If the start or end position of
the range is outside the document, the corresponding offset should be
set to 0 or the document's text length, respectively.
But in case of the end range we were setting it to `0`.

To avoid that mistake:
* I changed `Position::toOffset` method name to
`Position::toOffsetOrZero` to be very explicit of what it does. It
should not be used for ranges, but sometimes we only have a single
position (like for a text insert).
* I added `Range::toOffsetRange` which returns pair of offsets,
correcting them appropriately depending if position is a start or end

## Test plan

1. Login to Cody with Free/Pro user
2. Create new empty python file called quicksort.py, make sure it's
empty
3. Ask Cody to "create a quicksort function"
4. Click on Apply, make sure code is added to source code file
5. Click on Reject

Edit should be properly rejected and file should be empty again.

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()