Convert CustomListLocations into flat LazyColumn

3 files changed, 130 insertions, 34 deletions

diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/screen/CustomListLocationsScreen.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/screen/CustomListLocationsScreen.kt
index 19b548153a..cb3767784e 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/screen/CustomListLocationsScreen.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/screen/CustomListLocationsScreen.kt
@@ -1,7 +1,6 @@
 package net.mullvad.mullvadvpn.compose.screen
 
 import android.content.Context
-import androidx.compose.animation.animateContentSize
 import androidx.compose.foundation.layout.Column
 import androidx.compose.foundation.layout.Spacer
 import androidx.compose.foundation.layout.fillMaxWidth
@@ -9,8 +8,10 @@ import androidx.compose.foundation.layout.height
 import androidx.compose.foundation.layout.padding
 import androidx.compose.foundation.lazy.LazyColumn
 import androidx.compose.foundation.lazy.LazyListScope
+import androidx.compose.foundation.lazy.itemsIndexed
 import androidx.compose.foundation.lazy.rememberLazyListState
 import androidx.compose.material3.ButtonDefaults
+import androidx.compose.material3.HorizontalDivider
 import androidx.compose.material3.MaterialTheme
 import androidx.compose.material3.SnackbarDuration
 import androidx.compose.material3.SnackbarHostState
@@ -54,7 +55,6 @@ import net.mullvad.mullvadvpn.compose.util.LaunchedEffectCollect
 import net.mullvad.mullvadvpn.compose.util.showSnackbarImmediately
 import net.mullvad.mullvadvpn.lib.model.CustomListId
 import net.mullvad.mullvadvpn.lib.model.RelayItem
-import net.mullvad.mullvadvpn.lib.theme.AppTheme
 import net.mullvad.mullvadvpn.lib.theme.Dimens
 import net.mullvad.mullvadvpn.lib.theme.color.AlphaScrollbar
 import net.mullvad.mullvadvpn.viewmodel.CustomListLocationsSideEffect
