From 8c646c1078948f5d88f7e3bbaa54b6d0aad64498 Mon Sep 17 00:00:00 2001
From: Nikolay Rykunov <nikolay.rykunov@jetbrains.com>
Date: Thu, 2 Nov 2023 17:14:35 +0100
Subject: [PATCH] Fix events handling

---
 .../jewel/foundation/lazy/tree/KeyActions.kt  | 51 ++++++++++++++-----
 1 file changed, 38 insertions(+), 13 deletions(-)

diff --git a/foundation/src/main/kotlin/org/jetbrains/jewel/foundation/lazy/tree/KeyActions.kt b/foundation/src/main/kotlin/org/jetbrains/jewel/foundation/lazy/tree/KeyActions.kt
index 5053fbbf6..da45561e7 100644
--- a/foundation/src/main/kotlin/org/jetbrains/jewel/foundation/lazy/tree/KeyActions.kt
+++ b/foundation/src/main/kotlin/org/jetbrains/jewel/foundation/lazy/tree/KeyActions.kt
@@ -278,6 +278,24 @@ open class DefaultSelectableLazyColumnKeyActions(
         keys: List<SelectableLazyListKey>,
         state: SelectableLazyListState,
         selectionMode: SelectionMode,
+    ): Boolean {
+        val singleSelectionEventHandled = handleSingleSelectionEvents(keys, state)
+        if (singleSelectionEventHandled) {
+            return true
+        }
+        if (selectionMode == SelectionMode.Multiple) {
+            val multipleSelectionEventHandled = handleMultipleSelectionEvents(keys, state)
+            if (multipleSelectionEventHandled) {
+                return true
+            }
+        }
+        return false
+    }
+
+    context(SelectableColumnKeybindings, SelectableColumnOnKeyEvent)
+    private fun KeyEvent.handleSingleSelectionEvents(
+        keys: List<SelectableLazyListKey>,
+        state: SelectableLazyListState,
     ): Boolean {
         when {
             isSelectNextItem -> onSelectNextItem(keys, state)
@@ -285,20 +303,27 @@ open class DefaultSelectableLazyColumnKeyActions(
             isSelectFirstItem -> onSelectFirstItem(keys, state)
             isSelectLastItem -> onSelectLastItem(keys, state)
             isEdit -> onEdit()
+            else -> return false
         }
-        if (selectionMode == SelectionMode.Multiple) {
-            when {
-                isExtendSelectionToFirstItem -> onExtendSelectionToFirst(keys, state)
-                isExtendSelectionToLastItem -> onExtendSelectionToLastItem(keys, state)
-                isExtendSelectionWithNextItem -> onExtendSelectionWithNextItem(keys, state)
-                isExtendSelectionWithPreviousItem -> onExtendSelectionWithPreviousItem(keys, state)
-                isScrollPageDownAndExtendSelection -> onScrollPageDownAndExtendSelection(keys, state)
-                isScrollPageDownAndSelectItem -> onScrollPageDownAndSelectItem(keys, state)
-                isScrollPageUpAndExtendSelection -> onScrollPageUpAndExtendSelection(keys, state)
-                isScrollPageUpAndSelectItem -> onScrollPageUpAndSelectItem(keys, state)
-                isSelectAll -> onSelectAll(keys, state)
-                else -> return false
-            }
+        return true
+    }
+
+    context(SelectableColumnKeybindings, SelectableColumnOnKeyEvent)
+    private fun KeyEvent.handleMultipleSelectionEvents(
+        keys: List<SelectableLazyListKey>,
+        state: SelectableLazyListState,
+    ): Boolean {
+        when {
+            isExtendSelectionToFirstItem -> onExtendSelectionToFirst(keys, state)
+            isExtendSelectionToLastItem -> onExtendSelectionToLastItem(keys, state)
+            isExtendSelectionWithNextItem -> onExtendSelectionWithNextItem(keys, state)
+            isExtendSelectionWithPreviousItem -> onExtendSelectionWithPreviousItem(keys, state)
+            isScrollPageDownAndExtendSelection -> onScrollPageDownAndExtendSelection(keys, state)
+            isScrollPageDownAndSelectItem -> onScrollPageDownAndSelectItem(keys, state)
+            isScrollPageUpAndExtendSelection -> onScrollPageUpAndExtendSelection(keys, state)
+            isScrollPageUpAndSelectItem -> onScrollPageUpAndSelectItem(keys, state)
+            isSelectAll -> onSelectAll(keys, state)
+            else -> return false
         }
         return true
     }