Skip to content

Commit

Permalink
Optimize performance of SelectableLazyColumn (#240)
Browse files Browse the repository at this point in the history 
* Fix index out of bounds for items call in SelectableLazyColumn

* Remove redundant index from SelectableLazyColumn items

* Speed up scrolling in SelectableLazyColumn

Search in List shouldn't be performed inside Composable function

* Speed up selection in SelectableLazyColumn

* Rewrite key events handling, so it will require less iterations

* Rewrite multiple selection, speeding it up

Also, fix DefaultSelectableLazyColumnKeyActions check for multiple selection

* Cleanup code for lint checks

* Fix events handling

* Cleanup after merge

* Don't update state in loop

* Fix lint checks
  • Loading branch information
@Walingar
Walingar authored Nov 3, 2023
1 parent c345f7a commit 5e0ac3a
Show file tree
Hide file tree
Showing 6 changed files with 406 additions and 133 deletions.
144 changes: 71 additions & 73 deletions foundation/src/main/kotlin/org/jetbrains/jewel/foundation/lazy/SelectableColumnOnKeyEvent.kt
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
}
}
}

Expand All @@ -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
}

/**
Expand All @@ -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
}
}
}

/**
Expand All @@ -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
}

/**
Expand All @@ -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
}
}
}

Expand All @@ -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
}
}
}

Expand All @@ -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
}
}
}

Expand All @@ -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
}
}
}

Expand Down
40 changes: 27 additions & 13 deletions foundation/src/main/kotlin/org/jetbrains/jewel/foundation/lazy/SelectableLazyColumn.kt
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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 }
Expand All @@ -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,
Expand All @@ -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) {
Expand All @@ -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,
Expand All @@ -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,
Expand All @@ -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,
Expand Down
41 changes: 33 additions & 8 deletions foundation/src/main/kotlin/org/jetbrains/jewel/foundation/lazy/SelectableLazyListScope.kt
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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>()

Expand Down Expand Up @@ -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))
}

Expand All @@ -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))
}

Expand All @@ -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))
}
}
Expand Down
Loading

0 comments on commit 5e0ac3a

Please sign in to comment.