@@ -64,7 +64,7 @@ import org.koin.androidx.compose.koinViewModel
 @Composable
 @Preview
 private fun PreviewCustomListLocationScreen() {
-    AppTheme { CustomListLocationsScreen(state = CustomListLocationsUiState.Content.Data()) }
+    //    AppTheme { CustomListLocationsScreen(state = CustomListLocationsUiState.Content.Data()) }
 }
 
 data class CustomListLocationsNavArgs(
@@ -119,6 +119,7 @@ fun CustomListLocations(
         onSearchTermInput = customListsViewModel::onSearchTermInput,
         onSaveClick = customListsViewModel::save,
         onRelaySelectionClick = customListsViewModel::onRelaySelectionClick,
+        onExpand = customListsViewModel::onExpand,
         onBackClick =
             dropUnlessResumed {
                 if (state.hasUnsavedChanges) {
@@ -137,6 +138,7 @@ fun CustomListLocationsScreen(
     onSearchTermInput: (String) -> Unit = {},
     onSaveClick: () -> Unit = {},
     onRelaySelectionClick: (RelayItem.Location, selected: Boolean) -> Unit = { _, _ -> },
+    onExpand: (RelayItem.Location, selected: Boolean) -> Unit = { _, _ -> },
     onBackClick: () -> Unit = {}
 ) {
     ScaffoldWithSmallTopBar(
@@ -184,7 +186,11 @@ fun CustomListLocationsScreen(
                         empty(searchTerm = state.searchTerm)
                     }
                     is CustomListLocationsUiState.Content.Data -> {
-                        content(uiState = state, onRelaySelectedChanged = onRelaySelectionClick)
+                        content(
+                            uiState = state,
+                            onRelaySelectedChanged = onRelaySelectionClick,
+                            onExpand = onExpand
+                        )
                     }
                 }
             }
@@ -224,21 +230,27 @@ private fun LazyListScope.empty(searchTerm: String) {
 
 private fun LazyListScope.content(
     uiState: CustomListLocationsUiState.Content.Data,
+    onExpand: (RelayItem.Location, expand: Boolean) -> Unit,
     onRelaySelectedChanged: (RelayItem.Location, selected: Boolean) -> Unit,
 ) {
-    items(
-        count = uiState.availableLocations.size,
-        key = { index -> uiState.availableLocations[index].hashCode() },
-        contentType = { ContentType.ITEM },
-    ) { index ->
-        val country = uiState.availableLocations[index]
-        CheckableRelayLocationCell(
-            relay = country,
-            modifier = Modifier.animateContentSize(),
-            onRelayCheckedChange = { item, isChecked ->
-                onRelaySelectedChanged(item as RelayItem.Location, isChecked)
-            },
-            selectedRelays = uiState.selectedLocations,
-        )
+    itemsIndexed(
+        uiState.locations,
+        key = { index, listItem -> listItem.item.id },
+    ) { index, listItem ->
+        Column(modifier = Modifier.animateItem()) {
+            if (index != 0) {
+                HorizontalDivider()
+            }
+            CheckableRelayLocationCell(
+                item = listItem.item,
+                onRelayCheckedChange = { isChecked ->
+                    onRelaySelectedChanged(listItem.item, isChecked)
+                },
+                checked = listItem.checked,
+                depth = listItem.depth,
+                onExpand = { expand -> onExpand(listItem.item, expand) },
+                expanded = listItem.expanded,
+            )
+        }
     }
 }
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/state/CustomListLocationsUiState.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/state/CustomListLocationsUiState.kt
index f207d85359..c9c842ee0b 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/state/CustomListLocationsUiState.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/state/CustomListLocationsUiState.kt
@@ -22,11 +22,17 @@ sealed interface CustomListLocationsUiState {
 
         data class Data(
             override val newList: Boolean = false,
-            val availableLocations: List<RelayItem.Location.Country> = emptyList(),
-            val selectedLocations: Set<RelayItem> = emptySet(),
+            val locations: List<RelayLocationListItem>,
             override val searchTerm: String = "",
             override val saveEnabled: Boolean = false,
             override val hasUnsavedChanges: Boolean = false
         ) : Content
     }
 }
+
+data class RelayLocationListItem(
+    val item: RelayItem.Location,
+    val depth: Int = 0,
+    val checked: Boolean = false,
+    val expanded: Boolean = false
+)
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/viewmodel/CustomListLocationsViewModel.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/viewmodel/CustomListLocationsViewModel.kt
index 6d738a6417..3df2b2f623 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/viewmodel/CustomListLocationsViewModel.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/viewmodel/CustomListLocationsViewModel.kt
@@ -10,22 +10,28 @@ import kotlinx.coroutines.flow.SharedFlow
 import kotlinx.coroutines.flow.SharingStarted
 import kotlinx.coroutines.flow.combine
 import kotlinx.coroutines.flow.first
+import kotlinx.coroutines.flow.map
+import kotlinx.coroutines.flow.onEach
 import kotlinx.coroutines.flow.stateIn
 import kotlinx.coroutines.flow.update
 import kotlinx.coroutines.launch
 import net.mullvad.mullvadvpn.compose.communication.CustomListAction
 import net.mullvad.mullvadvpn.compose.communication.LocationsChanged
 import net.mullvad.mullvadvpn.compose.state.CustomListLocationsUiState
+import net.mullvad.mullvadvpn.compose.state.RelayLocationListItem
 import net.mullvad.mullvadvpn.lib.model.RelayItem
+import net.mullvad.mullvadvpn.lib.model.RelayItemId
+import net.mullvad.mullvadvpn.relaylist.MIN_SEARCH_LENGTH
+import net.mullvad.mullvadvpn.relaylist.ancestors
 import net.mullvad.mullvadvpn.relaylist.descendants
-import net.mullvad.mullvadvpn.relaylist.filterOnSearchTerm
+import net.mullvad.mullvadvpn.relaylist.newFilterOnSearch
 import net.mullvad.mullvadvpn.relaylist.withDescendants
 import net.mullvad.mullvadvpn.repository.RelayListRepository
 import net.mullvad.mullvadvpn.usecase.customlists.CustomListActionUseCase
 import net.mullvad.mullvadvpn.usecase.customlists.CustomListRelayItemsUseCase
 
 class CustomListLocationsViewModel(
-    relayListRepository: RelayListRepository,
+    private val relayListRepository: RelayListRepository,
     private val customListRelayItemsUseCase: CustomListRelayItemsUseCase,
     private val customListActionUseCase: CustomListActionUseCase,
     savedStateHandle: SavedStateHandle
@@ -39,18 +45,18 @@ class CustomListLocationsViewModel(
     private val _initialLocations = MutableStateFlow<Set<RelayItem.Location>>(emptySet())
     private val _selectedLocations = MutableStateFlow<Set<RelayItem.Location>?>(null)
     private val _searchTerm = MutableStateFlow(EMPTY_SEARCH_TERM)
+    private val _expandedItems = MutableStateFlow<Set<RelayItemId>>(setOf())
 
     val uiState =
-        combine(relayListRepository.relayList, _searchTerm, _selectedLocations) {
+        combine(searchRelayListLocations(), _searchTerm, _selectedLocations, _expandedItems) {
                 relayCountries,
                 searchTerm,
-                selectedLocations ->
-                val filteredRelayCountries = relayCountries.filterOnSearchTerm(searchTerm, null)
-
+                selectedLocations,
+                expandedLocations ->
                 when {
                     selectedLocations == null ->
                         CustomListLocationsUiState.Loading(newList = navArgs.newList)
-                    filteredRelayCountries.isEmpty() ->
+                    relayCountries.isEmpty() ->
                         CustomListLocationsUiState.Content.Empty(
                             newList = navArgs.newList,
                             searchTerm = searchTerm
@@ -59,8 +65,11 @@ class CustomListLocationsViewModel(
                         CustomListLocationsUiState.Content.Data(
                             newList = navArgs.newList,
                             searchTerm = searchTerm,
-                            availableLocations = filteredRelayCountries,
-                            selectedLocations = selectedLocations,
+                            locations =
+                                relayCountries.toRelayItems(
+                                    isSelected = { it in selectedLocations },
+                                    isExpanded = { it in expandedLocations },
+                                ),
                             saveEnabled =
                                 selectedLocations.isNotEmpty() &&
                                     selectedLocations != _initialLocations.value,
@@ -78,6 +87,23 @@ class CustomListLocationsViewModel(
         viewModelScope.launch { fetchInitialSelectedLocations() }
     }
 
+    fun searchRelayListLocations() =
+        combine(
+                _searchTerm,
+                relayListRepository.relayList,
+            ) { searchTerm, relayCountries ->
+                val isSearching = searchTerm.length >= MIN_SEARCH_LENGTH
+                if (isSearching) {
+                    val (exp, filteredRelayCountries) = relayCountries.newFilterOnSearch(searchTerm)
+                    exp.toSet() to filteredRelayCountries
+                } else {
+                    initialExpands((_selectedLocations.value ?: emptyList()).toSet()) to
+                        relayCountries
+                }
+            }
+            .onEach { _expandedItems.value = it.first }
+            .map { it.second }
+
     fun save() {
         viewModelScope.launch {
             _selectedLocations.value?.let { selectedLocations ->
@@ -113,6 +139,16 @@ class CustomListLocationsViewModel(
         }
     }
 
+    fun onExpand(relayItem: RelayItem.Location, expand: Boolean) {
+        _expandedItems.update {
+            if (expand) {
+                it + relayItem.id
+            } else {
+                it - relayItem.id
+            }
+        }
+    }
+
     fun onSearchTermInput(searchTerm: String) {
         viewModelScope.launch { _searchTerm.emit(searchTerm) }
     }
@@ -138,14 +174,10 @@ class CustomListLocationsViewModel(
         }
     }
 
-    private fun availableLocations(): List<RelayItem.Location.Country> =
-        (uiState.value as? CustomListLocationsUiState.Content.Data)?.availableLocations
-            ?: emptyList()
-
     private fun Set<RelayItem.Location>.deselectParents(
         relayItem: RelayItem.Location
     ): Set<RelayItem.Location> {
-        val availableLocations = availableLocations()
+        val availableLocations = relayListRepository.relayList.value
         val updateSelectionList = this.toMutableSet()
         when (relayItem) {
             is RelayItem.Location.City -> {
@@ -196,6 +228,52 @@ class CustomListLocationsViewModel(
 
         _initialLocations.value = selectedLocations
         _selectedLocations.value = selectedLocations
+        // Initial expand
+        _expandedItems.value = initialExpands(selectedLocations)
+    }
+
+    private fun initialExpands(locations: Set<RelayItem.Location>): Set<RelayItemId> =
+        locations.flatMap { it.id.ancestors() }.toSet()
+
+    private fun List<RelayItem.Location>.toRelayItems(
+        isSelected: (RelayItem) -> Boolean,
+        isExpanded: (RelayItemId) -> Boolean,
+        depth: Int = 0,
+    ): List<RelayLocationListItem> = flatMap { relayItem ->
+        buildList<RelayLocationListItem> {
+            val expanded = isExpanded(relayItem.id)
+            add(
+                RelayLocationListItem(
+                    item = relayItem,
+                    depth = depth,
+                    checked = isSelected(relayItem),
+                    expanded = expanded
+                )
+            )
+            if (expanded) {
+                when (relayItem) {
+                    is RelayItem.Location.City ->
+                        addAll(
+                            relayItem.relays.toRelayItems(
+                                isSelected = isSelected,
+                                isExpanded = isExpanded,
+                                depth = depth + 1
+                            )
+                        )
+                    is RelayItem.Location.Country ->
+                        addAll(
+                            relayItem.cities.toRelayItems(
+                                isSelected = isSelected,
+                                isExpanded = isExpanded,
+                                depth = depth + 1
+                            )
+                        )
+                    is RelayItem.Location.Relay -> {
+                        /* Do nothing */
+                    }
+                }
+            }
+        }
     }
 
     companion object {