From 1eabe1c32bf5464d24d57bca37f8b510c3a131b6 Mon Sep 17 00:00:00 2001
From: Akash Yadav <itsaky01@gmail.com>
Date: Mon, 22 Jan 2024 23:08:19 +0530
Subject: [PATCH] fix(editor): preserve selection after executing 'Organize
 imports' action (fixes #1598)

---
 .../itsaky/androidide/editor/api/IEditor.java | 24 +++++++++++++++++++
 .../androidide/editor/ui/EditorFeatures.kt    | 11 +++++++++
 .../actions/common/OrganizeImportsAction.kt   |  8 +++++++
 3 files changed, 43 insertions(+)

diff --git a/editor-api/src/main/java/com/itsaky/androidide/editor/api/IEditor.java b/editor-api/src/main/java/com/itsaky/androidide/editor/api/IEditor.java
index 2afd2c003..3cbc3a28e 100644
--- a/editor-api/src/main/java/com/itsaky/androidide/editor/api/IEditor.java
+++ b/editor-api/src/main/java/com/itsaky/androidide/editor/api/IEditor.java
@@ -23,6 +23,7 @@ import androidx.annotation.Nullable;
 import com.itsaky.androidide.models.Position;
 import com.itsaky.androidide.models.Range;
 
+import io.github.rosemoe.sora.text.CharPosition;
 import java.io.File;
 
 /**
@@ -66,6 +67,29 @@ public interface IEditor {
   }
 
   void setSelection(@NonNull Position start, @NonNull Position end);
+
+  /**
+   * Set selection around the given position.
+   * @param position The position to set selection around.
+   */
+  default void setSelectionAround(CharPosition position) {
+    setSelectionAround(position.getLine(), position.getColumn());
+  }
+
+  /**
+   * Set selection around the given position.
+   * @param position The position to set selection around.
+   */
+  default void setSelectionAround(Position position) {
+    setSelectionAround(position.getLine(), position.getColumn());
+  }
+
+  /**
+   * Set selection around the given position.
+   * @param line The line index.
+   * @param column The column index.
+   */
+  void setSelectionAround(int line, int column);
   
   /**
    * Get the cursor's selection range in the form of {@link Range}.
diff --git a/editor/src/main/java/com/itsaky/androidide/editor/ui/EditorFeatures.kt b/editor/src/main/java/com/itsaky/androidide/editor/ui/EditorFeatures.kt
index a33955adc..c359601e7 100644
--- a/editor/src/main/java/com/itsaky/androidide/editor/ui/EditorFeatures.kt
+++ b/editor/src/main/java/com/itsaky/androidide/editor/ui/EditorFeatures.kt
@@ -53,6 +53,17 @@ class EditorFeatures(
     }
   }
 
+  override fun setSelectionAround(line: Int, column: Int) {
+    withEditor {
+      if (line < lineCount) {
+        val columnCount = text.getColumnCount(line)
+        setSelection(line, if (column > columnCount) columnCount else column)
+      } else {
+        setSelection(lineCount - 1, text.getColumnCount(lineCount - 1))
+      }
+    }
+  }
+
   override fun getCursorLSPRange(): Range = withEditor {
     val end = cursor.right().let {
       Position(line = it.line, column = it.column, index = it.index)
diff --git a/lsp/java/src/main/java/com/itsaky/androidide/lsp/java/actions/common/OrganizeImportsAction.kt b/lsp/java/src/main/java/com/itsaky/androidide/lsp/java/actions/common/OrganizeImportsAction.kt
index f4bf76e36..f027b44d0 100644
--- a/lsp/java/src/main/java/com/itsaky/androidide/lsp/java/actions/common/OrganizeImportsAction.kt
+++ b/lsp/java/src/main/java/com/itsaky/androidide/lsp/java/actions/common/OrganizeImportsAction.kt
@@ -6,6 +6,7 @@ import com.itsaky.androidide.actions.ActionData
 import com.itsaky.androidide.actions.hasRequiredData
 import com.itsaky.androidide.actions.markInvisible
 import com.itsaky.androidide.actions.requireEditor
+import com.itsaky.androidide.editor.api.IEditor
 import com.itsaky.androidide.lsp.java.JavaLanguageServer
 import com.itsaky.androidide.lsp.java.actions.BaseJavaCodeAction
 import com.itsaky.androidide.lsp.java.models.JavaServerSettings
@@ -56,10 +57,17 @@ class OrganizeImportsAction : BaseJavaCodeAction() {
     if (result is String) {
       if (result.isNotEmpty()) {
         val editor = data.requireEditor()
+        val cursor = editor.cursor.left()
+
         editor.text.apply {
           val endLine = getLine(lineCount - 1)
           replace(0, 0, lineCount - 1, endLine.length + endLine.lineSeparator.length, result)
         }
+
+        (editor as? IEditor?)?.also {
+          it.setSelectionAround(cursor)
+          editor.ensureSelectionVisible()
+        }
       }
     }
   }