Rewrite key events handling, so it will require less iterations

foundation/src/main/kotlin/org/jetbrains/jewel/foundation/lazy/SelectableColumnOnKeyEvent.kt

@@ -12,10 +12,13 @@ interface SelectableColumnOnKeyEvent {
      * Select First Node
      */
     fun onSelectFirstItem(allKeys: List<SelectableLazyListKey>, state: SelectableLazyListState) {
-        val firstSelectable = allKeys.withIndex().firstOrNull { it.value is Selectable }
-        if (firstSelectable != null) {
-            state.selectedKeys = listOf(firstSelectable.value.key)
-            state.lastActiveItemIndex = firstSelectable.index
+        for (index in allKeys.indices) {
+            val key = allKeys[index]
+            if (key is Selectable) {
+                state.selectedKeys = listOf(key.key)
+                state.lastActiveItemIndex = index
+                return
+            }
         }
     }
 
@@ -48,12 +51,14 @@ interface SelectableColumnOnKeyEvent {
      * Select Last Node inherited from Move Caret to Text End
      */
     fun onSelectLastItem(keys: List<SelectableLazyListKey>, state: SelectableLazyListState) {
-        keys.withIndex()
-            .lastOrNull { it.value is Selectable }
-            ?.let {
-                state.selectedKeys = listOf(it)
-                state.lastActiveItemIndex = it.index
+        for (index in keys.lastIndex downTo 0) {
+            val key = keys[index]
+            if (key is Selectable) {
+                state.selectedKeys = listOf(key.key)
+                state.lastActiveItemIndex = index
+                return
             }
+        }
     }
 
     /**
@@ -76,18 +81,14 @@ interface SelectableColumnOnKeyEvent {
      * Select Previous Node inherited from Up
      */
     fun onSelectPreviousItem(keys: List<SelectableLazyListKey>, state: SelectableLazyListState) {
-        state.lastActiveItemIndex?.let { lastActiveIndex ->
-            if (lastActiveIndex == 0) return@let
-            keys
-                .withIndex()
-                .toList()
-                .dropLastWhile { it.index >= lastActiveIndex }
-                .reversed()
-                .firstOrNull { it.value is Selectable }
-                ?.let { (index, selectableKey) ->
-                    state.selectedKeys = listOf(selectableKey.key)
-                    state.lastActiveItemIndex = index
-                }
+        val initialIndex = state.lastActiveItemIndex ?: return
+        for (index in initialIndex - 1 downTo 0) {
+            val key = keys[index]
+            if (key is Selectable) {
+                state.selectedKeys = listOf(key.key)
+                state.lastActiveItemIndex = index
+                return
+            }
         }
     }
 
@@ -114,16 +115,14 @@ interface SelectableColumnOnKeyEvent {
      * Select Next Node inherited from Down
      */
     fun onSelectNextItem(keys: List<SelectableLazyListKey>, state: SelectableLazyListState) {
-        state.lastActiveItemIndex?.let { lastActiveIndex ->
-            if (lastActiveIndex == keys.lastIndex) return@let
-            keys
-                .withIndex()
-                .dropWhile { it.index <= lastActiveIndex }
-                .firstOrNull { it.value is Selectable }
-                ?.let { (index, selectableKey) ->
-                    state.selectedKeys = listOf(selectableKey.key)
-                    state.lastActiveItemIndex = index
-                }
+        val initialIndex = state.lastActiveItemIndex ?: return
+        for (index in initialIndex + 1..keys.lastIndex) {
+            val key = keys[index]
+            if (key is Selectable) {
+                state.selectedKeys = listOf(key.key)
+                state.lastActiveItemIndex = index
+                return
+            }
         }
     }
 

foundation/src/test/kotlin/org/jetbrains/jewel/foundation/lazy/SelectableLazyColumnTest.kt

@@ -5,11 +5,17 @@ import androidx.compose.foundation.layout.requiredHeight
 import androidx.compose.foundation.lazy.LazyListState
 import androidx.compose.foundation.text.BasicText
 import androidx.compose.ui.Modifier
+import androidx.compose.ui.input.key.Key
 import androidx.compose.ui.platform.testTag
+import androidx.compose.ui.test.ExperimentalTestApi
 import androidx.compose.ui.test.junit4.createComposeRule
 import androidx.compose.ui.test.onNodeWithTag
+import androidx.compose.ui.test.performClick
+import androidx.compose.ui.test.performKeyInput
 import androidx.compose.ui.unit.dp
 import kotlinx.coroutines.runBlocking
+import org.junit.Assert.assertEquals
+import org.junit.Assert.assertTrue
 import org.junit.Rule
 import org.junit.Test
 
@@ -53,4 +59,69 @@ internal class SelectableLazyColumnTest {
         scrollState.scrollToItem(20)
         composeRule.onNodeWithTag("Item 20").assertExists()
     }
+
+    @OptIn(ExperimentalTestApi::class)
+    @Test
+    fun `selection with arrow keys`() = runBlocking {
+        val items = (0..10).toList()
+        val state = SelectableLazyListState(LazyListState())
+        composeRule.setContent {
+            Box(modifier = Modifier.requiredHeight(300.dp)) {
+                SelectableLazyColumn(state = state, modifier = Modifier.testTag("list")) {
+                    items(
+                        items.size,
+                        key = {
+                            items[it]
+                        },
+                    ) {
+                        val itemText = "Item ${items[it]}"
+                        BasicText(itemText, modifier = Modifier.testTag(itemText))
+                    }
+                }
+            }
+        }
+        composeRule.awaitIdle()
+        // select item 5 by click
+        composeRule.onNodeWithTag("Item 5").assertExists()
+        composeRule.onNodeWithTag("Item 5").performClick()
+
+        // check that 5th element is selected
+        assertEquals(1, state.selectedKeys.size)
+        assertEquals(items[5], state.selectedKeys.single())
+
+        // press arrow up and check that selected key is changed
+        repeat(20) { step ->
+            composeRule.onNodeWithTag("list").performKeyInput {
+                keyDown(Key.DirectionUp)
+                keyUp(Key.DirectionUp)
+            }
+
+            // check that previous element is selected
+            // when started from 5th element
+            assertTrue(state.selectedKeys.size == 1)
+            val expectedSelectedIndex = (5 - step - 1).takeIf { it >= 0 } ?: 0
+            assertEquals(items[expectedSelectedIndex], state.selectedKeys.single())
+        }
+
+        // since amount of arrow up is bigger than amount of items -> first element should be selected
+        assertTrue(state.selectedKeys.size == 1)
+        assertEquals(items[0], state.selectedKeys.single())
+
+        // press arrow down and check that selected key is changed
+        repeat(40) { step ->
+            composeRule.onNodeWithTag("list").performKeyInput {
+                keyDown(Key.DirectionDown)
+                keyUp(Key.DirectionDown)
+            }
+
+            // check that next element is selected
+            assertTrue(state.selectedKeys.size == 1)
+            val expectedSelectedIndex = (step + 1).takeIf { it in items.indices } ?: items.lastIndex
+            assertEquals(items[expectedSelectedIndex], state.selectedKeys.single())
+        }
+
+        // since amount of arrow down is bigger than amount of items -> last element should be selected
+        assertTrue(state.selectedKeys.size == 1)
+        assertEquals(items.last(), state.selectedKeys.single())
+    }
 }