Optimize performance of SelectableLazyColumn

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

@@ -15,10 +15,13 @@ public interface SelectableColumnOnKeyEvent {
         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
+            }
         }
     }
 
@@ -30,23 +33,20 @@ public interface SelectableColumnOnKeyEvent {
         keys: List<SelectableLazyListKey>,
         state: SelectableLazyListState,
     ) {
-        state.lastActiveItemIndex?.let {
-            val iterator = keys.listIterator(it)
-            val list = buildList {
-                while (iterator.hasPrevious()) {
-                    val previous = iterator.previous()
-                    if (previous is Selectable) {
-                        add(previous.key)
-                        state.lastActiveItemIndex = (iterator.previousIndex() + 1).coerceAtMost(keys.size)
-                    }
-                }
-            }
-            if (list.isNotEmpty()) {
-                state.selectedKeys =
-                    state.selectedKeys.toMutableList()
-                        .also { selectionList -> selectionList.addAll(list) }
+        val initialIndex = state.lastActiveItemIndex ?: return
+        val newSelection = ArrayList<Any>(max(initialIndex, state.selectedKeys.size)).apply {
+            addAll(state.selectedKeys)
+        }
+        var lastActiveItemIndex = initialIndex
+        for (index in initialIndex - 1 downTo 0) {
+            val key = keys[index]
+            if (key is Selectable) {
+                newSelection.add(key.key)
+                lastActiveItemIndex = index
             }
         }
+        state.lastActiveItemIndex = lastActiveItemIndex
+        state.selectedKeys = newSelection
     }
 
     /**
@@ -56,12 +56,14 @@ public interface SelectableColumnOnKeyEvent {
         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
             }
+        }
     }
 
     /**
@@ -72,16 +74,20 @@ public interface SelectableColumnOnKeyEvent {
         keys: List<SelectableLazyListKey>,
         state: SelectableLazyListState,
     ) {
-        state.lastActiveItemIndex?.let {
-            val list = mutableListOf<Any>(state.selectedKeys)
-            keys.subList(it, keys.lastIndex).forEachIndexed { index, selectableLazyListKey ->
-                if (selectableLazyListKey is Selectable) {
-                    list.add(selectableLazyListKey.key)
-                }
-                state.lastActiveItemIndex = index
+        val initialIndex = state.lastActiveItemIndex ?: return
+        val newSelection = ArrayList<Any>(max(keys.size - initialIndex, state.selectedKeys.size)).apply {
+            addAll(state.selectedKeys)
+        }
+        var lastActiveItemIndex = initialIndex
+        for (index in initialIndex + 1..keys.lastIndex) {
+            val key = keys[index]
+            if (key is Selectable) {
+                newSelection.add(key.key)
+                lastActiveItemIndex = index
             }
-            state.selectedKeys = list
         }
+        state.lastActiveItemIndex = lastActiveItemIndex
+        state.selectedKeys = newSelection
     }
 
     /**
@@ -91,17 +97,14 @@ public interface SelectableColumnOnKeyEvent {
         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
+            }
         }
     }
 
@@ -112,17 +115,15 @@ public interface SelectableColumnOnKeyEvent {
         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 = state.selectedKeys + selectableKey.key
-                    state.lastActiveItemIndex = index
-                }
+        // todo we need deselect if we are changing direction
+        val initialIndex = state.lastActiveItemIndex ?: return
+        for (index in initialIndex - 1 downTo 0) {
+            val key = keys[index]
+            if (key is Selectable) {
+                state.selectedKeys += key.key
+                state.lastActiveItemIndex = index
+                return
+            }
         }
     }
 
@@ -133,15 +134,14 @@ public interface SelectableColumnOnKeyEvent {
         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
+            }
         }
     }
 
@@ -153,16 +153,14 @@ public interface SelectableColumnOnKeyEvent {
         state: SelectableLazyListState,
     ) {
         // todo we need deselect if we are changing direction
-        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 = state.selectedKeys + 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 += key.key
+                state.lastActiveItemIndex = index
+                return
+            }
         }
     }
 

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

@@ -10,11 +10,14 @@ import androidx.compose.foundation.layout.PaddingValues
 import androidx.compose.foundation.lazy.LazyColumn
 import androidx.compose.foundation.lazy.LazyListScope
 import androidx.compose.runtime.Composable
+import androidx.compose.runtime.LaunchedEffect
 import androidx.compose.runtime.getValue
 import androidx.compose.runtime.mutableStateOf
 import androidx.compose.runtime.remember
 import androidx.compose.runtime.rememberCoroutineScope
+import androidx.compose.runtime.rememberUpdatedState
 import androidx.compose.runtime.setValue
+import androidx.compose.runtime.snapshotFlow
 import androidx.compose.ui.Alignment
 import androidx.compose.ui.Modifier
 import androidx.compose.ui.focus.FocusRequester
@@ -57,12 +60,13 @@ public fun SelectableLazyColumn(
     val keys = remember(container) { container.getKeys() }
     var isFocused by remember { mutableStateOf(false) }
 
-    fun evaluateIndexes(): List<Int> {
-        val keyToIndexMap = keys.withIndex().associateBy({ it.value.key }, { it.index })
-        return state.selectedKeys.mapNotNull { selected -> keyToIndexMap[selected] }.sorted()
+    val latestOnSelectedIndexesChanged = rememberUpdatedState(onSelectedIndexesChanged)
+    LaunchedEffect(state, container) {
+        snapshotFlow { state.selectedKeys }.collect { selectedKeys ->
+            val indices = selectedKeys.map { key -> container.getKeyIndex(key) }
+            latestOnSelectedIndexesChanged.value.invoke(indices)
+        }
     }
-
-    remember(state.selectedKeys) { onSelectedIndexesChanged(evaluateIndexes()) }
     val focusRequester = remember { FocusRequester() }
     LazyColumn(
         modifier = modifier.onFocusChanged { isFocused = it.hasFocus }
@@ -87,12 +91,22 @@ public fun SelectableLazyColumn(
         flingBehavior = flingBehavior,
     ) {
         container.getEntries().forEach { entry ->
-            AppendEntry(entry, state, isFocused, keys, focusRequester, keyActions, pointerEventActions, selectionMode)
+            appendEntry(
+                entry,
+                state,
+                isFocused,
+                keys,
+                focusRequester,
+                keyActions,
+                pointerEventActions,
+                selectionMode,
+                container::isKeySelectable,
+            )
         }
     }
 }
 
-private fun LazyListScope.AppendEntry(
+private fun LazyListScope.appendEntry(
     entry: Entry,
     state: SelectableLazyListState,
     isFocused: Boolean,
@@ -101,6 +115,7 @@ private fun LazyListScope.AppendEntry(
     keyActions: KeyActions,
     pointerEventActions: PointerEventActions,
     selectionMode: SelectionMode,
+    isKeySelectable: (Any) -> Boolean,
 ) {
     when (entry) {
         is Entry.Item -> item(entry.key, entry.contentType) {
@@ -109,7 +124,7 @@ private fun LazyListScope.AppendEntry(
                     isSelected = entry.key in state.selectedKeys,
                     isActive = isFocused,
                 )
-            if (keys.any { it.key == entry.key && it is SelectableLazyListKey.Selectable }) {
+            if (isKeySelectable(entry.key)) {
                 Box(
                     modifier = Modifier.selectable(
                         requester = focusRequester,
@@ -133,10 +148,9 @@ private fun LazyListScope.AppendEntry(
             key = { entry.key(it) },
             contentType = { entry.contentType(it) },
         ) { index ->
-            val itemScope =
-                SelectableLazyItemScope(entry.key(index) in state.selectedKeys, isFocused)
-
-            if (keys.any { it.key == entry.key(index) && it is SelectableLazyListKey.Selectable }) {
+            val key = remember(entry, index) { entry.key(index) }
+            val itemScope = SelectableLazyItemScope(key in state.selectedKeys, isFocused)
+            if (isKeySelectable(key)) {
                 Box(
                     modifier = Modifier.selectable(
                         requester = focusRequester,
@@ -158,7 +172,7 @@ private fun LazyListScope.AppendEntry(
         is Entry.StickyHeader -> stickyHeader(entry.key, entry.contentType) {
             val itemScope = SelectableLazyItemScope(entry.key in state.selectedKeys, isFocused)
 
-            if (keys.any { it.key == entry.key && it is SelectableLazyListKey.Selectable }) {
+            if (isKeySelectable(entry.key)) {
                 Box(
                     modifier = Modifier.selectable(
                         keybindings = keyActions.keybindings,

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

@@ -67,6 +67,18 @@ public interface SelectableLazyListScope {
 
 internal class SelectableLazyListScopeContainer : SelectableLazyListScope {
 
+    /**
+     * Provides a set of keys that cannot be selected.
+     * Here we use an assumption that amount of selectable items >> amount of non-selectable items.
+     * So, for optimization we will keep only this set.
+     *
+     * @see [isKeySelectable]
+     */
+    private val nonSelectableKeys = hashSetOf<Any>()
+
+    // TODO: [performance] we can get rid of that map if indices won't be used at all in the API
+    private val keyToIndex = hashMapOf<Any, Int>()
+
     private val keys = mutableListOf<SelectableLazyListKey>()
     private val entries = mutableListOf<Entry>()
 
@@ -95,13 +107,21 @@ internal class SelectableLazyListScopeContainer : SelectableLazyListScope {
         ) : Entry
     }
 
+    internal fun getKeyIndex(key: Any): Int = keyToIndex[key] ?: error("Cannot find index of '$key'")
+
+    internal fun isKeySelectable(key: Any): Boolean = key !in nonSelectableKeys
+
     override fun item(
         key: Any,
         contentType: Any?,
         selectable: Boolean,
         content: @Composable (SelectableLazyItemScope.() -> Unit),
     ) {
+        keyToIndex[key] = keys.size
         keys.add(if (selectable) Selectable(key) else NotSelectable(key))
+        if (!selectable) {
+            nonSelectableKeys.add(key)
+        }
         entries.add(Entry.Item(key, contentType, content))
     }
 
@@ -112,15 +132,16 @@ internal class SelectableLazyListScopeContainer : SelectableLazyListScope {
         selectable: (index: Int) -> Boolean,
         itemContent: @Composable (SelectableLazyItemScope.(index: Int) -> Unit),
     ) {
-        val selectableKeys: List<SelectableLazyListKey> =
-            List(count) {
-                if (selectable(it)) {
-                    Selectable(key(it))
-                } else {
-                    NotSelectable(key(it))
-                }
+        // TODO: [performance] now the implementation requires O(count) operations but should be done in ~ O(1)
+        for (index in 0 until count) {
+            val isSelectable = selectable(index)
+            val currentKey = key(index)
+            if (!isSelectable) {
+                nonSelectableKeys.add(currentKey)
             }
-        keys.addAll(selectableKeys)
+            keyToIndex[currentKey] = keys.size
+            keys.add(if (isSelectable) Selectable(currentKey) else NotSelectable(currentKey))
+        }
         entries.add(Entry.Items(count, key, contentType, itemContent))
     }
 
@@ -131,7 +152,11 @@ internal class SelectableLazyListScopeContainer : SelectableLazyListScope {
         selectable: Boolean,
         content: @Composable (SelectableLazyItemScope.() -> Unit),
     ) {
+        keyToIndex[key] = keys.size
         keys.add(if (selectable) Selectable(key) else NotSelectable(key))
+        if (!selectable) {
+            nonSelectableKeys.add(key)
+        }
         entries.add(Entry.StickyHeader(key, contentType, content))
     }
 }