From 266a1a4480a6dc481c8be40d0578b950dfbf7703 Mon Sep 17 00:00:00 2001
From: Nikolay Rykunov <nikolay.rykunov@jetbrains.com>
Date: Wed, 1 Nov 2023 17:58:26 +0100
Subject: [PATCH] Speed up selection in SelectableLazyColumn

---
 .../foundation/lazy/SelectableLazyColumn.kt   | 19 ++++++++++---------
 .../lazy/SelectableLazyListScope.kt           | 10 ++++++++++
 2 files changed, 20 insertions(+), 9 deletions(-)

diff --git a/foundation/src/main/kotlin/org/jetbrains/jewel/foundation/lazy/SelectableLazyColumn.kt b/foundation/src/main/kotlin/org/jetbrains/jewel/foundation/lazy/SelectableLazyColumn.kt
index 08da9d3c4..81afdd871 100644
--- a/foundation/src/main/kotlin/org/jetbrains/jewel/foundation/lazy/SelectableLazyColumn.kt
+++ b/foundation/src/main/kotlin/org/jetbrains/jewel/foundation/lazy/SelectableLazyColumn.kt
@@ -9,11 +9,14 @@ import androidx.compose.foundation.layout.Box
 import androidx.compose.foundation.layout.PaddingValues
 import androidx.compose.foundation.lazy.LazyColumn
 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
@@ -23,6 +26,7 @@ import androidx.compose.ui.input.key.onPreviewKeyEvent
 import androidx.compose.ui.input.pointer.PointerEventType
 import androidx.compose.ui.input.pointer.pointerInput
 import androidx.compose.ui.unit.dp
+import kotlinx.coroutines.flow.collectLatest
 import kotlinx.coroutines.launch
 import org.jetbrains.jewel.foundation.lazy.SelectableLazyListScopeContainer.Entry
 import org.jetbrains.jewel.foundation.lazy.tree.DefaultSelectableLazyColumnEventAction
@@ -59,15 +63,12 @@ fun SelectableLazyColumn(
     }
     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()
-    }
-
-    remember(state.selectedKeys) {
-        onSelectedIndexesChanged(evaluateIndexes())
+    val latestOnSelectedIndexesChanged = rememberUpdatedState(onSelectedIndexesChanged)
+    LaunchedEffect(state, container) {
+        snapshotFlow { state.selectedKeys }.collect { selectedKeys ->
+            val indices = selectedKeys.map { key -> container.getKeyIndex(key) }
+            latestOnSelectedIndexesChanged.value.invoke(indices)
+        }
     }
     val focusRequester = remember { FocusRequester() }
     LazyColumn(
diff --git a/foundation/src/main/kotlin/org/jetbrains/jewel/foundation/lazy/SelectableLazyListScope.kt b/foundation/src/main/kotlin/org/jetbrains/jewel/foundation/lazy/SelectableLazyListScope.kt
index ebaef98ab..747264536 100644
--- a/foundation/src/main/kotlin/org/jetbrains/jewel/foundation/lazy/SelectableLazyListScope.kt
+++ b/foundation/src/main/kotlin/org/jetbrains/jewel/foundation/lazy/SelectableLazyListScope.kt
@@ -74,6 +74,9 @@ internal class SelectableLazyListScopeContainer : SelectableLazyListScope {
      */
     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>()
 
@@ -101,6 +104,10 @@ internal class SelectableLazyListScopeContainer : SelectableLazyListScope {
         ) : Entry
     }
 
+    internal fun getKeyIndex(key: Any): Int {
+        return keyToIndex[key] ?: error("Cannot find index of '$key'")
+    }
+
     internal fun isKeySelectable(key: Any): Boolean {
         return key !in nonSelectableKeys
     }
@@ -111,6 +118,7 @@ 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)
@@ -132,6 +140,7 @@ internal class SelectableLazyListScopeContainer : SelectableLazyListScope {
             if (!isSelectable) {
                 nonSelectableKeys.add(currentKey)
             }
+            keyToIndex[currentKey] = keys.size
             keys.add(if (isSelectable) Selectable(currentKey) else NotSelectable(currentKey))
         }
         entries.add(Entry.Items(count, key, contentType, itemContent))
@@ -144,6 +153,7 @@ 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)