From d105d9ee03391186f921aa8f1c1f77a49421de7f Mon Sep 17 00:00:00 2001
From: Nikolai Rykunov <Nikolay.Rykunov@jetbrains.com>
Date: Fri, 3 Nov 2023 16:27:16 +0100
Subject: [PATCH] Fix error appearing when selectedKeys contain invalid key
 (#244)

* Fix error appearing when selectedKeys contain invalid key

In this case, we should just skip invalid keys when receiving their indices

* Cleanup
---
 .../foundation/lazy/SelectableLazyColumn.kt   |  2 +-
 .../lazy/SelectableLazyListScope.kt           |  2 +-
 .../lazy/SelectableLazyColumnTest.kt          | 46 +++++++++++++++++++
 3 files changed, 48 insertions(+), 2 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 dc2f4c1ca..4961ebbb9 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
@@ -63,7 +63,7 @@ public fun SelectableLazyColumn(
     val latestOnSelectedIndexesChanged = rememberUpdatedState(onSelectedIndexesChanged)
     LaunchedEffect(state, container) {
         snapshotFlow { state.selectedKeys }.collect { selectedKeys ->
-            val indices = selectedKeys.map { key -> container.getKeyIndex(key) }
+            val indices = selectedKeys.mapNotNull { key -> container.getKeyIndex(key) }
             latestOnSelectedIndexesChanged.value.invoke(indices)
         }
     }
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 dc16a5e9b..bf9cce09b 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
@@ -107,7 +107,7 @@ internal class SelectableLazyListScopeContainer : SelectableLazyListScope {
         ) : Entry
     }
 
-    internal fun getKeyIndex(key: Any): Int = keyToIndex[key] ?: error("Cannot find index of '$key'")
+    internal fun getKeyIndex(key: Any): Int? = keyToIndex[key]
 
     internal fun isKeySelectable(key: Any): Boolean = key !in nonSelectableKeys
 
diff --git a/foundation/src/test/kotlin/org/jetbrains/jewel/foundation/lazy/SelectableLazyColumnTest.kt b/foundation/src/test/kotlin/org/jetbrains/jewel/foundation/lazy/SelectableLazyColumnTest.kt
index 77615063c..f1207e9d2 100644
--- a/foundation/src/test/kotlin/org/jetbrains/jewel/foundation/lazy/SelectableLazyColumnTest.kt
+++ b/foundation/src/test/kotlin/org/jetbrains/jewel/foundation/lazy/SelectableLazyColumnTest.kt
@@ -4,6 +4,7 @@ import androidx.compose.foundation.layout.Box
 import androidx.compose.foundation.layout.requiredHeight
 import androidx.compose.foundation.lazy.LazyListState
 import androidx.compose.foundation.text.BasicText
+import androidx.compose.runtime.mutableStateOf
 import androidx.compose.ui.Modifier
 import androidx.compose.ui.input.key.Key
 import androidx.compose.ui.platform.testTag
@@ -255,4 +256,49 @@ internal class SelectableLazyColumnTest {
         assertEquals(expectedElementsAfterPageDown.size, state.selectedKeys.size)
         assertEquals(expectedElementsAfterPageDown.toSet(), state.selectedKeys.toSet())
     }
+
+    @Test
+    fun `changing items model with selection shouldn't fail`() = runBlocking {
+        val items1 = (0..50).toList()
+        val items2 = (70..80).toList()
+        val currentItems = mutableStateOf(items1)
+        val state = SelectableLazyListState(LazyListState())
+        composeRule.setContent {
+            Box(modifier = Modifier.requiredHeight(300.dp)) {
+                val items = currentItems.value
+                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(currentItems.value[5], state.selectedKeys.single())
+
+        // change items from 0..50 to 70..80, so "Item 5" doesn't exist
+        currentItems.value = items2
+        composeRule.awaitIdle()
+        // TODO: should the selectedKeys be cleared when items are changed
+        //  https://github.com/JetBrains/jewel/issues/242
+        // assertEquals(0, state.selectedKeys.size)
+
+        composeRule.onNodeWithTag("Item 75").assertExists()
+        composeRule.onNodeWithTag("Item 75").performClick()
+
+        assertEquals(1, state.selectedKeys.size)
+        assertEquals(currentItems.value[5], state.selectedKeys.single())
+    }
 }