Merge branch 'custom-list-added-animation-causes-custom-list-description-droid-901'

25 files changed, 1175 insertions, 892 deletions

diff --git a/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/data/DummyRelayItems.kt b/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/data/DummyRelayItems.kt
index 0218e06afd..052f2d897a 100644
--- a/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/data/DummyRelayItems.kt
+++ b/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/data/DummyRelayItems.kt
@@ -40,29 +40,25 @@ private val DUMMY_RELAY_2 =
 private val DUMMY_RELAY_CITY_1 =
     RelayItem.Location.City(
         name = "Relay City 1",
-        id = GeoLocationId.City(countryCode = GeoLocationId.Country("RCo1"), cityCode = "RCi1"),
+        id = GeoLocationId.City(country = GeoLocationId.Country("RCo1"), code = "RCi1"),
         relays = listOf(DUMMY_RELAY_1),
-        expanded = false
     )
 private val DUMMY_RELAY_CITY_2 =
     RelayItem.Location.City(
         name = "Relay City 2",
-        id = GeoLocationId.City(countryCode = GeoLocationId.Country("RCo2"), cityCode = "RCi2"),
+        id = GeoLocationId.City(country = GeoLocationId.Country("RCo2"), code = "RCi2"),
         relays = listOf(DUMMY_RELAY_2),
-        expanded = false
     )
 private val DUMMY_RELAY_COUNTRY_1 =
     RelayItem.Location.Country(
         name = "Relay Country 1",
         id = GeoLocationId.Country("RCo1"),
-        expanded = false,
         cities = listOf(DUMMY_RELAY_CITY_1)
     )
 private val DUMMY_RELAY_COUNTRY_2 =
     RelayItem.Location.Country(
         name = "Relay Country 2",
         id = GeoLocationId.Country("RCo2"),
-        expanded = false,
         cities = listOf(DUMMY_RELAY_CITY_2)
     )
 
@@ -80,15 +76,21 @@ val DUMMY_RELAY_LIST =
 val DUMMY_RELAY_ITEM_CUSTOM_LISTS =
     listOf(
         RelayItem.CustomList(
-            customListName = CustomListName.fromString("First list"),
-            expanded = false,
-            id = CustomListId("1"),
+            customList =
+                CustomList(
+                    name = CustomListName.fromString("First list"),
+                    id = CustomListId("1"),
+                    locations = emptyList()
+                ),
             locations = DUMMY_RELAY_COUNTRIES
         ),
         RelayItem.CustomList(
-            customListName = CustomListName.fromString("Empty list"),
-            expanded = false,
-            id = CustomListId("2"),
+            customList =
+                CustomList(
+                    name = CustomListName.fromString("Empty list"),
+                    id = CustomListId("2"),
+                    locations = emptyList()
+                ),
             locations = emptyList()
         )
     )
diff --git a/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/CustomListLocationsScreenTest.kt b/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/CustomListLocationsScreenTest.kt
index 4f4db0a529..1a8d35a5a9 100644
--- a/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/CustomListLocationsScreenTest.kt
+++ b/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/CustomListLocationsScreenTest.kt
@@ -12,6 +12,7 @@ import net.mullvad.mullvadvpn.compose.createEdgeToEdgeComposeExtension
 import net.mullvad.mullvadvpn.compose.data.DUMMY_RELAY_COUNTRIES
 import net.mullvad.mullvadvpn.compose.setContentWithTheme
 import net.mullvad.mullvadvpn.compose.state.CustomListLocationsUiState
+import net.mullvad.mullvadvpn.compose.state.RelayLocationListItem
 import net.mullvad.mullvadvpn.compose.test.CIRCULAR_PROGRESS_INDICATOR
 import net.mullvad.mullvadvpn.compose.test.SAVE_BUTTON_TEST_TAG
 import net.mullvad.mullvadvpn.lib.model.RelayItem
@@ -80,8 +81,14 @@ class CustomListLocationsScreenTest {
                 CustomListLocationsScreen(
                     state =
                         CustomListLocationsUiState.Content.Data(
-                            availableLocations = DUMMY_RELAY_COUNTRIES,
-                            selectedLocations = emptySet(),
+                            locations =
+                                listOf(
+                                    RelayLocationListItem(DUMMY_RELAY_COUNTRIES[0], checked = true),
+                                    RelayLocationListItem(
+                                        DUMMY_RELAY_COUNTRIES[1],
+                                        checked = false
+                                    ),
+                                ),
                             searchTerm = ""
                         ),
                 )
@@ -89,11 +96,7 @@ class CustomListLocationsScreenTest {
 
             // Assert
             onNodeWithText("Relay Country 1").assertExists()
-            onNodeWithText("Relay City 1").assertDoesNotExist()
-            onNodeWithText("Relay host 1").assertDoesNotExist()
             onNodeWithText("Relay Country 2").assertExists()
-            onNodeWithText("Relay City 2").assertDoesNotExist()
-            onNodeWithText("Relay host 2").assertDoesNotExist()
         }
 
     @Test
@@ -107,8 +110,8 @@ class CustomListLocationsScreenTest {
                     state =
                         CustomListLocationsUiState.Content.Data(
                             newList = false,
-                            availableLocations = DUMMY_RELAY_COUNTRIES,
-                            selectedLocations = setOf(selectedCountry)
+                            locations =
+                                listOf(RelayLocationListItem(selectedCountry, checked = true))
                         ),
                     onRelaySelectionClick = mockedOnRelaySelectionClicked
                 )
@@ -131,7 +134,7 @@ class CustomListLocationsScreenTest {
                     state =
                         CustomListLocationsUiState.Content.Data(
                             newList = false,
-                            availableLocations = DUMMY_RELAY_COUNTRIES,
+                            locations = emptyList(),
                         ),
                     onSearchTermInput = mockedSearchTermInput
                 )
@@ -197,7 +200,7 @@ class CustomListLocationsScreenTest {
                     state =
                         CustomListLocationsUiState.Content.Data(
                             newList = false,
-                            availableLocations = DUMMY_RELAY_COUNTRIES,
+                            locations = emptyList(),
                             saveEnabled = true,
                         ),
                     onSaveClick = mockOnSaveClick
@@ -221,7 +224,7 @@ class CustomListLocationsScreenTest {
                     state =
                         CustomListLocationsUiState.Content.Data(
                             newList = false,
-                            availableLocations = DUMMY_RELAY_COUNTRIES,
+                            locations = emptyList(),
                             saveEnabled = false,
                         ),
                     onSaveClick = mockOnSaveClick
diff --git a/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/SelectLocationScreenTest.kt b/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/SelectLocationScreenTest.kt
index 4fcee479d6..4f3bac57e2 100644
--- a/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/SelectLocationScreenTest.kt
+++ b/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/SelectLocationScreenTest.kt
@@ -12,6 +12,7 @@ import net.mullvad.mullvadvpn.compose.createEdgeToEdgeComposeExtension
 import net.mullvad.mullvadvpn.compose.data.DUMMY_RELAY_COUNTRIES
 import net.mullvad.mullvadvpn.compose.data.DUMMY_RELAY_ITEM_CUSTOM_LISTS
 import net.mullvad.mullvadvpn.compose.setContentWithTheme
+import net.mullvad.mullvadvpn.compose.state.RelayListItem
 import net.mullvad.mullvadvpn.compose.state.SelectLocationUiState
 import net.mullvad.mullvadvpn.compose.test.CIRCULAR_PROGRESS_INDICATOR
 import net.mullvad.mullvadvpn.compose.test.SELECT_LOCATION_CUSTOM_LIST_BOTTOM_SHEET_TEST_TAG
@@ -54,13 +55,13 @@ class SelectLocationScreenTest {
                 SelectLocationScreen(
                     state =
                         SelectLocationUiState.Content(
+                            searchTerm = "",
+                            filterChips = emptyList(),
+                            relayListItems =
+                                DUMMY_RELAY_COUNTRIES.map {
+                                    RelayListItem.GeoLocationItem(item = it)
+                                },
                             customLists = emptyList(),
-                            filteredCustomLists = emptyList(),
-                            countries = DUMMY_RELAY_COUNTRIES,
-                            selectedItem = null,
-                            selectedOwnership = null,
-                            selectedProvidersCount = 0,
-                            searchTerm = ""
                         ),
                 )
             }
@@ -75,45 +76,6 @@ class SelectLocationScreenTest {
         }
 
     @Test
-    fun testShowRelayListStateSelected() =
-        composeExtension.use {
-            val updatedDummyList =
-                DUMMY_RELAY_COUNTRIES.let {
-                    val cities = it[0].cities.toMutableList()
-                    val city = cities.removeAt(0)
-                    cities.add(0, city.copy(expanded = true))
-
-                    val mutableRelayList = it.toMutableList()
-                    mutableRelayList[0] = it[0].copy(expanded = true, cities = cities.toList())
-                    mutableRelayList
-                }
-
-            // Arrange
-            setContentWithTheme {
-                SelectLocationScreen(
-                    state =
-                        SelectLocationUiState.Content(
-                            customLists = emptyList(),
-                            filteredCustomLists = emptyList(),
-                            countries = updatedDummyList,
-                            selectedItem = updatedDummyList[0].cities[0].relays[0].id,
-                            selectedOwnership = null,
-                            selectedProvidersCount = 0,
-                            searchTerm = ""
-                        ),
-                )
-            }
-
-            // Assert
-            onNodeWithText("Relay Country 1").assertExists()
-            onNodeWithText("Relay City 1").assertExists()
-            onNodeWithText("Relay host 1").assertExists()
-            onNodeWithText("Relay Country 2").assertExists()
-            onNodeWithText("Relay City 2").assertDoesNotExist()
-            onNodeWithText("Relay host 2").assertDoesNotExist()
-        }
-
-    @Test
     fun testSearchInput() =
         composeExtension.use {
             // Arrange
@@ -122,13 +84,10 @@ class SelectLocationScreenTest {
                 SelectLocationScreen(
                     state =
                         SelectLocationUiState.Content(
-                            customLists = emptyList(),
-                            filteredCustomLists = emptyList(),
-                            countries = emptyList(),
-                            selectedItem = null,
-                            selectedOwnership = null,
-                            selectedProvidersCount = 0,
-                            searchTerm = ""
+                            searchTerm = "",
+                            filterChips = emptyList(),
+                            relayListItems = emptyList(),
+                            customLists = emptyList()
                         ),
                     onSearchTermInput = mockedSearchTermInput
                 )
@@ -152,13 +111,11 @@ class SelectLocationScreenTest {
                 SelectLocationScreen(
                     state =
                         SelectLocationUiState.Content(
+                            searchTerm = mockSearchString,
+                            filterChips = emptyList(),
+                            relayListItems =
+                                listOf(RelayListItem.LocationsEmptyText(mockSearchString)),
                             customLists = emptyList(),
-                            filteredCustomLists = emptyList(),
-                            countries = emptyList(),
-                            selectedItem = null,
-                            selectedOwnership = null,
-                            selectedProvidersCount = 0,
-                            searchTerm = mockSearchString
                         ),
                     onSearchTermInput = mockedSearchTermInput
                 )
@@ -170,7 +127,7 @@ class SelectLocationScreenTest {
         }
 
     @Test
-    fun givenNoCustomListsAndSearchIsTermIsEmptyShouldShowCustomListsEmptyText() =
+    fun customListFooterShouldShowEmptyTextWhenNoCustomList() =
         composeExtension.use {
             // Arrange
             val mockSearchString = ""
@@ -178,13 +135,10 @@ class SelectLocationScreenTest {
                 SelectLocationScreen(
                     state =
                         SelectLocationUiState.Content(
+                            searchTerm = mockSearchString,
+                            filterChips = emptyList(),
+                            relayListItems = listOf(RelayListItem.CustomListFooter(false)),
                             customLists = emptyList(),
-                            filteredCustomLists = emptyList(),
-                            countries = emptyList(),
-                            selectedItem = null,
-                            selectedOwnership = null,
-                            selectedProvidersCount = 0,
-                            searchTerm = mockSearchString
                         ),
                 )
             }
@@ -202,13 +156,10 @@ class SelectLocationScreenTest {
                 SelectLocationScreen(
                     state =
                         SelectLocationUiState.Content(
+                            searchTerm = mockSearchString,
+                            filterChips = emptyList(),
+                            relayListItems = emptyList(),
                             customLists = DUMMY_RELAY_ITEM_CUSTOM_LISTS,
-                            filteredCustomLists = emptyList(),
-                            countries = emptyList(),
-                            selectedItem = null,
-                            selectedOwnership = null,
-                            selectedProvidersCount = 0,
-                            searchTerm = mockSearchString
                         ),
                 )
             }
@@ -228,13 +179,10 @@ class SelectLocationScreenTest {
                 SelectLocationScreen(
                     state =
                         SelectLocationUiState.Content(
-                            customLists = DUMMY_RELAY_ITEM_CUSTOM_LISTS,
-                            filteredCustomLists = DUMMY_RELAY_ITEM_CUSTOM_LISTS,
-                            countries = emptyList(),
-                            selectedItem = null,
-                            selectedOwnership = null,
-                            selectedProvidersCount = 0,
-                            searchTerm = ""
+                            searchTerm = "",
+                            filterChips = emptyList(),
+                            relayListItems = listOf(RelayListItem.CustomListItem(customList)),
+                            customLists = DUMMY_RELAY_ITEM_CUSTOM_LISTS
                         ),
                     onSelectRelay = mockedOnSelectRelay
                 )
@@ -257,13 +205,11 @@ class SelectLocationScreenTest {
                 SelectLocationScreen(
                     state =
                         SelectLocationUiState.Content(
-                            customLists = DUMMY_RELAY_ITEM_CUSTOM_LISTS,
-                            filteredCustomLists = DUMMY_RELAY_ITEM_CUSTOM_LISTS,
-                            countries = emptyList(),
-                            selectedItem = null,
-                            selectedOwnership = null,
-                            selectedProvidersCount = 0,
-                            searchTerm = ""
+                            searchTerm = "",
+                            filterChips = emptyList(),
+                            relayListItems =
+                                listOf(RelayListItem.CustomListItem(item = customList)),
+                            customLists = DUMMY_RELAY_ITEM_CUSTOM_LISTS
                         ),
                     onSelectRelay = mockedOnSelectRelay
                 )
@@ -286,13 +232,10 @@ class SelectLocationScreenTest {
                 SelectLocationScreen(
                     state =
                         SelectLocationUiState.Content(
+                            searchTerm = "",
+                            filterChips = emptyList(),
+                            relayListItems = listOf(RelayListItem.GeoLocationItem(relayItem)),
                             customLists = emptyList(),
-                            filteredCustomLists = emptyList(),
-                            countries = DUMMY_RELAY_COUNTRIES,
-                            selectedItem = null,
-                            selectedOwnership = null,
-                            selectedProvidersCount = 0,
-                            searchTerm = ""
                         ),
                     onSelectRelay = mockedOnSelectRelay
                 )
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/cell/FilterCell.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/cell/FilterRow.kt
index 6dfd8f3eb1..39f6396132 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/cell/FilterCell.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/cell/FilterRow.kt
@@ -1,11 +1,10 @@
 package net.mullvad.mullvadvpn.compose.cell
 
 import androidx.compose.foundation.horizontalScroll
+import androidx.compose.foundation.layout.Arrangement
 import androidx.compose.foundation.layout.Row
-import androidx.compose.foundation.layout.Spacer
 import androidx.compose.foundation.layout.fillMaxWidth
 import androidx.compose.foundation.layout.padding
-import androidx.compose.foundation.layout.size
 import androidx.compose.foundation.rememberScrollState
 import androidx.compose.material3.MaterialTheme
 import androidx.compose.material3.Text
@@ -16,6 +15,7 @@ import androidx.compose.ui.res.stringResource
 import androidx.compose.ui.tooling.preview.Preview
 import net.mullvad.mullvadvpn.R
 import net.mullvad.mullvadvpn.compose.component.MullvadFilterChip
+import net.mullvad.mullvadvpn.compose.state.FilterChip
 import net.mullvad.mullvadvpn.lib.model.Ownership
 import net.mullvad.mullvadvpn.lib.theme.AppTheme
 import net.mullvad.mullvadvpn.lib.theme.Dimens
@@ -24,21 +24,19 @@ import net.mullvad.mullvadvpn.lib.theme.Dimens
 @Composable
 private fun PreviewFilterCell() {
     AppTheme {
-        FilterCell(
-            ownershipFilter = Ownership.MullvadOwned,
-            selectedProviderFilter = 3,
-            removeOwnershipFilter = {},
-            removeProviderFilter = {}
+        FilterRow(
+            listOf(FilterChip.Ownership(Ownership.MullvadOwned), FilterChip.Provider(2)),
+            {},
+            {}
         )
     }
 }
 
 @Composable
-fun FilterCell(
-    ownershipFilter: Ownership?,
-    selectedProviderFilter: Int?,
-    removeOwnershipFilter: () -> Unit,
-    removeProviderFilter: () -> Unit
+fun FilterRow(
+    filters: List<FilterChip>,
+    onRemoveOwnershipFilter: () -> Unit,
+    onRemoveProviderFilter: () -> Unit
 ) {
     val scrollState = rememberScrollState()
     Row(
@@ -49,31 +47,39 @@ fun FilterCell(
                     horizontal = Dimens.searchFieldHorizontalPadding,
                 )
                 .fillMaxWidth(),
+        horizontalArrangement = Arrangement.spacedBy(Dimens.chipSpace)
     ) {
         Text(
-            modifier = Modifier.padding(end = Dimens.filterTittlePadding),
             text = stringResource(id = R.string.filtered),
             color = MaterialTheme.colorScheme.onPrimary,
             style = MaterialTheme.typography.labelMedium
         )
-
-        if (selectedProviderFilter != null) {
-            MullvadFilterChip(
-                text = stringResource(id = R.string.number_of_providers, selectedProviderFilter),
-                onRemoveClick = removeProviderFilter
-            )
-            Spacer(modifier = Modifier.size(Dimens.chipSpace))
-        }
-
-        if (ownershipFilter != null) {
-            MullvadFilterChip(
-                text = stringResource(ownershipFilter.stringResources()),
-                onRemoveClick = removeOwnershipFilter
-            )
+        filters.forEach {
+            when (it) {
+                is FilterChip.Ownership ->
+                    OwnershipFilterChip(it.ownership, onRemoveOwnershipFilter)
+                is FilterChip.Provider -> ProviderFilterChip(it.count, onRemoveProviderFilter)
+            }
         }
     }
 }
 
+@Composable
+fun ProviderFilterChip(providers: Int, onRemoveClick: () -> Unit) {
+    MullvadFilterChip(
+        text = stringResource(id = R.string.number_of_providers, providers),
+        onRemoveClick = onRemoveClick
+    )
+}
+
+@Composable
+fun OwnershipFilterChip(ownership: Ownership, onRemoveClick: () -> Unit) {
+    MullvadFilterChip(
+        text = stringResource(ownership.stringResources()),
+        onRemoveClick = onRemoveClick
+    )
+}
+
 private fun Ownership.stringResources(): Int =
     when (this) {
         Ownership.MullvadOwned -> R.string.owned
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/cell/RelayLocationCell.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/cell/RelayLocationCell.kt
index 9f1fef4e08..cf9682bcbe 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/cell/RelayLocationCell.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/cell/RelayLocationCell.kt
@@ -1,12 +1,10 @@
 package net.mullvad.mullvadvpn.compose.cell
 
-import androidx.compose.animation.animateContentSize
 import androidx.compose.foundation.ExperimentalFoundationApi
 import androidx.compose.foundation.background
 import androidx.compose.foundation.clickable
 import androidx.compose.foundation.combinedClickable
 import androidx.compose.foundation.layout.Box
-import androidx.compose.foundation.layout.BoxScope
 import androidx.compose.foundation.layout.Column
 import androidx.compose.foundation.layout.IntrinsicSize
 import androidx.compose.foundation.layout.Row
@@ -16,15 +14,12 @@ import androidx.compose.foundation.layout.fillMaxWidth
 import androidx.compose.foundation.layout.height
 import androidx.compose.foundation.layout.padding
 import androidx.compose.foundation.layout.size
-import androidx.compose.foundation.layout.wrapContentHeight
 import androidx.compose.foundation.shape.CircleShape
 import androidx.compose.material3.Icon
 import androidx.compose.material3.MaterialTheme
 import androidx.compose.material3.Text
 import androidx.compose.material3.VerticalDivider
 import androidx.compose.runtime.Composable
-import androidx.compose.runtime.mutableStateOf
-import androidx.compose.runtime.saveable.rememberSaveable
 import androidx.compose.ui.Alignment
 import androidx.compose.ui.Modifier
 import androidx.compose.ui.draw.alpha
@@ -33,34 +28,17 @@ import androidx.compose.ui.res.painterResource
 import androidx.compose.ui.text.style.TextOverflow
 import androidx.compose.ui.tooling.preview.Preview
 import androidx.compose.ui.tooling.preview.PreviewParameter
-import androidx.compose.ui.unit.Dp
+import androidx.compose.ui.unit.dp
 import net.mullvad.mullvadvpn.R
 import net.mullvad.mullvadvpn.compose.component.Chevron
 import net.mullvad.mullvadvpn.compose.component.MullvadCheckbox
 import net.mullvad.mullvadvpn.compose.preview.RelayItemCheckableCellPreviewParameterProvider
-import net.mullvad.mullvadvpn.compose.preview.RelayItemStatusCellPreviewParameterProvider
 import net.mullvad.mullvadvpn.lib.model.RelayItem
-import net.mullvad.mullvadvpn.lib.model.RelayItemId
 import net.mullvad.mullvadvpn.lib.theme.AppTheme
 import net.mullvad.mullvadvpn.lib.theme.Dimens
 import net.mullvad.mullvadvpn.lib.theme.color.AlphaInactive
-import net.mullvad.mullvadvpn.lib.theme.color.AlphaInvisible
 import net.mullvad.mullvadvpn.lib.theme.color.AlphaVisible
 import net.mullvad.mullvadvpn.lib.theme.color.selected
-import net.mullvad.mullvadvpn.relaylist.children
-
-@Composable
-@Preview
-private fun PreviewStatusRelayLocationCell(
-    @PreviewParameter(RelayItemStatusCellPreviewParameterProvider::class)
-    relayItems: List<RelayItem.Location.Country>
-) {
-    AppTheme {
-        Column(Modifier.background(color = MaterialTheme.colorScheme.background)) {
-            relayItems.map { StatusRelayLocationCell(relay = it) }
-        }
-    }
-}
 
 @Composable
 @Preview
@@ -70,169 +48,149 @@ private fun PreviewCheckableRelayLocationCell(
 ) {
     AppTheme {
         Column(Modifier.background(color = MaterialTheme.colorScheme.background)) {
-            relayItems.map { CheckableRelayLocationCell(relay = it) }
+            relayItems.map {
+                CheckableRelayLocationCell(
+                    item = it,
+                    checked = false,
+                    expanded = false,
+                    depth = 0,
+                    onExpand = {}
+                )
+            }
         }
     }
 }
 
+@OptIn(ExperimentalFoundationApi::class)
 @Composable
-fun StatusRelayLocationCell(
-    relay: RelayItem,
+fun StatusRelayItemCell(
+    item: RelayItem,
+    isSelected: Boolean,
     modifier: Modifier = Modifier,
+    onClick: () -> Unit = {},
+    onLongClick: (() -> Unit)? = null,
+    onToggleExpand: ((Boolean) -> Unit) = {},
+    isExpanded: Boolean = false,
+    depth: Int = 0,
     activeColor: Color = MaterialTheme.colorScheme.selected,
     inactiveColor: Color = MaterialTheme.colorScheme.error,
     disabledColor: Color = MaterialTheme.colorScheme.onSecondary,
-    selectedItem: RelayItemId? = null,
-    onSelectRelay: (item: RelayItem) -> Unit = {},
-    onLongClick: (item: RelayItem) -> Unit = {},
 ) {
-    RelayLocationCell(
-        relay = relay,
-        leadingContent = { relayItem ->
-            val selected = selectedItem == relayItem.id
-            Box(
-                modifier =
-                    Modifier.align(Alignment.CenterStart)
-                        .size(Dimens.relayCircleSize)
-                        .background(
-                            color =
-                                when {
-                                    selected -> Color.Unspecified
-                                    relayItem is RelayItem.CustomList && !relayItem.hasChildren ->
-                                        disabledColor
-                                    relayItem.active -> activeColor
-                                    else -> inactiveColor
-                                },
-                            shape = CircleShape
-                        )
-            )
-            Icon(
-                painter = painterResource(id = R.drawable.icon_tick),
-                modifier =
-                    Modifier.align(Alignment.CenterStart)
-                        .alpha(
-                            if (selected) {
-                                AlphaVisible
-                            } else {
-                                AlphaInvisible
-                            }
-                        ),
-                tint = Color.Unspecified,
-                contentDescription = null
-            )
-        },
+
+    RelayItemCell(
         modifier = modifier,
-        specialBackgroundColor = { relayItem ->
-            when {
-                selectedItem == relayItem.id -> MaterialTheme.colorScheme.selected
-                relayItem is RelayItem.CustomList && !relayItem.active ->
-                    MaterialTheme.colorScheme.surfaceTint
-                else -> null
+        item,
+        isSelected,
+        leadingContent = {
+            if (isSelected) {
+                Icon(
+                    painter = painterResource(id = R.drawable.icon_tick),
+                    contentDescription = null
+                )
+            } else {
+                Box(
+                    modifier =
+                        Modifier.padding(4.dp)
+                            .size(Dimens.relayCircleSize)
+                            .background(
+                                color =
+                                    when {
+                                        isSelected -> Color.Unspecified
+                                        item is RelayItem.CustomList && item.locations.isEmpty() ->
+                                            disabledColor
+                                        item.active -> activeColor
+                                        else -> inactiveColor
+                                    },
+                                shape = CircleShape
+                            )
+                )
             }
         },
-        onClick = onSelectRelay,
+        onClick = onClick,
         onLongClick = onLongClick,
-        depth = 0
-    )
-}
-
-@Composable
-fun CheckableRelayLocationCell(
-    relay: RelayItem,
-    modifier: Modifier = Modifier,
-    onRelayCheckedChange: (item: RelayItem, isChecked: Boolean) -> Unit = { _, _ -> },
-    selectedRelays: Set<RelayItem> = emptySet(),
-) {
-    RelayLocationCell(
-        relay = relay,
-        leadingContent = { relayItem ->
-            val checked = selectedRelays.contains(relayItem)
-            MullvadCheckbox(
-                checked = checked,
-                onCheckedChange = { isChecked -> onRelayCheckedChange(relayItem, isChecked) }
-            )
-        },
-        leadingContentStartPadding = Dimens.cellStartPaddingInteractive,
-        modifier = modifier,
-        onClick = { onRelayCheckedChange(it, !selectedRelays.contains(it)) },
-        onLongClick = null,
-        depth = 0
+        onToggleExpand = onToggleExpand,
+        isExpanded = isExpanded,
+        depth = depth,
     )
 }
 
 @OptIn(ExperimentalFoundationApi::class)
 @Composable
-private fun RelayLocationCell(
-    relay: RelayItem,
-    leadingContent: @Composable BoxScope.(relay: RelayItem) -> Unit,
+fun RelayItemCell(
     modifier: Modifier = Modifier,
-    leadingContentStartPadding: Dp = Dimens.cellStartPadding,
-    leadingContentStarPaddingModifier: Dp = Dimens.mediumPadding,
-    specialBackgroundColor: @Composable (relayItem: RelayItem) -> Color? = { null },
-    onClick: (item: RelayItem) -> Unit,
-    onLongClick: ((item: RelayItem) -> Unit)?,
-    depth: Int
+    item: RelayItem,
+    isSelected: Boolean,
+    leadingContent: (@Composable RowScope.() -> Unit)? = null,
+    onClick: () -> Unit,
+    onLongClick: (() -> Unit)? = null,
+    onToggleExpand: ((Boolean) -> Unit),
+    isExpanded: Boolean,
+    depth: Int,
 ) {
+
+    val leadingContentStartPadding = Dimens.cellStartPadding
+    val leadingContentStarPaddingModifier = Dimens.mediumPadding
     val startPadding = leadingContentStartPadding + leadingContentStarPaddingModifier * depth
-    val expanded =
-        rememberSaveable(key = relay.expanded.toString()) { mutableStateOf(relay.expanded) }
-    Column(
+    Row(
         modifier =
             modifier
                 .fillMaxWidth()
-                .padding(top = Dimens.listItemDivider)
-                .wrapContentHeight()
-                .fillMaxWidth()
-    ) {
-        Row(
-            modifier =
-                Modifier.align(Alignment.Start)
-                    .wrapContentHeight()
-                    .height(IntrinsicSize.Min)
-                    .fillMaxWidth()
-                    .background(specialBackgroundColor.invoke(relay) ?: depth.toBackgroundColor())
-        ) {
-            Row(
-                modifier =
-                    Modifier.weight(1f)
-                        .combinedClickable(
-                            enabled = relay.active,
-                            onClick = { onClick(relay) },
-                            onLongClick = { onLongClick?.invoke(relay) },
-                        )
-            ) {
-                Box(
-                    modifier =
-                        Modifier.align(Alignment.CenterVertically).padding(start = startPadding)
-                ) {
-                    leadingContent(relay)
-                }
-                Name(
-                    modifier = Modifier.weight(1f).align(Alignment.CenterVertically),
-                    relay = relay
+                .height(IntrinsicSize.Min)
+                .background(
+                    when {
+                        isSelected -> MaterialTheme.colorScheme.selected
+                        item is RelayItem.CustomList && !item.active ->
+                            MaterialTheme.colorScheme.surfaceTint
+                        else -> depth.toBackgroundColor()
+                    }
                 )
-            }
-            if (relay.hasChildren) {
-                ExpandButton(isExpanded = expanded.value) { expand -> expanded.value = expand }
-            }
-        }
-        if (expanded.value) {
-            relay.children().forEach {
-                RelayLocationCell(
-                    relay = it,
+                .combinedClickable(
+                    enabled = item.active,
                     onClick = onClick,
-                    modifier = Modifier.animateContentSize(),
-                    leadingContent = leadingContent,
-                    specialBackgroundColor = specialBackgroundColor,
                     onLongClick = onLongClick,
-                    depth = depth + 1,
                 )
-            }
+                .padding(start = startPadding),
+        verticalAlignment = Alignment.CenterVertically
+    ) {
+        if (leadingContent != null) {
+            leadingContent()
+        }
+        Name(modifier = Modifier.weight(1f), relay = item)
+
+        if (item.hasChildren) {
+            ExpandButton(isExpanded = isExpanded, onClick = { onToggleExpand(!isExpanded) })
         }
     }
 }
 
 @Composable
+fun CheckableRelayLocationCell(
+    item: RelayItem,
+    modifier: Modifier = Modifier,
+    checked: Boolean,
+    onRelayCheckedChange: (isChecked: Boolean) -> Unit = { _ -> },
+    expanded: Boolean,
+    onExpand: (Boolean) -> Unit,
+    depth: Int
+) {
+    RelayItemCell(
+        modifier = modifier,
+        item = item,
+        isSelected = false,
+        leadingContent = {
+            MullvadCheckbox(
+                checked = checked,
+                onCheckedChange = { isChecked -> onRelayCheckedChange(isChecked) }
+            )
+        },
+        onClick = { onRelayCheckedChange(!checked) },
+        onToggleExpand = onExpand,
+        isExpanded = expanded,
+        depth = depth
+    )
+}
+
+@Composable
 private fun Name(modifier: Modifier = Modifier, relay: RelayItem) {
     Text(
         text = relay.name,
@@ -275,6 +233,5 @@ private fun Int.toBackgroundColor(): Color =
         0 -> MaterialTheme.colorScheme.surfaceContainerHighest
         1 -> MaterialTheme.colorScheme.surfaceContainerHigh
         2 -> MaterialTheme.colorScheme.surfaceContainerLow
-        3 -> MaterialTheme.colorScheme.surfaceContainerLowest
         else -> MaterialTheme.colorScheme.surfaceContainerLowest
     }
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/preview/RelayItemCheckableCellPreviewParameterProvider.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/preview/RelayItemCheckableCellPreviewParameterProvider.kt
index c0cae0128f..bdf1ace173 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/preview/RelayItemCheckableCellPreviewParameterProvider.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/preview/RelayItemCheckableCellPreviewParameterProvider.kt
@@ -18,14 +18,11 @@ class RelayItemCheckableCellPreviewParameterProvider :
                     name = "Relay country Expanded",
                     cityNames = listOf("Normal city"),
                     relaysPerCity = 2,
-                    expanded = true
                 ),
                 generateRelayItemCountry(
                     name = "Country and city Expanded",
                     cityNames = listOf("Expanded city A", "Expanded city B"),
                     relaysPerCity = 2,
-                    expanded = true,
-                    expandChildren = true
                 )
             )
         )
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/preview/RelayItemPreviewData.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/preview/RelayItemPreviewData.kt
index afaf81ac55..c1b42c9415 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/preview/RelayItemPreviewData.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/preview/RelayItemPreviewData.kt
@@ -12,8 +12,6 @@ internal object RelayItemPreviewData {
         cityNames: List<String>,
         relaysPerCity: Int,
         active: Boolean = true,
-        expanded: Boolean = false,
-        expandChildren: Boolean = false,
     ) =
         RelayItem.Location.Country(
             name = name,
@@ -25,10 +23,8 @@ internal object RelayItemPreviewData {
                         name.generateCountryCode(),
                         relaysPerCity,
                         active,
-                        expandChildren
                     )
                 },
-            expanded = expanded,
         )
 }
 
@@ -37,7 +33,6 @@ private fun generateRelayItemCity(
     countryCode: GeoLocationId.Country,
     numberOfRelays: Int,
     active: Boolean = true,
-    expanded: Boolean = false,
 ) =
     RelayItem.Location.City(
         name = name,
@@ -50,7 +45,6 @@ private fun generateRelayItemCity(
                     active
                 )
             },
-        expanded = expanded,
     )
 
 private fun generateRelayItemRelay(
@@ -62,7 +56,7 @@ private fun generateRelayItemRelay(
         id =
             GeoLocationId.Hostname(
                 city = cityCode,
-                hostname = hostName,
+                code = hostName,
             ),
         active = active,
         provider = Provider(ProviderId("Provider"), Ownership.MullvadOwned),
@@ -75,6 +69,6 @@ private fun String.generateCityCode(countryCode: GeoLocationId.Country) =
     GeoLocationId.City(countryCode, take(CITY_CODE_LENGTH).lowercase())
 
 private fun generateHostname(city: GeoLocationId.City, index: Int) =
-    "${city.countryCode.countryCode}-${city.cityCode}-wg-${index+1}"
+    "${city.country.code}-${city.code}-wg-${index+1}"
 
 private const val CITY_CODE_LENGTH = 3
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/preview/RelayItemStatusCellPreviewParameterProvider.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/preview/RelayItemStatusCellPreviewParameterProvider.kt
index 26ea644185..a825975b0f 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/preview/RelayItemStatusCellPreviewParameterProvider.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/preview/RelayItemStatusCellPreviewParameterProvider.kt
@@ -24,14 +24,11 @@ class RelayItemStatusCellPreviewParameterProvider :
                     name = "Relay country Expanded",
                     cityNames = listOf("Normal city"),
                     relaysPerCity = 2,
-                    expanded = true
                 ),
                 generateRelayItemCountry(
                     name = "Country and city Expanded",
                     cityNames = listOf("Expanded city A", "Expanded city B"),
                     relaysPerCity = 2,
-                    expanded = true,
-                    expandChildren = true
                 )
             )
         )
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/screen/EditCustomListScreen.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/screen/EditCustomListScreen.kt
index c3f7662ea2..365328ea93 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/screen/EditCustomListScreen.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/screen/EditCustomListScreen.kt
@@ -65,10 +65,7 @@ private fun PreviewEditCustomListScreen() {
                     locations =
                         listOf(
                             GeoLocationId.Hostname(
-                                GeoLocationId.City(
-                                    GeoLocationId.Country("country"),
-                                    cityCode = "city"
-                                ),
+                                GeoLocationId.City(GeoLocationId.Country("country"), code = "city"),
                                 "hostname",
                             )
                         )
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/screen/SelectLocationScreen.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/screen/SelectLocationScreen.kt
index 555983d51d..430c1130d5 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/screen/SelectLocationScreen.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/screen/SelectLocationScreen.kt
@@ -1,7 +1,6 @@
 package net.mullvad.mullvadvpn.compose.screen
 
 import android.content.Context
-import androidx.compose.animation.animateContentSize
 import androidx.compose.foundation.ExperimentalFoundationApi
 import androidx.compose.foundation.background
 import androidx.compose.foundation.gestures.animateScrollBy
@@ -13,9 +12,10 @@ import androidx.compose.foundation.layout.fillMaxWidth
 import androidx.compose.foundation.layout.height
 import androidx.compose.foundation.layout.padding
 import androidx.compose.foundation.lazy.LazyColumn
+import androidx.compose.foundation.lazy.LazyItemScope
 import androidx.compose.foundation.lazy.LazyListScope
 import androidx.compose.foundation.lazy.LazyListState
-import androidx.compose.foundation.lazy.items
+import androidx.compose.foundation.lazy.itemsIndexed
 import androidx.compose.foundation.lazy.rememberLazyListState
 import androidx.compose.material3.ExperimentalMaterial3Api
 import androidx.compose.material3.HorizontalDivider
@@ -45,7 +45,9 @@ import androidx.compose.ui.res.painterResource
 import androidx.compose.ui.res.stringResource
 import androidx.compose.ui.text.style.TextAlign
 import androidx.compose.ui.tooling.preview.Preview
+import androidx.lifecycle.compose.LocalLifecycleOwner
 import androidx.lifecycle.compose.collectAsStateWithLifecycle
+import androidx.lifecycle.compose.currentStateAsState
 import androidx.lifecycle.compose.dropUnlessResumed
 import com.ramcosta.composedestinations.annotation.Destination
 import com.ramcosta.composedestinations.annotation.RootGraph
@@ -62,10 +64,10 @@ import com.ramcosta.composedestinations.result.ResultRecipient
 import com.ramcosta.composedestinations.spec.DestinationSpec
 import kotlinx.coroutines.launch
 import net.mullvad.mullvadvpn.R
-import net.mullvad.mullvadvpn.compose.cell.FilterCell
+import net.mullvad.mullvadvpn.compose.cell.FilterRow
 import net.mullvad.mullvadvpn.compose.cell.HeaderCell
 import net.mullvad.mullvadvpn.compose.cell.IconCell
-import net.mullvad.mullvadvpn.compose.cell.StatusRelayLocationCell
+import net.mullvad.mullvadvpn.compose.cell.StatusRelayItemCell
 import net.mullvad.mullvadvpn.compose.cell.SwitchComposeSubtitleCell
 import net.mullvad.mullvadvpn.compose.cell.ThreeDotCell
 import net.mullvad.mullvadvpn.compose.communication.Created
@@ -81,6 +83,11 @@ import net.mullvad.mullvadvpn.compose.component.MullvadSnackbar
 import net.mullvad.mullvadvpn.compose.component.drawVerticalScrollbar
 import net.mullvad.mullvadvpn.compose.constant.ContentType
 import net.mullvad.mullvadvpn.compose.extensions.dropUnlessResumed
+import net.mullvad.mullvadvpn.compose.screen.BottomSheetState.ShowCustomListsBottomSheet
+import net.mullvad.mullvadvpn.compose.screen.BottomSheetState.ShowCustomListsEntryBottomSheet
+import net.mullvad.mullvadvpn.compose.screen.BottomSheetState.ShowEditCustomListBottomSheet
+import net.mullvad.mullvadvpn.compose.screen.BottomSheetState.ShowLocationBottomSheet
+import net.mullvad.mullvadvpn.compose.state.RelayListItem
 import net.mullvad.mullvadvpn.compose.state.SelectLocationUiState
 import net.mullvad.mullvadvpn.compose.test.CIRCULAR_PROGRESS_INDICATOR
 import net.mullvad.mullvadvpn.compose.test.SELECT_LOCATION_CUSTOM_LIST_BOTTOM_SHEET_TEST_TAG
@@ -92,7 +99,6 @@ import net.mullvad.mullvadvpn.compose.util.CollectSideEffectWithLifecycle
 import net.mullvad.mullvadvpn.compose.util.RunOnKeyChange
 import net.mullvad.mullvadvpn.compose.util.showSnackbarImmediately
 import net.mullvad.mullvadvpn.lib.model.CustomListId
-import net.mullvad.mullvadvpn.lib.model.GeoLocationId
 import net.mullvad.mullvadvpn.lib.model.RelayItem
 import net.mullvad.mullvadvpn.lib.model.RelayItemId
 import net.mullvad.mullvadvpn.lib.theme.AppTheme
@@ -111,20 +117,9 @@ private fun PreviewSelectLocationScreen() {
     val state =
         SelectLocationUiState.Content(
             searchTerm = "",
-            selectedOwnership = null,
-            selectedProvidersCount = 0,
-            countries =
-                listOf(
-                    RelayItem.Location.Country(
-                        GeoLocationId.Country("Country 1"),
-                        "Code 1",
-                        false,
-                        emptyList()
-                    )
-                ),
-            selectedItem = null,
+            emptyList(),
+            relayListItems = emptyList(),
             customLists = emptyList(),
-            filteredCustomLists = emptyList()
         )
     AppTheme {
         SelectLocationScreen(
@@ -147,14 +142,16 @@ fun SelectLocation(
         ResultRecipient<CustomListLocationsDestination, LocationsChanged>
 ) {
     val vm = koinViewModel<SelectLocationViewModel>()
-    val state = vm.uiState.collectAsStateWithLifecycle().value
+    val state = vm.uiState.collectAsStateWithLifecycle()
 
     val snackbarHostState = remember { SnackbarHostState() }
     val context = LocalContext.current
-
+    val lazyListState = rememberLazyListState()
     CollectSideEffectWithLifecycle(vm.uiSideEffect) {
         when (it) {
-            SelectLocationSideEffect.CloseScreen -> backNavigator.navigateBack(result = true)
+            SelectLocationSideEffect.CloseScreen -> {
+                backNavigator.navigateBack(result = true)
+            }
             is SelectLocationSideEffect.LocationAddedToCustomList ->
                 launch {
                     snackbarHostState.showResultSnackbar(
@@ -181,6 +178,15 @@ fun SelectLocation(
         }
     }
 
+    val stateActual = state.value
+    RunOnKeyChange(stateActual is SelectLocationUiState.Content) {
+        val index = stateActual.indexOfSelectedRelayItem()
+        if (index != -1) {
+            lazyListState.scrollToItem(index)
+            lazyListState.animateScrollAndCentralizeItem(index)
+        }
+    }
+
     createCustomListDialogResultRecipient.OnCustomListNavResult(
         snackbarHostState,
         vm::performAction
@@ -199,7 +205,8 @@ fun SelectLocation(
     updateCustomListResultRecipient.OnCustomListNavResult(snackbarHostState, vm::performAction)
 
     SelectLocationScreen(
-        state = state,
+        state = state.value,
+        lazyListState = lazyListState,
         snackbarHostState = snackbarHostState,
         onSelectRelay = vm::selectRelay,
         onSearchTermInput = vm::onSearchTermInput,
@@ -211,6 +218,7 @@ fun SelectLocation(
                     CreateCustomListDestination(locationCode = relayItem?.id),
                 )
             },
+        onToggleExpand = vm::onToggleExpand,
         onEditCustomLists = dropUnlessResumed { navigator.navigate(CustomListsDestination()) },
         removeOwnershipFilter = vm::removeOwnerFilter,
         removeProviderFilter = vm::removeProviderFilter,
@@ -221,7 +229,7 @@ fun SelectLocation(
                 navigator.navigate(
                     EditCustomListNameDestination(
                         customListId = customList.id,
-                        initialName = customList.customListName
+                        initialName = customList.customList.name
                     ),
                 )
             },
@@ -236,7 +244,7 @@ fun SelectLocation(
                 navigator.navigate(
                     DeleteCustomListDestination(
                         customListId = customList.id,
-                        name = customList.customListName
+                        name = customList.customList.name
                     ),
                 )
             }
@@ -248,6 +256,7 @@ fun SelectLocation(
 @Composable
 fun SelectLocationScreen(
     state: SelectLocationUiState,
+    lazyListState: LazyListState = rememberLazyListState(),
     snackbarHostState: SnackbarHostState = remember { SnackbarHostState() },
     onSelectRelay: (item: RelayItem) -> Unit = {},
     onSearchTermInput: (searchTerm: String) -> Unit = {},
@@ -260,13 +269,13 @@ fun SelectLocationScreen(
     onAddLocationToList: (location: RelayItem.Location, customList: RelayItem.CustomList) -> Unit =
         { _, _ ->
         },
-    onRemoveLocationFromList:
-        (location: RelayItem.Location, customList: RelayItem.CustomList) -> Unit =
+    onRemoveLocationFromList: (location: RelayItem.Location, customListId: CustomListId) -> Unit =
         { _, _ ->
         },
     onEditCustomListName: (RelayItem.CustomList) -> Unit = {},
     onEditLocationsCustomList: (RelayItem.CustomList) -> Unit = {},
-    onDeleteCustomList: (RelayItem.CustomList) -> Unit = {}
+    onDeleteCustomList: (RelayItem.CustomList) -> Unit = {},
+    onToggleExpand: (RelayItemId, CustomListId?, Boolean) -> Unit = { _, _, _ -> },
 ) {
     val backgroundColor = MaterialTheme.colorScheme.background
 
@@ -278,6 +287,8 @@ fun SelectLocationScreen(
             )
         }
     ) {
+        val lifecycleState = LocalLifecycleOwner.current.lifecycle.currentStateAsState()
+        Text(text = lifecycleState.value.toString())
         var bottomSheetState by remember { mutableStateOf<BottomSheetState?>(null) }
         BottomSheets(
             bottomSheetState = bottomSheetState,
@@ -294,18 +305,8 @@ fun SelectLocationScreen(
         Column(modifier = Modifier.padding(it).background(backgroundColor).fillMaxSize()) {
             SelectLocationTopBar(onBackClick = onBackClick, onFilterClick = onFilterClick)
 
-            when (state) {
-                SelectLocationUiState.Loading -> {}
-                is SelectLocationUiState.Content -> {
-                    if (state.hasFilter) {
-                        FilterCell(
-                            ownershipFilter = state.selectedOwnership,
-                            selectedProviderFilter = state.selectedProvidersCount,
-                            removeOwnershipFilter = removeOwnershipFilter,
-                            removeProviderFilter = removeProviderFilter,
-                        )
-                    }
-                }
+            if (state is SelectLocationUiState.Content && state.filterChips.isNotEmpty()) {
+                FilterRow(filters = state.filterChips, removeOwnershipFilter, removeProviderFilter)
             }
 
             SearchTextField(
@@ -319,16 +320,7 @@ fun SelectLocationScreen(
                 onSearchTermInput.invoke(searchString)
             }
             Spacer(modifier = Modifier.height(height = Dimens.verticalSpace))
-            val lazyListState = rememberLazyListState()
-            val selectedItemCode = (state as? SelectLocationUiState.Content)?.selectedItem ?: ""
-            RunOnKeyChange(key = selectedItemCode) {
-                val index = state.indexOfSelectedRelayItem()
 
-                if (index >= 0) {
-                    lazyListState.scrollToItem(index)
-                    lazyListState.animateScrollAndCentralizeItem(index)
-                }
-            }
             LazyColumn(
                 modifier =
                     Modifier.fillMaxSize()
@@ -344,58 +336,87 @@ fun SelectLocationScreen(
                         loading()
                     }
                     is SelectLocationUiState.Content -> {
-                        if (state.showCustomLists) {
-                            customLists(
-                                customLists = state.filteredCustomLists,
-                                selectedItem = state.selectedItem,
-                                backgroundColor = backgroundColor,
-                                onSelectRelay = onSelectRelay,
-                                onShowCustomListBottomSheet = {
-                                    bottomSheetState =
-                                        BottomSheetState.ShowCustomListsBottomSheet(
-                                            state.customLists.isNotEmpty()
-                                        )
-                                },
-                                onShowEditBottomSheet = { customList ->
-                                    bottomSheetState =
-                                        BottomSheetState.ShowEditCustomListBottomSheet(customList)
-                                },
-                                onShowEditCustomListEntryBottomSheet = {
-                                    item: RelayItem.Location,
-                                    customList: RelayItem.CustomList ->
-                                    bottomSheetState =
-                                        BottomSheetState.ShowCustomListsEntryBottomSheet(
-                                            customList,
-                                            item,
-                                        )
+
+                        itemsIndexed(
+                            items = state.relayListItems,
+                            key = { index: Int, item: RelayListItem -> item.key },
+                            contentType = { _, item -> item.contentType },
+                            itemContent = { index: Int, listItem: RelayListItem ->
+                                Column(modifier = Modifier.animateItem()) {
+                                    if (index != 0) {
+                                        HorizontalDivider(color = backgroundColor)
+                                    }
+                                    when (listItem) {
+                                        RelayListItem.CustomListHeader ->
+                                            CustomListHeader(
+                                                onShowCustomListBottomSheet = {
+                                                    bottomSheetState =
+                                                        ShowCustomListsBottomSheet(
+                                                            editListEnabled =
+                                                                state.customLists.isNotEmpty()
+                                                        )
+                                                }
+                                            )
+                                        is RelayListItem.CustomListItem ->
+                                            CustomListItem(
+                                                listItem,
+                                                onSelectRelay,
+                                                {
+                                                    bottomSheetState =
+                                                        ShowEditCustomListBottomSheet(it)
+                                                },
+                                                { customListId, expand ->
+                                                    onToggleExpand(customListId, null, expand)
+                                                }
+                                            )
+                                        is RelayListItem.CustomListEntryItem ->
+                                            CustomListEntryItem(
+                                                listItem,
+                                                { onSelectRelay(listItem.item) },
+                                                if (listItem.depth == 1) {
+                                                    {
+                                                        bottomSheetState =
+                                                            ShowCustomListsEntryBottomSheet(
+                                                                listItem.parentId,
+                                                                listItem.item
+                                                            )
+                                                    }
+                                                } else {
+                                                    null
+                                                },
+                                                { expand: Boolean ->
+                                                    onToggleExpand(
+                                                        listItem.item.id,
+                                                        listItem.parentId,
+                                                        expand
+                                                    )
+                                                }
+                                            )
+                                        is RelayListItem.CustomListFooter ->
+                                            CustomListFooter(listItem)
+                                        RelayListItem.LocationHeader -> RelayLocationHeader()
+                                        is RelayListItem.GeoLocationItem ->
+                                            RelayLocationItem(
+                                                listItem,
+                                                { onSelectRelay(listItem.item) },
+                                                {
+                                                    // Only direct children can be removed
+                                                    bottomSheetState =
+                                                        ShowLocationBottomSheet(
+                                                            state.customLists,
+                                                            listItem.item
+                                                        )
+                                                },
+                                                { expand ->
+                                                    onToggleExpand(listItem.item.id, null, expand)
+                                                }
+                                            )
+                                        is RelayListItem.LocationsEmptyText ->
+                                            LocationsEmptyText(listItem.searchTerm)
+                                    }
                                 }
-                            )
-                            item {
-                                Spacer(
-                                    modifier =
-                                        Modifier.height(Dimens.mediumPadding)
-                                            .animateItemPlacement()
-                                            .animateContentSize()
-                                )
                             }
-                        }
-                        if (state.countries.isNotEmpty()) {
-                            relayList(
-                                countries = state.countries,
-                                selectedItem = state.selectedItem,
-                                onSelectRelay = onSelectRelay,
-                                onShowLocationBottomSheet = { location ->
-                                    bottomSheetState =
-                                        BottomSheetState.ShowLocationBottomSheet(
-                                            customLists = state.customLists,
-                                            item = location
-                                        )
-                                }
-                            )
-                        }
-                        if (state.showEmpty) {
-                            item { LocationsEmptyText(searchTerm = state.searchTerm) }
-                        }
+                        )
                     }
                 }
             }
@@ -404,6 +425,80 @@ fun SelectLocationScreen(
 }
 
 @Composable
+fun LazyItemScope.RelayLocationHeader() {
+    HeaderCell(text = stringResource(R.string.all_locations))
+}
+
+@Composable
+fun LazyItemScope.RelayLocationItem(
+    relayItem: RelayListItem.GeoLocationItem,
+    onSelectRelay: () -> Unit,
+    onLongClick: () -> Unit,
+    onExpand: (Boolean) -> Unit,
+) {
+    val location = relayItem.item
+    StatusRelayItemCell(
+        location,
+        relayItem.isSelected,
+        onClick = { onSelectRelay() },
+        onLongClick = { onLongClick() },
+        onToggleExpand = { onExpand(it) },
+        isExpanded = relayItem.expanded,
+        depth = relayItem.depth
+    )
+}
+
+@Composable
+fun LazyItemScope.CustomListItem(
+    itemState: RelayListItem.CustomListItem,
+    onSelectRelay: (item: RelayItem) -> Unit,
+    onShowEditBottomSheet: (RelayItem.CustomList) -> Unit,
+    onExpand: ((CustomListId, Boolean) -> Unit),
+) {
+    val customListItem = itemState.item
+    StatusRelayItemCell(
+        customListItem,
+        itemState.isSelected,
+        onClick = { onSelectRelay(customListItem) },
+        onLongClick = { onShowEditBottomSheet(customListItem) },
+        onToggleExpand = { onExpand(customListItem.id, it) },
+        isExpanded = itemState.expanded
+    )
+}
+
+@Composable
+fun LazyItemScope.CustomListEntryItem(
+    itemState: RelayListItem.CustomListEntryItem,
+    onSelectRelay: () -> Unit,
+    onShowEditCustomListEntryBottomSheet: (() -> Unit)?,
+    onToggleExpand: (Boolean) -> Unit,
+) {
+    val customListEntryItem = itemState.item
+    StatusRelayItemCell(
+        customListEntryItem,
+        false,
+        onClick = onSelectRelay,
+        onLongClick = onShowEditCustomListEntryBottomSheet,
+        onToggleExpand = onToggleExpand,
+        isExpanded = itemState.expanded,
+        depth = itemState.depth
+    )
+}
+
+@Composable
+fun LazyItemScope.CustomListFooter(item: RelayListItem.CustomListFooter) {
+    SwitchComposeSubtitleCell(
+        text =
+            if (item.hasCustomList) {
+                stringResource(R.string.to_add_locations_to_a_list)
+            } else {
+                stringResource(R.string.to_create_a_custom_list)
+            },
+        modifier = Modifier.background(MaterialTheme.colorScheme.background)
+    )
+}
+
+@Composable
 private fun SelectLocationTopBar(onBackClick: () -> Unit, onFilterClick: () -> Unit) {
     Row(modifier = Modifier.fillMaxWidth()) {
         IconButton(onClick = onBackClick) {
@@ -437,95 +532,13 @@ private fun LazyListScope.loading() {
     }
 }
 
-@OptIn(ExperimentalFoundationApi::class)
-private fun LazyListScope.customLists(
-    customLists: List<RelayItem.CustomList>,
-    selectedItem: RelayItemId?,
-    backgroundColor: Color,
-    onSelectRelay: (item: RelayItem) -> Unit,
-    onShowCustomListBottomSheet: () -> Unit,
-    onShowEditBottomSheet: (RelayItem.CustomList) -> Unit,
-    onShowEditCustomListEntryBottomSheet: (item: RelayItem.Location, RelayItem.CustomList) -> Unit
-) {
-    item(
-        contentType = { ContentType.HEADER },
-    ) {
-        ThreeDotCell(
-            text = stringResource(R.string.custom_lists),
-            onClickDots = onShowCustomListBottomSheet,
-            modifier =
-                Modifier.testTag(SELECT_LOCATION_CUSTOM_LIST_HEADER_TEST_TAG)
-                    .animateItemPlacement()
-                    .animateContentSize()
-        )
-    }
-    if (customLists.isNotEmpty()) {
-        items(
-            items = customLists,
-            key = { item -> item.id },
-            contentType = { ContentType.ITEM },
-        ) { customList ->
-            StatusRelayLocationCell(
-                relay = customList,
-                // Do not show selection for locations in custom lists
-                selectedItem = selectedItem as? CustomListId,
-                onSelectRelay = onSelectRelay,
-                onLongClick = {
-                    if (it is RelayItem.CustomList) {
-                        onShowEditBottomSheet(it)
-                    } else if (it is RelayItem.Location && it in customList.locations) {
-                        onShowEditCustomListEntryBottomSheet(it, customList)
-                    }
-                },
-                modifier = Modifier.animateContentSize().animateItemPlacement(),
-            )
-        }
-        item {
-            SwitchComposeSubtitleCell(
-                text = stringResource(R.string.to_add_locations_to_a_list),
-                modifier =
-                    Modifier.background(backgroundColor).animateItemPlacement().animateContentSize()
-            )
-        }
-    } else {
-        item(contentType = ContentType.EMPTY_TEXT) {
-            SwitchComposeSubtitleCell(
-                text = stringResource(R.string.to_create_a_custom_list),
-                modifier =
-                    Modifier.background(backgroundColor).animateItemPlacement().animateContentSize()
-            )
-        }
-    }
-}
-
-@OptIn(ExperimentalFoundationApi::class)
-private fun LazyListScope.relayList(
-    countries: List<RelayItem.Location.Country>,
-    selectedItem: RelayItemId?,
-    onSelectRelay: (item: RelayItem) -> Unit,
-    onShowLocationBottomSheet: (item: RelayItem.Location) -> Unit,
-) {
-    item(
-        contentType = ContentType.HEADER,
-    ) {
-        HeaderCell(
-            text = stringResource(R.string.all_locations),
-            modifier = Modifier.animateItemPlacement().animateContentSize()
-        )
-    }
-    items(
-        items = countries,
-        key = { item -> item.id },
-        contentType = { ContentType.ITEM },
-    ) { country ->
-        StatusRelayLocationCell(
-            relay = country,
-            selectedItem = selectedItem,
-            onSelectRelay = onSelectRelay,
-            onLongClick = { onShowLocationBottomSheet(it as RelayItem.Location) },
-            modifier = Modifier.animateContentSize().animateItemPlacement(),
-        )
-    }
+@Composable
+private fun LazyItemScope.CustomListHeader(onShowCustomListBottomSheet: () -> Unit) {
+    ThreeDotCell(
+        text = stringResource(R.string.custom_lists),
+        onClickDots = onShowCustomListBottomSheet,
+        modifier = Modifier.testTag(SELECT_LOCATION_CUSTOM_LIST_HEADER_TEST_TAG)
+    )
 }
 
 @OptIn(ExperimentalMaterial3Api::class)
@@ -535,7 +548,7 @@ private fun BottomSheets(
     onCreateCustomList: (RelayItem.Location?) -> Unit,
     onEditCustomLists: () -> Unit,
     onAddLocationToList: (RelayItem.Location, RelayItem.CustomList) -> Unit,
-    onRemoveLocationFromList: (RelayItem.Location, RelayItem.CustomList) -> Unit,
+    onRemoveLocationFromList: (location: RelayItem.Location, parent: CustomListId) -> Unit,
     onEditCustomListName: (RelayItem.CustomList) -> Unit,
     onEditLocationsCustomList: (RelayItem.CustomList) -> Unit,
     onDeleteCustomList: (RelayItem.CustomList) -> Unit,
@@ -559,7 +572,7 @@ private fun BottomSheets(
     val onBackgroundColor: Color = MaterialTheme.colorScheme.onSurface
 
     when (bottomSheetState) {
-        is BottomSheetState.ShowCustomListsBottomSheet -> {
+        is ShowCustomListsBottomSheet -> {
             CustomListsBottomSheet(
                 sheetState = sheetState,
                 onBackgroundColor = onBackgroundColor,
@@ -569,7 +582,7 @@ private fun BottomSheets(
                 closeBottomSheet = onCloseBottomSheet
             )
         }
-        is BottomSheetState.ShowLocationBottomSheet -> {
+        is ShowLocationBottomSheet -> {
             LocationBottomSheet(
                 sheetState = sheetState,
                 onBackgroundColor = onBackgroundColor,
@@ -580,7 +593,7 @@ private fun BottomSheets(
                 closeBottomSheet = onCloseBottomSheet
             )
         }
-        is BottomSheetState.ShowEditCustomListBottomSheet -> {
+        is ShowEditCustomListBottomSheet -> {
             EditCustomListBottomSheet(
                 sheetState = sheetState,
                 onBackgroundColor = onBackgroundColor,
@@ -591,11 +604,11 @@ private fun BottomSheets(
                 closeBottomSheet = onCloseBottomSheet
             )
         }
-        is BottomSheetState.ShowCustomListsEntryBottomSheet -> {
+        is ShowCustomListsEntryBottomSheet -> {
             CustomListEntryBottomSheet(
                 sheetState = sheetState,
                 onBackgroundColor = onBackgroundColor,
-                customList = bottomSheetState.customList,
+                customListId = bottomSheetState.parentId,
                 item = bottomSheetState.item,
                 onRemoveLocationFromList = onRemoveLocationFromList,
                 closeBottomSheet = onCloseBottomSheet
@@ -609,14 +622,16 @@ private fun BottomSheets(
 
 private fun SelectLocationUiState.indexOfSelectedRelayItem(): Int =
     if (this is SelectLocationUiState.Content) {
-        when (selectedItem) {
-            is CustomListId ->
-                filteredCustomLists.indexOfFirst { it.id == selectedItem } + EXTRA_ITEM_CUSTOM_LIST
-            is GeoLocationId ->
-                countries.indexOfFirst { it.id == selectedItem.country } +
-                    customLists.size +
-                    EXTRA_ITEMS_LOCATION
-            else -> -1
+        relayListItems.indexOfFirst {
+            when (it) {
+                is RelayListItem.CustomListItem -> it.isSelected
+                is RelayListItem.GeoLocationItem -> it.isSelected
+                is RelayListItem.CustomListEntryItem -> false
+                is RelayListItem.CustomListFooter -> false
+                RelayListItem.CustomListHeader -> false
+                RelayListItem.LocationHeader -> false
+                is RelayListItem.LocationsEmptyText -> false
+            }
         }
     } else {
         -1
@@ -627,7 +642,7 @@ private fun SelectLocationUiState.indexOfSelectedRelayItem(): Int =
 private fun CustomListsBottomSheet(
     onBackgroundColor: Color,
     sheetState: SheetState,
-    bottomSheetState: BottomSheetState.ShowCustomListsBottomSheet,
+    bottomSheetState: ShowCustomListsBottomSheet,
     onCreateCustomList: () -> Unit,
     onEditCustomLists: () -> Unit,
     closeBottomSheet: (animate: Boolean) -> Unit
@@ -787,10 +802,9 @@ private fun EditCustomListBottomSheet(
 private fun CustomListEntryBottomSheet(
     onBackgroundColor: Color,
     sheetState: SheetState,
-    customList: RelayItem.CustomList,
+    customListId: CustomListId,
     item: RelayItem.Location,
-    onRemoveLocationFromList:
-        (location: RelayItem.Location, customList: RelayItem.CustomList) -> Unit,
+    onRemoveLocationFromList: (location: RelayItem.Location, customListId: CustomListId) -> Unit,
     closeBottomSheet: (animate: Boolean) -> Unit
 ) {
     MullvadModalBottomSheet(
@@ -809,7 +823,7 @@ private fun CustomListEntryBottomSheet(
             title = stringResource(id = R.string.remove_button),
             titleColor = onBackgroundColor,
             onClick = {
-                onRemoveLocationFromList(item, customList)
+                onRemoveLocationFromList(item, customListId)
                 closeBottomSheet(true)
             },
             background = Color.Unspecified
@@ -879,16 +893,12 @@ private fun <D : DestinationSpec, R : CustomListSuccess> ResultRecipient<D, R>
     }
 }
 
-private const val EXTRA_ITEMS_LOCATION =
-    4 // Custom lists header, custom lists description, spacer, all locations header
-private const val EXTRA_ITEM_CUSTOM_LIST = 1 // Custom lists header
-
 sealed interface BottomSheetState {
 
     data class ShowCustomListsBottomSheet(val editListEnabled: Boolean) : BottomSheetState
 
     data class ShowCustomListsEntryBottomSheet(
-        val customList: RelayItem.CustomList,
+        val parentId: CustomListId,
         val item: RelayItem.Location
     ) : BottomSheetState
 
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/compose/state/FilterConstrainExtensions.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/state/FilterConstrainExtensions.kt
index 52ef7445b0..01fe84f76c 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/state/FilterConstrainExtensions.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/state/FilterConstrainExtensions.kt
@@ -5,12 +5,6 @@ import net.mullvad.mullvadvpn.lib.model.Ownership
 import net.mullvad.mullvadvpn.lib.model.Provider
 import net.mullvad.mullvadvpn.lib.model.Providers
 
-fun Constraint<Ownership>.toNullableOwnership(): Ownership? =
-    when (this) {
-        Constraint.Any -> null
-        is Constraint.Only -> this.value
-    }
-
 fun Ownership?.toOwnershipConstraint(): Constraint<Ownership> =
     when (this) {
         null -> Constraint.Any
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/state/SelectLocationUiState.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/state/SelectLocationUiState.kt
index 79f434aad1..5d6b683116 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/state/SelectLocationUiState.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/state/SelectLocationUiState.kt
@@ -1,9 +1,9 @@
 package net.mullvad.mullvadvpn.compose.state
 
-import net.mullvad.mullvadvpn.lib.model.Ownership
+import net.mullvad.mullvadvpn.lib.model.CustomListId
 import net.mullvad.mullvadvpn.lib.model.RelayItem
-import net.mullvad.mullvadvpn.lib.model.RelayItemId
-import net.mullvad.mullvadvpn.relaylist.MIN_SEARCH_LENGTH
+
+typealias ModelOwnership = net.mullvad.mullvadvpn.lib.model.Ownership
 
 sealed interface SelectLocationUiState {
 
@@ -11,18 +11,88 @@ sealed interface SelectLocationUiState {
 
     data class Content(
         val searchTerm: String,
-        val selectedOwnership: Ownership?,
-        val selectedProvidersCount: Int?,
-        val filteredCustomLists: List<RelayItem.CustomList>,
+        val filterChips: List<FilterChip>,
+        val relayListItems: List<RelayListItem>,
         val customLists: List<RelayItem.CustomList>,
-        val countries: List<RelayItem.Location.Country>,
-        val selectedItem: RelayItemId?
-    ) : SelectLocationUiState {
-        val hasFilter: Boolean = (selectedProvidersCount != null || selectedOwnership != null)
-        val inSearch = searchTerm.length >= MIN_SEARCH_LENGTH
-        val showCustomLists = inSearch.not() || filteredCustomLists.isNotEmpty()
-        // Show empty state if we don't have any relays or if we are searching and no custom list or
-        // relay is found
-        val showEmpty = countries.isEmpty() && (inSearch.not() || filteredCustomLists.isEmpty())
+    ) : SelectLocationUiState
+}
+
+sealed interface FilterChip {
+    data class Ownership(val ownership: ModelOwnership) : FilterChip
+
+    data class Provider(val count: Int) : FilterChip
+}
+
+enum class RelayListItemContentType {
+    CUSTOM_LIST_HEADER,
+    CUSTOM_LIST_ITEM,
+    CUSTOM_LIST_ENTRY_ITEM,
+    CUSTOM_LIST_FOOTER,
+    LOCATION_HEADER,
+    LOCATION_ITEM,
+    LOCATIONS_EMPTY_TEXT
+}
+
+sealed interface RelayListItem {
+    val key: Any
+    val contentType: RelayListItemContentType
+
+    data object CustomListHeader : RelayListItem {
+        override val key = "custom_list_header"
+        override val contentType = RelayListItemContentType.CUSTOM_LIST_HEADER
+    }
+
+    sealed interface SelectableItem : RelayListItem {
+        val depth: Int
+        val isSelected: Boolean
+        val expanded: Boolean
+    }
+
+    data class CustomListItem(
+        val item: RelayItem.CustomList,
+        override val isSelected: Boolean = false,
+        override val expanded: Boolean = false,
+    ) : SelectableItem {
+        override val key = item.id
+        override val depth: Int = 0
+        override val contentType = RelayListItemContentType.CUSTOM_LIST_ITEM
+    }
+
+    data class CustomListEntryItem(
+        val parentId: CustomListId,
+        val item: RelayItem.Location,
+        override val expanded: Boolean,
+        override val depth: Int = 0
+    ) : SelectableItem {
+        override val key = parentId to item.id
+
+        // Can't be displayed as selected
+        override val isSelected: Boolean = false
+        override val contentType = RelayListItemContentType.CUSTOM_LIST_ENTRY_ITEM
+    }
+
+    data class CustomListFooter(val hasCustomList: Boolean) : RelayListItem {
+        override val key = "custom_list_footer"
+        override val contentType = RelayListItemContentType.CUSTOM_LIST_FOOTER
+    }
+
+    data object LocationHeader : RelayListItem {
+        override val key: Any = "location_header"
+        override val contentType = RelayListItemContentType.LOCATION_HEADER
+    }
+
+    data class GeoLocationItem(
+        val item: RelayItem.Location,
+        override val isSelected: Boolean = false,
+        override val depth: Int = 0,
+        override val expanded: Boolean = false,
+    ) : SelectableItem {
+        override val key = item.id
+        override val contentType = RelayListItemContentType.LOCATION_ITEM
+    }
+
+    data class LocationsEmptyText(val searchTerm: String) : RelayListItem {
+        override val key: Any = "locations_empty_text"
+        override val contentType = RelayListItemContentType.LOCATIONS_EMPTY_TEXT
     }
 }
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/di/UiModule.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/di/UiModule.kt
index 371a30bdf1..6494cbb167 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/di/UiModule.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/di/UiModule.kt
@@ -44,6 +44,7 @@ import net.mullvad.mullvadvpn.usecase.VersionNotificationUseCase
 import net.mullvad.mullvadvpn.usecase.customlists.CustomListActionUseCase
 import net.mullvad.mullvadvpn.usecase.customlists.CustomListRelayItemsUseCase
 import net.mullvad.mullvadvpn.usecase.customlists.CustomListsRelayItemUseCase
+import net.mullvad.mullvadvpn.usecase.customlists.FilterCustomListsRelayItemUseCase
 import net.mullvad.mullvadvpn.util.ChangelogDataProvider
 import net.mullvad.mullvadvpn.util.IChangelogDataProvider
 import net.mullvad.mullvadvpn.viewmodel.AccountViewModel
@@ -134,6 +135,7 @@ val uiModule = module {
     single { CustomListActionUseCase(get(), get()) }
     single { SelectedLocationTitleUseCase(get(), get()) }
     single { AvailableProvidersUseCase(get()) }
+    single { FilterCustomListsRelayItemUseCase(get(), get()) }
     single { CustomListsRelayItemUseCase(get(), get()) }
     single { CustomListRelayItemsUseCase(get(), get()) }
     single { FilteredRelayListUseCase(get(), get()) }
@@ -183,7 +185,7 @@ val uiModule = module {
     viewModel { DnsDialogViewModel(get(), get(), get()) }
     viewModel { LoginViewModel(get(), get(), get()) }
     viewModel { PrivacyDisclaimerViewModel(get(), IS_PLAY_BUILD) }
-    viewModel { SelectLocationViewModel(get(), get(), get(), get(), get(), get()) }
+    viewModel { SelectLocationViewModel(get(), get(), get(), get(), get(), get(), get(), get()) }
     viewModel { SettingsViewModel(get(), get(), IS_PLAY_BUILD) }
     viewModel { SplashViewModel(get(), get(), get(), get()) }
     viewModel { VoucherDialogViewModel(get()) }
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/relaylist/CustomListExtensions.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/relaylist/CustomListExtensions.kt
index 2a7eeddb69..ac03080e21 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/relaylist/CustomListExtensions.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/relaylist/CustomListExtensions.kt
@@ -8,9 +8,7 @@ fun CustomList.toRelayItemCustomList(
     relayCountries: List<RelayItem.Location.Country>
 ): RelayItem.CustomList =
     RelayItem.CustomList(
-        id = id,
-        customListName = name,
-        expanded = false,
+        customList = this,
         locations = locations.mapNotNull { relayCountries.findByGeoLocationId(it) },
     )
 
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/relaylist/RelayListExtensions.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/relaylist/RelayListExtensions.kt
index 069f0e1a08..ea017339f6 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/relaylist/RelayListExtensions.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/relaylist/RelayListExtensions.kt
@@ -1,9 +1,7 @@
 package net.mullvad.mullvadvpn.relaylist
 
-import net.mullvad.mullvadvpn.lib.model.CustomListId
 import net.mullvad.mullvadvpn.lib.model.GeoLocationId
 import net.mullvad.mullvadvpn.lib.model.RelayItem
-import net.mullvad.mullvadvpn.lib.model.RelayItemId
 
 fun List<RelayItem.Location.Country>.findByGeoLocationId(geoLocationId: GeoLocationId) =
     withDescendants().firstOrNull { it.id == geoLocationId }
@@ -11,119 +9,46 @@ fun List<RelayItem.Location.Country>.findByGeoLocationId(geoLocationId: GeoLocat
 fun List<RelayItem.Location.Country>.findByGeoLocationId(geoLocationId: GeoLocationId.City) =
     flatMap { it.cities }.firstOrNull { it.id == geoLocationId }
 
-/**
- * Filter and expand the list based on search terms If a country is matched, that country and all
- * its children are added to the list, but the country is not expanded If a city is matched, its
- * parent country is added and expanded if needed and its children are added, but the city is not
- * expanded If a relay is matched, its parents are added and expanded and itself is also added.
- */
-@Suppress("NestedBlockDepth")
-fun List<RelayItem.Location.Country>.filterOnSearchTerm(
-    searchTerm: String,
-    selectedItem: RelayItemId?
-): List<RelayItem.Location.Country> {
-    return if (searchTerm.length >= MIN_SEARCH_LENGTH) {
-        val filteredCountries = mutableMapOf<GeoLocationId.Country, RelayItem.Location.Country>()
-        this.forEach { relayCountry ->
-            val cities = mutableListOf<RelayItem.Location.City>()
+fun List<RelayItem.Location.Country>.search(searchTerm: String): List<GeoLocationId> =
+    withDescendants().filter { it.name.contains(searchTerm, ignoreCase = true) }.map { it.id }
 
-            // Try to match the search term with a country
-            // If we match a country, add that country and all cities and relays in that country
-            // Do not currently expand the country or any city
-            if (relayCountry.name.contains(other = searchTerm, ignoreCase = true)) {
-                cities.addAll(relayCountry.cities.map { city -> city.copy(expanded = false) })
-                filteredCountries[relayCountry.id] =
-                    relayCountry.copy(expanded = false, cities = cities)
-            }
+fun List<GeoLocationId>.expansionSet() = flatMap { it.ancestors() }.toSet()
 
-            // Go through and try to match the search term with every city
-            relayCountry.cities.forEach { relayCity ->
-                val relays = mutableListOf<RelayItem.Location.Relay>()
-                // If we match and we already added the country to the filtered list just expand the
-                // country.
-                // If the country is not currently in the filtered list, add it and expand it.
-                // Finally if the city has not already been added to the filtered list, add it, but
-                // do not expand it yet.
-                if (relayCity.name.contains(other = searchTerm, ignoreCase = true)) {
-                    val value = filteredCountries[relayCountry.id]
-                    if (value != null) {
-                        filteredCountries[relayCountry.id] = value.copy(expanded = true)
-                    } else {
-                        filteredCountries[relayCountry.id] =
-                            relayCountry.copy(expanded = true, cities = cities)
-                    }
-                    if (cities.none { city -> city.id == relayCity.id }) {
-                        cities.add(relayCity.copy(expanded = false))
-                    }
-                }
-
-                // Go through and try to match the search term with every relay
-                relayCity.relays.forEach { relay ->
-                    // If we match a relay, check if the county the relay is in already is added, if
-                    // so expand, if not add and expand the country.
-                    // Check if the city that the relay is in is already added to the filtered list,
-                    // if so expand it, if not add it to the filtered list and expand it.
-                    // Finally add the relay to the list.
-                    if (relay.name.contains(other = searchTerm, ignoreCase = true)) {
-                        val value = filteredCountries[relayCountry.id]
-                        if (value != null) {
-                            filteredCountries[relayCountry.id] = value.copy(expanded = true)
-                        } else {
-                            filteredCountries[relayCountry.id] =
-                                relayCountry.copy(expanded = true, cities = cities)
-                        }
-                        val cityIndex = cities.indexOfFirst { it.id == relayCity.id }
+fun List<RelayItem.Location.Country>.newFilterOnSearch(
+    searchTerm: String
+): Pair<Set<GeoLocationId>, List<RelayItem.Location.Country>> {
+    val matchesIds = search(searchTerm)
+    val expansionSet = matchesIds.expansionSet()
 
-                        // No city found
-                        if (cityIndex < 0) {
-                            cities.add(relayCity.copy(expanded = true, relays = relays))
-                        } else {
-                            // Update found city as expanded
-                            cities[cityIndex] = cities[cityIndex].copy(expanded = true)
-                        }
-
-                        relays.add(relay.copy())
+    val filteredCountryList = mapNotNull { country ->
+        if (country.id in matchesIds) {
+            country
+        } else if (country.id in expansionSet) {
+            country.copy(
+                cities =
+                    country.cities.mapNotNull { city ->
+                        if (city.id in matchesIds) {
+                            city
+                        } else if (city.id in expansionSet) {
+                            city.copy(
+                                relays = city.relays.filter { relay -> relay.id in matchesIds }
+                            )
+                        } else null
                     }
-                }
-            }
+            )
+        } else {
+            null
         }
-        filteredCountries.values.sortedBy { it.name }
-    } else {
-        this.expandItemForSelection(selectedItem)
     }
+    return expansionSet to filteredCountryList
 }
 
-/** Expand the parent(s), if any, for the current selected item */
-private fun List<RelayItem.Location.Country>.expandItemForSelection(
-    selectedItem: RelayItemId?
-): List<RelayItem.Location.Country> {
-    selectedItem ?: return this
-    return when (selectedItem) {
-        is CustomListId,
-        is GeoLocationId.Country -> this
-        is GeoLocationId.City ->
-            map { if (it.id == selectedItem.country) it.copy(expanded = true) else it }
-        is GeoLocationId.Hostname -> {
-            map { country ->
-                if (country.id == selectedItem.country) {
-                    country.copy(
-                        expanded = true,
-                        cities =
-                            country.cities.map { city ->
-                                if (city.id == selectedItem.city) {
-                                    city.copy(expanded = true)
-                                } else {
-                                    city
-                                }
-                            },
-                    )
-                } else {
-                    country
-                }
-            }
-        }
+fun GeoLocationId.ancestors(): List<GeoLocationId> =
+    when (this) {
+        is GeoLocationId.City -> listOf(country)
+        is GeoLocationId.Country -> emptyList()
+        is GeoLocationId.Hostname -> listOf(country, city)
     }
-}
 
 fun List<RelayItem.Location.Country>.getRelayItemsByCodes(
     codes: List<GeoLocationId>
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/usecase/customlists/FilterCustomListsRelayItemUseCase.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/usecase/customlists/FilterCustomListsRelayItemUseCase.kt
new file mode 100644
index 0000000000..f82d9eed5e
--- /dev/null
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/usecase/customlists/FilterCustomListsRelayItemUseCase.kt
@@ -0,0 +1,30 @@
+package net.mullvad.mullvadvpn.usecase.customlists
+
+import kotlin.collections.mapNotNull
+import kotlinx.coroutines.flow.combine
+import net.mullvad.mullvadvpn.lib.model.Constraint
+import net.mullvad.mullvadvpn.lib.model.Ownership
+import net.mullvad.mullvadvpn.lib.model.Providers
+import net.mullvad.mullvadvpn.lib.model.RelayItem
+import net.mullvad.mullvadvpn.relaylist.filterOnOwnershipAndProvider
+import net.mullvad.mullvadvpn.repository.RelayListFilterRepository
+
+class FilterCustomListsRelayItemUseCase(
+    private val customListsRelayItemUseCase: CustomListsRelayItemUseCase,
+    private val relayListFilterRepository: RelayListFilterRepository
+) {
+
+    operator fun invoke() =
+        combine(
+            customListsRelayItemUseCase(),
+            relayListFilterRepository.selectedOwnership,
+            relayListFilterRepository.selectedProviders,
+        ) { customLists, selectedOwnership, selectedProviders ->
+            customLists.filterOnOwnershipAndProvider(selectedOwnership, selectedProviders)
+        }
+
+    private fun List<RelayItem.CustomList>.filterOnOwnershipAndProvider(
+        ownership: Constraint<Ownership>,
+        providers: Constraint<Providers>
+    ) = mapNotNull { it.filterOnOwnershipAndProvider(ownership, providers) }
+}
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/util/FlowUtils.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/util/FlowUtils.kt
index 13561aa7f8..368c614aae 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/util/FlowUtils.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/util/FlowUtils.kt
@@ -28,58 +28,6 @@ inline fun <T1, T2, T3, T4, T5, T6, R> combine(
     }
 }
 
-inline fun <T1, T2, T3, T4, T5, T6, T7, R> combine(
-    flow: Flow<T1>,
-    flow2: Flow<T2>,
-    flow3: Flow<T3>,
-    flow4: Flow<T4>,
-    flow5: Flow<T5>,
-    flow6: Flow<T6>,
-    flow7: Flow<T7>,
-    crossinline transform: suspend (T1, T2, T3, T4, T5, T6, T7) -> R
-): Flow<R> {
-    return kotlinx.coroutines.flow.combine(flow, flow2, flow3, flow4, flow5, flow6, flow7) {
-        args: Array<*> ->
-        @Suppress("UNCHECKED_CAST")
-        transform(
-            args[0] as T1,
-            args[1] as T2,
-            args[2] as T3,
-            args[3] as T4,
-            args[4] as T5,
-            args[5] as T6,
-            args[6] as T7
-        )
-    }
-}
-
-inline fun <T1, T2, T3, T4, T5, T6, T7, T8, R> combine(
-    flow: Flow<T1>,
-    flow2: Flow<T2>,
-    flow3: Flow<T3>,
-    flow4: Flow<T4>,
-    flow5: Flow<T5>,
-    flow6: Flow<T6>,
-    flow7: Flow<T7>,
-    flow8: Flow<T8>,
-    crossinline transform: suspend (T1, T2, T3, T4, T5, T6, T7, T8) -> R
-): Flow<R> {
-    return kotlinx.coroutines.flow.combine(flow, flow2, flow3, flow4, flow5, flow6, flow7, flow8) {
-        args: Array<*> ->
-        @Suppress("UNCHECKED_CAST")
-        transform(
-            args[0] as T1,
-            args[1] as T2,
-            args[2] as T3,
-            args[3] as T4,
-            args[4] as T5,
-            args[5] as T6,
-            args[6] as T7,
-            args[7] as T8
-        )
-    }
-}
-
 fun <T> Deferred<T>.getOrDefault(default: T) =
     try {
         getCompleted()
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 {
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/viewmodel/FilterViewModel.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/viewmodel/FilterViewModel.kt
index 728142b3ff..1529bb4221 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/viewmodel/FilterViewModel.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/viewmodel/FilterViewModel.kt
@@ -13,7 +13,6 @@ import kotlinx.coroutines.flow.stateIn
 import kotlinx.coroutines.launch
 import net.mullvad.mullvadvpn.compose.state.RelayFilterState
 import net.mullvad.mullvadvpn.compose.state.toConstraintProviders
-import net.mullvad.mullvadvpn.compose.state.toNullableOwnership
 import net.mullvad.mullvadvpn.compose.state.toOwnershipConstraint
 import net.mullvad.mullvadvpn.compose.state.toSelectedProviders
 import net.mullvad.mullvadvpn.lib.model.Ownership
@@ -43,7 +42,7 @@ class FilterViewModel(
                     .first()
 
             val ownershipConstraint = relayListFilterRepository.selectedOwnership.first()
-            selectedOwnership.value = ownershipConstraint.toNullableOwnership()
+            selectedOwnership.value = ownershipConstraint.getOrNull()
         }
     }
 
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/viewmodel/SelectLocationViewModel.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/viewmodel/SelectLocationViewModel.kt
index 2509fdc876..f31fcb3078 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/viewmodel/SelectLocationViewModel.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/viewmodel/SelectLocationViewModel.kt
@@ -2,93 +2,71 @@ package net.mullvad.mullvadvpn.viewmodel
 
 import androidx.lifecycle.ViewModel
 import androidx.lifecycle.viewModelScope
+import arrow.core.raise.either
 import kotlinx.coroutines.channels.Channel
 import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.flow.SharingStarted
+import kotlinx.coroutines.flow.combine
+import kotlinx.coroutines.flow.map
+import kotlinx.coroutines.flow.onEach
 import kotlinx.coroutines.flow.receiveAsFlow
 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.FilterChip
+import net.mullvad.mullvadvpn.compose.state.RelayListItem
+import net.mullvad.mullvadvpn.compose.state.RelayListItem.CustomListHeader
 import net.mullvad.mullvadvpn.compose.state.SelectLocationUiState
-import net.mullvad.mullvadvpn.compose.state.toNullableOwnership
+import net.mullvad.mullvadvpn.compose.state.SelectLocationUiState.Content
 import net.mullvad.mullvadvpn.compose.state.toSelectedProviders
 import net.mullvad.mullvadvpn.lib.model.Constraint
+import net.mullvad.mullvadvpn.lib.model.CustomListId
+import net.mullvad.mullvadvpn.lib.model.GeoLocationId
 import net.mullvad.mullvadvpn.lib.model.Ownership
 import net.mullvad.mullvadvpn.lib.model.Provider
-import net.mullvad.mullvadvpn.lib.model.Providers
 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.descendants
-import net.mullvad.mullvadvpn.relaylist.filterOnOwnershipAndProvider
 import net.mullvad.mullvadvpn.relaylist.filterOnSearchTerm
+import net.mullvad.mullvadvpn.relaylist.newFilterOnSearch
+import net.mullvad.mullvadvpn.repository.CustomListsRepository
 import net.mullvad.mullvadvpn.repository.RelayListFilterRepository
 import net.mullvad.mullvadvpn.repository.RelayListRepository
 import net.mullvad.mullvadvpn.usecase.AvailableProvidersUseCase
 import net.mullvad.mullvadvpn.usecase.FilteredRelayListUseCase
 import net.mullvad.mullvadvpn.usecase.customlists.CustomListActionUseCase
 import net.mullvad.mullvadvpn.usecase.customlists.CustomListsRelayItemUseCase
-import net.mullvad.mullvadvpn.util.combine
+import net.mullvad.mullvadvpn.usecase.customlists.FilterCustomListsRelayItemUseCase
 
 class SelectLocationViewModel(
     private val relayListFilterRepository: RelayListFilterRepository,
-    availableProvidersUseCase: AvailableProvidersUseCase,
-    customListsRelayItemUseCase: CustomListsRelayItemUseCase,
+    private val availableProvidersUseCase: AvailableProvidersUseCase,
+    private val customListsRelayItemUseCase: CustomListsRelayItemUseCase,
+    private val filteredCustomListRelayItemsUseCase: FilterCustomListsRelayItemUseCase,
+    private val customListsRepository: CustomListsRepository,
     private val customListActionUseCase: CustomListActionUseCase,
-    filteredRelayListUseCase: FilteredRelayListUseCase,
+    private val filteredRelayListUseCase: FilteredRelayListUseCase,
     private val relayListRepository: RelayListRepository
 ) : ViewModel() {
     private val _searchTerm = MutableStateFlow(EMPTY_SEARCH_TERM)
 
+    private val _expandedItems = MutableStateFlow(initialExpand())
+
     @Suppress("DestructuringDeclarationWithTooManyEntries")
     val uiState =
-        combine(
-                filteredRelayListUseCase(),
-                customListsRelayItemUseCase(),
-                relayListRepository.selectedLocation,
-                _searchTerm,
-                relayListFilterRepository.selectedOwnership,
-                availableProvidersUseCase(),
-                relayListFilterRepository.selectedProviders,
-            ) {
-                relayCountries,
-                customLists,
-                selectedItem,
+        combine(_searchTerm, relayListItems(), filterChips(), customListsRelayItemUseCase()) {
                 searchTerm,
-                selectedOwnership,
-                allProviders,
-                selectedConstraintProviders ->
-                val selectRelayItemId = selectedItem.getOrNull()
-                val selectedOwnershipItem = selectedOwnership.toNullableOwnership()
-                val selectedProvidersCount =
-                    when (selectedConstraintProviders) {
-                        is Constraint.Any -> null
-                        is Constraint.Only ->
-                            filterSelectedProvidersByOwnership(
-                                    selectedConstraintProviders.toSelectedProviders(allProviders),
-                                    selectedOwnershipItem,
-                                )
-                                .size
-                    }
-
-                val filteredRelayCountries =
-                    relayCountries.filterOnSearchTerm(searchTerm, selectRelayItemId)
-
-                val filteredCustomLists =
-                    customLists
-                        .filterOnSearchTerm(searchTerm)
-                        .filterOnOwnershipAndProvider(
-                            ownership = selectedOwnership,
-                            providers = selectedConstraintProviders,
-                        )
-
-                SelectLocationUiState.Content(
+                relayListItems,
+                filterChips,
+                customLists ->
+                Content(
                     searchTerm = searchTerm,
-                    selectedOwnership = selectedOwnershipItem,
-                    selectedProvidersCount = selectedProvidersCount,
-                    filteredCustomLists = filteredCustomLists,
-                    customLists = customLists,
-                    countries = filteredRelayCountries,
-                    selectedItem = selectRelayItemId,
+                    filterChips = filterChips,
+                    relayListItems = relayListItems,
+                    customLists = customLists
                 )
             }
             .stateIn(
@@ -100,18 +78,268 @@ class SelectLocationViewModel(
     private val _uiSideEffect = Channel<SelectLocationSideEffect>()
     val uiSideEffect = _uiSideEffect.receiveAsFlow()
 
+    private fun initialExpand(): Set<String> = buildSet {
+        val item = relayListRepository.selectedLocation.value.getOrNull()
+        when (item) {
+            is GeoLocationId.City -> {
+                add(item.country.code)
+            }
+            is GeoLocationId.Hostname -> {
+                add(item.country.code)
+                add(item.city.code)
+            }
+            is CustomListId,
+            is GeoLocationId.Country,
+            null -> {
+                /* No expands */
+            }
+        }
+    }
+
+    private fun searchRelayListLocations() =
+        combine(
+                _searchTerm,
+                filteredRelayListUseCase(),
+            ) { searchTerm, relayCountries ->
+                val isSearching = searchTerm.length >= MIN_SEARCH_LENGTH
+                if (isSearching) {
+                    val (exp, filteredRelayCountries) = relayCountries.newFilterOnSearch(searchTerm)
+                    exp.map { it.expandKey() }.toSet() to filteredRelayCountries
+                } else {
+                    initialExpand() to relayCountries
+                }
+            }
+            .onEach { _expandedItems.value = it.first }
+            .map { it.second }
+
+    private fun filterChips() =
+        combine(
+            relayListFilterRepository.selectedOwnership,
+            relayListFilterRepository.selectedProviders,
+            availableProvidersUseCase(),
+        ) { selectedOwnership, selectedConstraintProviders, allProviders,
+            ->
+            val ownershipFilter = selectedOwnership.getOrNull()
+            val providerCountFilter =
+                when (selectedConstraintProviders) {
+                    is Constraint.Any -> null
+                    is Constraint.Only ->
+                        filterSelectedProvidersByOwnership(
+                                selectedConstraintProviders.toSelectedProviders(allProviders),
+                                ownershipFilter,
+                            )
+                            .size
+                }
+
+            buildList<FilterChip> {
+                if (ownershipFilter != null) {
+                    add(FilterChip.Ownership(ownershipFilter))
+                }
+                if (providerCountFilter != null) {
+                    add(FilterChip.Provider(providerCountFilter))
+                }
+            }
+        }
+
+    private fun relayListItems() =
+        combine(
+            _searchTerm,
+            searchRelayListLocations(),
+            filteredCustomListRelayItemsUseCase(),
+            relayListRepository.selectedLocation,
+            _expandedItems,
+        ) { searchTerm, relayCountries, customLists, selectedItem, expandedItems ->
+            val filteredCustomLists = customLists.filterOnSearchTerm(searchTerm)
+
+            buildList {
+                val relayItems =
+                    createRelayListItems(
+                        searchTerm.length >= MIN_SEARCH_LENGTH,
+                        selectedItem.getOrNull(),
+                        filteredCustomLists,
+                        relayCountries,
+                        { it in expandedItems }
+                    )
+                if (relayItems.isEmpty()) {
+                    add(RelayListItem.LocationsEmptyText(searchTerm))
+                } else {
+                    addAll(relayItems)
+                }
+            }
+        }
+
+    private fun createRelayListItems(
+        isSearching: Boolean,
+        selectedItem: RelayItemId?,
+        customLists: List<RelayItem.CustomList>,
+        countries: List<RelayItem.Location.Country>,
+        isExpanded: (String) -> Boolean
+    ): List<RelayListItem> =
+        createCustomListSection(isSearching, selectedItem, customLists, isExpanded) +
+            createLocationSection(isSearching, selectedItem, countries, isExpanded)
+
+    private fun createCustomListSection(
+        isSearching: Boolean,
+        selectedItem: RelayItemId?,
+        customLists: List<RelayItem.CustomList>,
+        isExpanded: (String) -> Boolean
+    ): List<RelayListItem> = buildList {
+        if (isSearching && customLists.isEmpty()) {
+            // If we are searching and no results are found don't show header or footer
+        } else {
+            add(CustomListHeader)
+            val customListItems = createCustomListRelayItems(customLists, selectedItem, isExpanded)
+            addAll(customListItems)
+            add(RelayListItem.CustomListFooter(customListItems.isNotEmpty()))
+        }
+    }
+
+    private fun createCustomListRelayItems(
+        customLists: List<RelayItem.CustomList>,
+        selectedItem: RelayItemId?,
+        isExpanded: (String) -> Boolean
+    ): List<RelayListItem> =
+        customLists.flatMap { customList ->
+            val expanded = isExpanded(customList.id.expandKey())
+            buildList<RelayListItem> {
+                add(
+                    RelayListItem.CustomListItem(
+                        customList,
+                        isSelected = selectedItem == customList.id,
+                        expanded
+                    )
+                )
+
+                if (expanded) {
+                    addAll(
+                        customList.locations.flatMap {
+                            createCustomListEntry(parent = customList.id, item = it, 1, isExpanded)
+                        }
+                    )
+                }
+            }
+        }
+
+    private fun createLocationSection(
+        isSearching: Boolean,
+        selectedItem: RelayItemId?,
+        countries: List<RelayItem.Location.Country>,
+        isExpanded: (String) -> Boolean
+    ): List<RelayListItem> = buildList {
+        if (isSearching && countries.isEmpty()) {
+            // If we are searching and no results are found don't show header or footer
+        } else {
+            add(RelayListItem.LocationHeader)
+            addAll(
+                countries.flatMap { country ->
+                    createGeoLocationEntry(country, selectedItem, isExpanded = isExpanded)
+                }
+            )
+        }
+    }
+
+    private fun createCustomListEntry(
+        parent: CustomListId,
+        item: RelayItem.Location,
+        depth: Int = 1,
+        isExpanded: (String) -> Boolean,
+    ): List<RelayListItem.CustomListEntryItem> =
+        buildList<RelayListItem.CustomListEntryItem> {
+            val expanded = isExpanded(item.id.expandKey(parent))
+            add(
+                RelayListItem.CustomListEntryItem(
+                    parentId = parent,
+                    item = item,
+                    expanded = expanded,
+                    depth
+                )
+            )
+
+            if (expanded) {
+                when (item) {
+                    is RelayItem.Location.City ->
+                        addAll(
+                            item.relays.flatMap {
+                                createCustomListEntry(parent, it, depth + 1, isExpanded)
+                            }
+                        )
+                    is RelayItem.Location.Country ->
+                        addAll(
+                            item.cities.flatMap {
+                                createCustomListEntry(parent, it, depth + 1, isExpanded)
+                            }
+                        )
+                    is RelayItem.Location.Relay -> {} // No children to add
+                }
+            }
+        }
+
+    private fun createGeoLocationEntry(
+        item: RelayItem.Location,
+        selectedItem: RelayItemId?,
+        depth: Int = 0,
+        isExpanded: (String) -> Boolean
+    ): List<RelayListItem.GeoLocationItem> = buildList {
+        val expanded = isExpanded(item.id.expandKey())
+
+        add(
+            RelayListItem.GeoLocationItem(
+                item = item,
+                isSelected = selectedItem == item.id,
+                depth = depth,
+                expanded = expanded,
+            )
+        )
+
+        if (expanded) {
+            when (item) {
+                is RelayItem.Location.City ->
+                    addAll(
+                        item.relays.flatMap {
+                            createGeoLocationEntry(it, selectedItem, depth + 1, isExpanded)
+                        }
+                    )
+                is RelayItem.Location.Country ->
+                    addAll(
+                        item.cities.flatMap {
+                            createGeoLocationEntry(it, selectedItem, depth + 1, isExpanded)
+                        }
+                    )
+                is RelayItem.Location.Relay -> {} // Do nothing
+            }
+        }
+    }
+
+    private fun RelayItemId.expandKey(parent: CustomListId? = null) =
+        (parent?.value ?: "") +
+            when (this) {
+                is CustomListId -> value
+                is GeoLocationId -> code
+            }
+
     fun selectRelay(relayItem: RelayItem) {
         viewModelScope.launch {
             val locationConstraint = relayItem.id
             relayListRepository
                 .updateSelectedRelayLocation(locationConstraint)
                 .fold(
-                    { _uiSideEffect.trySend(SelectLocationSideEffect.GenericError) },
-                    { _uiSideEffect.trySend(SelectLocationSideEffect.CloseScreen) },
+                    { _uiSideEffect.send(SelectLocationSideEffect.GenericError) },
+                    { _uiSideEffect.send(SelectLocationSideEffect.CloseScreen) },
                 )
         }
     }
 
+    fun onToggleExpand(item: RelayItemId, parent: CustomListId? = null, expand: Boolean) {
+        _expandedItems.update {
+            val key = item.expandKey(parent)
+            if (expand) {
+                it + key
+            } else {
+                it - key
+            }
+        }
+    }
+
     fun onSearchTermInput(searchTerm: String) {
         viewModelScope.launch { _searchTerm.emit(searchTerm) }
     }
@@ -147,28 +375,27 @@ class SelectLocationViewModel(
         viewModelScope.launch { customListActionUseCase(action) }
     }
 
-    fun removeLocationFromList(item: RelayItem.Location, customList: RelayItem.CustomList) {
+    fun removeLocationFromList(item: RelayItem.Location, customListId: CustomListId) {
         viewModelScope.launch {
-            val newLocations = (customList.locations - item).map { it.id }
-            customListActionUseCase(CustomListAction.UpdateLocations(customList.id, newLocations))
-                .fold(
-                    { _uiSideEffect.send(SelectLocationSideEffect.GenericError) },
-                    {
-                        _uiSideEffect.send(
-                            SelectLocationSideEffect.LocationRemovedFromCustomList(it)
-                        )
+            val result =
+                either {
+                        val customList =
+                            customListsRepository.getCustomListById(customListId).bind()
+                        val newLocations = (customList.locations - item.id)
+
+                        customListActionUseCase(
+                                CustomListAction.UpdateLocations(customList.id, newLocations)
+                            )
+                            .bind()
                     }
-                )
+                    .fold(
+                        { SelectLocationSideEffect.GenericError },
+                        { SelectLocationSideEffect.LocationRemovedFromCustomList(it) }
+                    )
+            _uiSideEffect.send(result)
         }
     }
 
-    private fun List<RelayItem.CustomList>.filterOnOwnershipAndProvider(
-        ownership: Constraint<Ownership>,
-        providers: Constraint<Providers>
-    ): List<RelayItem.CustomList> = map { item ->
-        item.filterOnOwnershipAndProvider(ownership, providers)
-    }
-
     companion object {
         private const val EMPTY_SEARCH_TERM = ""
     }
diff --git a/android/app/src/test/kotlin/net/mullvad/mullvadvpn/usecase/CustomListActionUseCaseTest.kt b/android/app/src/test/kotlin/net/mullvad/mullvadvpn/usecase/CustomListActionUseCaseTest.kt
index d98292ccd9..d72d183202 100644
--- a/android/app/src/test/kotlin/net/mullvad/mullvadvpn/usecase/CustomListActionUseCaseTest.kt
+++ b/android/app/src/test/kotlin/net/mullvad/mullvadvpn/usecase/CustomListActionUseCaseTest.kt
@@ -71,7 +71,6 @@ class CustomListActionUseCaseTest {
                 RelayItem.Location.Country(
                     id = locationId,
                     name = locationName,
-                    expanded = false,
                     cities = emptyList()
                 )
             )
@@ -151,7 +150,7 @@ class CustomListActionUseCaseTest {
         val action = CustomListAction.Delete(id = customListId)
         val expectedResult =
             Deleted(undo = action.not(name = name, locations = listOf(location))).right()
-        every { mockLocation.countryCode } returns location.countryCode
+        every { mockLocation.code } returns location.code
         coEvery { mockCustomListsRepository.deleteCustomList(id = customListId) } returns
             Unit.right()
         coEvery { mockCustomListsRepository.getCustomListById(customListId) } returns
diff --git a/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/CustomListLocationsViewModelTest.kt b/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/CustomListLocationsViewModelTest.kt
index 4879031ce7..d2eaedd8c2 100644
--- a/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/CustomListLocationsViewModelTest.kt
+++ b/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/CustomListLocationsViewModelTest.kt
@@ -13,7 +13,9 @@ import net.mullvad.mullvadvpn.compose.communication.CustomListAction
 import net.mullvad.mullvadvpn.compose.communication.LocationsChanged
 import net.mullvad.mullvadvpn.compose.screen.CustomListLocationsNavArgs
 import net.mullvad.mullvadvpn.compose.state.CustomListLocationsUiState
+import net.mullvad.mullvadvpn.compose.state.RelayLocationListItem
 import net.mullvad.mullvadvpn.lib.common.test.TestCoroutineRule
+import net.mullvad.mullvadvpn.lib.common.test.assertLists
 import net.mullvad.mullvadvpn.lib.model.CustomList
 import net.mullvad.mullvadvpn.lib.model.CustomListId
 import net.mullvad.mullvadvpn.lib.model.CustomListName
@@ -23,6 +25,7 @@ import net.mullvad.mullvadvpn.lib.model.Provider
 import net.mullvad.mullvadvpn.lib.model.ProviderId
 import net.mullvad.mullvadvpn.lib.model.RelayItem
 import net.mullvad.mullvadvpn.relaylist.descendants
+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
@@ -66,15 +69,20 @@ class CustomListLocationsViewModelTest {
     fun `when selected locations is not null and relay countries is not empty should return ui state content`() =
         runTest {
             // Arrange
-            val expectedList = DUMMY_COUNTRIES
+            val expectedList =
+                DUMMY_COUNTRIES.map {
+                    RelayLocationListItem(
+                        item = it,
+                        depth = it.toDepth(),
+                        checked = false,
+                        expanded = false
+                    )
+                }
             val customListId = CustomListId("id")
             val expectedState =
-                CustomListLocationsUiState.Content.Data(
-                    newList = true,
-                    availableLocations = expectedList
-                )
+                CustomListLocationsUiState.Content.Data(newList = true, locations = expectedList)
             val viewModel = createViewModel(customListId, true)
-            relayListFlow.value = expectedList
+            relayListFlow.value = DUMMY_COUNTRIES
 
             // Act, Assert
             viewModel.uiState.test { assertEquals(expectedState, awaitItem()) }
@@ -85,8 +93,8 @@ class CustomListLocationsViewModelTest {
         // Arrange
         val expectedList = DUMMY_COUNTRIES
         val customListId = CustomListId("id")
-        val expectedSelection =
-            (DUMMY_COUNTRIES + DUMMY_COUNTRIES.flatMap { it.descendants() }).toSet()
+        val expectedSelection = (DUMMY_COUNTRIES.take(1).withDescendants()).map { it.id }
+
         val viewModel = createViewModel(customListId, true)
         relayListFlow.value = expectedList
 
@@ -95,12 +103,19 @@ class CustomListLocationsViewModelTest {
             // Check no selected
             val firstState = awaitItem()
             assertIs<CustomListLocationsUiState.Content.Data>(firstState)
-            assertEquals(emptySet<RelayItem>(), firstState.selectedLocations)
+            assertEquals(emptyList<RelayItem>(), firstState.selectedLocations())
+            // Expand country
+            viewModel.onExpand(DUMMY_COUNTRIES[0], true)
+            awaitItem()
+            // Expand city
+            viewModel.onExpand(DUMMY_COUNTRIES[0].cities[0], expand = true)
+            awaitItem()
+            // Select country
             viewModel.onRelaySelectionClick(DUMMY_COUNTRIES[0], true)
             // Check all items selected
             val secondState = awaitItem()
             assertIs<CustomListLocationsUiState.Content.Data>(secondState)
-            assertEquals(expectedSelection, secondState.selectedLocations)
+            assertLists(expectedSelection, secondState.selectedLocations())
         }
     }
 
@@ -108,25 +123,29 @@ class CustomListLocationsViewModelTest {
     fun `when deselecting child should deselect parent`() = runTest {
         // Arrange
         val expectedList = DUMMY_COUNTRIES
-        val initialSelection =
-            (DUMMY_COUNTRIES + DUMMY_COUNTRIES.flatMap { it.descendants() }).toSet()
+        val initialSelection = DUMMY_COUNTRIES.withDescendants()
+        val initialSelectionIds = initialSelection.map { it.id }
         val customListId = CustomListId("id")
-        val expectedSelection = emptySet<RelayItem>()
+        val expectedSelection = emptyList<RelayItem>()
         relayListFlow.value = expectedList
-        selectedLocationsFlow.value = initialSelection.toList()
+        selectedLocationsFlow.value = initialSelection
         val viewModel = createViewModel(customListId, true)
 
         // Act, Assert
         viewModel.uiState.test {
+            // Expand country
+            viewModel.onExpand(DUMMY_COUNTRIES[0], true)
+            // Expand city
+            viewModel.onExpand(DUMMY_COUNTRIES[0].cities[0], true)
             // Check initial selected
             val firstState = awaitItem()
             assertIs<CustomListLocationsUiState.Content.Data>(firstState)
-            assertEquals(initialSelection, firstState.selectedLocations)
+            assertEquals(initialSelectionIds, firstState.selectedLocations())
             viewModel.onRelaySelectionClick(DUMMY_COUNTRIES[0].cities[0].relays[0], false)
             // Check all items selected
             val secondState = awaitItem()
             assertIs<CustomListLocationsUiState.Content.Data>(secondState)
-            assertEquals(expectedSelection, secondState.selectedLocations)
+            assertEquals(expectedSelection, secondState.selectedLocations())
         }
     }
 
@@ -136,23 +155,28 @@ class CustomListLocationsViewModelTest {
         val expectedList = DUMMY_COUNTRIES
         val initialSelection =
             (DUMMY_COUNTRIES + DUMMY_COUNTRIES.flatMap { it.descendants() }).toSet()
+        val initialSelectionIds = initialSelection.map { it.id }
         val customListId = CustomListId("id")
-        val expectedSelection = emptySet<RelayItem>()
+        val expectedSelection = emptyList<RelayItem>()
         relayListFlow.value = expectedList
         selectedLocationsFlow.value = initialSelection.toList()
         val viewModel = createViewModel(customListId, true)
 
         // Act, Assert
         viewModel.uiState.test {
+            // Expand country
+            viewModel.onExpand(DUMMY_COUNTRIES[0], true)
+            // Expand city
+            viewModel.onExpand(DUMMY_COUNTRIES[0].cities[0], true)
             // Check initial selected
             val firstState = awaitItem()
             assertIs<CustomListLocationsUiState.Content.Data>(firstState)
-            assertEquals(initialSelection, firstState.selectedLocations)
+            assertEquals(initialSelectionIds, firstState.selectedLocations())
             viewModel.onRelaySelectionClick(DUMMY_COUNTRIES[0], false)
             // Check all items selected
             val secondState = awaitItem()
             assertIs<CustomListLocationsUiState.Content.Data>(secondState)
-            assertEquals(expectedSelection, secondState.selectedLocations)
+            assertEquals(expectedSelection, secondState.selectedLocations())
         }
     }
 
@@ -161,21 +185,27 @@ class CustomListLocationsViewModelTest {
         // Arrange
         val expectedList = DUMMY_COUNTRIES
         val customListId = CustomListId("id")
-        val expectedSelection = DUMMY_COUNTRIES[0].cities[0].relays.toSet()
+        val expectedSelection = DUMMY_COUNTRIES[0].cities[0].relays.map { it.id }
         val viewModel = createViewModel(customListId, true)
         relayListFlow.value = expectedList
 
         // Act, Assert
         viewModel.uiState.test {
+            awaitItem() // Initial item
+            // Expand country
+            viewModel.onExpand(DUMMY_COUNTRIES[0], true)
+            awaitItem()
+            // Expand city
+            viewModel.onExpand(DUMMY_COUNTRIES[0].cities[0], true)
             // Check no selected
             val firstState = awaitItem()
             assertIs<CustomListLocationsUiState.Content.Data>(firstState)
-            assertEquals(emptySet<RelayItem>(), firstState.selectedLocations)
+            assertEquals(emptyList<RelayItem>(), firstState.selectedLocations())
             viewModel.onRelaySelectionClick(DUMMY_COUNTRIES[0].cities[0].relays[0], true)
             // Check all items selected
             val secondState = awaitItem()
             assertIs<CustomListLocationsUiState.Content.Data>(secondState)
-            assertEquals(expectedSelection, secondState.selectedLocations)
+            assertEquals(expectedSelection, secondState.selectedLocations())
         }
     }
 
@@ -235,18 +265,26 @@ class CustomListLocationsViewModelTest {
         )
     }
 
+    private fun CustomListLocationsUiState.Content.Data.selectedLocations() =
+        this.locations.filter { it.checked }.map { it.item.id }
+
+    private fun RelayItem.Location.toDepth() =
+        when (this) {
+            is RelayItem.Location.Country -> 0
+            is RelayItem.Location.City -> 1
+            is RelayItem.Location.Relay -> 2
+        }
+
     companion object {
         private val DUMMY_COUNTRIES =
             listOf(
                 RelayItem.Location.Country(
                     name = "Sweden",
                     id = GeoLocationId.Country("SE"),
-                    expanded = false,
                     cities =
                         listOf(
                             RelayItem.Location.City(
                                 name = "Gothenburg",
-                                expanded = false,
                                 id = GeoLocationId.City(GeoLocationId.Country("SE"), "GBG"),
                                 relays =
                                     listOf(
diff --git a/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/SelectLocationViewModelTest.kt b/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/SelectLocationViewModelTest.kt
index 28f52ba261..d15e460e5d 100644
--- a/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/SelectLocationViewModelTest.kt
+++ b/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/SelectLocationViewModelTest.kt
@@ -11,15 +11,18 @@ import io.mockk.mockkStatic
 import io.mockk.unmockkAll
 import kotlin.test.assertEquals
 import kotlin.test.assertIs
+import kotlin.test.assertTrue
 import kotlinx.coroutines.cancel
 import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.test.runTest
 import net.mullvad.mullvadvpn.compose.communication.CustomListAction
 import net.mullvad.mullvadvpn.compose.communication.LocationsChanged
+import net.mullvad.mullvadvpn.compose.state.RelayListItem
 import net.mullvad.mullvadvpn.compose.state.SelectLocationUiState
 import net.mullvad.mullvadvpn.lib.common.test.TestCoroutineRule
 import net.mullvad.mullvadvpn.lib.common.test.assertLists
 import net.mullvad.mullvadvpn.lib.model.Constraint
+import net.mullvad.mullvadvpn.lib.model.CustomList
 import net.mullvad.mullvadvpn.lib.model.CustomListId
 import net.mullvad.mullvadvpn.lib.model.CustomListName
 import net.mullvad.mullvadvpn.lib.model.GeoLocationId
@@ -29,13 +32,14 @@ import net.mullvad.mullvadvpn.lib.model.Providers
 import net.mullvad.mullvadvpn.lib.model.RelayItem
 import net.mullvad.mullvadvpn.lib.model.RelayItemId
 import net.mullvad.mullvadvpn.relaylist.descendants
-import net.mullvad.mullvadvpn.relaylist.filterOnSearchTerm
+import net.mullvad.mullvadvpn.repository.CustomListsRepository
 import net.mullvad.mullvadvpn.repository.RelayListFilterRepository
 import net.mullvad.mullvadvpn.repository.RelayListRepository
 import net.mullvad.mullvadvpn.usecase.AvailableProvidersUseCase
 import net.mullvad.mullvadvpn.usecase.FilteredRelayListUseCase
 import net.mullvad.mullvadvpn.usecase.customlists.CustomListActionUseCase
 import net.mullvad.mullvadvpn.usecase.customlists.CustomListsRelayItemUseCase
+import net.mullvad.mullvadvpn.usecase.customlists.FilterCustomListsRelayItemUseCase
 import org.junit.jupiter.api.AfterEach
 import org.junit.jupiter.api.BeforeEach
 import org.junit.jupiter.api.Test
@@ -47,9 +51,11 @@ class SelectLocationViewModelTest {
     private val mockRelayListFilterRepository: RelayListFilterRepository = mockk()
     private val mockAvailableProvidersUseCase: AvailableProvidersUseCase = mockk(relaxed = true)
     private val mockCustomListActionUseCase: CustomListActionUseCase = mockk(relaxed = true)
-    private val mockCustomListsRelayItemUseCase: CustomListsRelayItemUseCase = mockk()
+    private val mockFilteredCustomListRelayItemsUseCase: FilterCustomListsRelayItemUseCase = mockk()
     private val mockFilteredRelayListUseCase: FilteredRelayListUseCase = mockk()
     private val mockRelayListRepository: RelayListRepository = mockk()
+    private val mockCustomListsRepository: CustomListsRepository = mockk()
+    private val mockCustomListsRelayItemUseCase: CustomListsRelayItemUseCase = mockk()
 
     private lateinit var viewModel: SelectLocationViewModel
 
@@ -58,7 +64,9 @@ class SelectLocationViewModelTest {
     private val selectedProviders = MutableStateFlow<Constraint<Providers>>(Constraint.Any)
     private val selectedRelayItemFlow = MutableStateFlow<Constraint<RelayItemId>>(Constraint.Any)
     private val filteredRelayList = MutableStateFlow<List<RelayItem.Location.Country>>(emptyList())
-    private val customRelayListItems = MutableStateFlow<List<RelayItem.CustomList>>(emptyList())
+    private val filteredCustomRelayListItems =
+        MutableStateFlow<List<RelayItem.CustomList>>(emptyList())
+    private val customListsRelayItem = MutableStateFlow<List<RelayItem.CustomList>>(emptyList())
 
     @BeforeEach
     fun setup() {
@@ -68,7 +76,8 @@ class SelectLocationViewModelTest {
         every { mockAvailableProvidersUseCase() } returns allProviders
         every { mockRelayListRepository.selectedLocation } returns selectedRelayItemFlow
         every { mockFilteredRelayListUseCase() } returns filteredRelayList
-        every { mockCustomListsRelayItemUseCase() } returns customRelayListItems
+        every { mockFilteredCustomListRelayItemsUseCase() } returns filteredCustomRelayListItems
+        every { mockCustomListsRelayItemUseCase() } returns customListsRelayItem
 
         mockkStatic(RELAY_LIST_EXTENSIONS)
         mockkStatic(RELAY_ITEM_EXTENSIONS)
@@ -77,10 +86,12 @@ class SelectLocationViewModelTest {
             SelectLocationViewModel(
                 relayListFilterRepository = mockRelayListFilterRepository,
                 availableProvidersUseCase = mockAvailableProvidersUseCase,
-                customListsRelayItemUseCase = mockCustomListsRelayItemUseCase,
+                filteredCustomListRelayItemsUseCase = mockFilteredCustomListRelayItemsUseCase,
                 customListActionUseCase = mockCustomListActionUseCase,
                 filteredRelayListUseCase = mockFilteredRelayListUseCase,
-                relayListRepository = mockRelayListRepository
+                relayListRepository = mockRelayListRepository,
+                customListsRepository = mockCustomListsRepository,
+                customListsRelayItemUseCase = mockCustomListsRelayItemUseCase,
             )
     }
 
@@ -96,41 +107,50 @@ class SelectLocationViewModelTest {
     }
 
     @Test
-    fun `given relayListWithSelection emits update uiState should contain new update`() = runTest {
+    fun `given filteredRelayList emits update uiState should contain new update`() = runTest {
         // Arrange
-        val mockCountries = listOf<RelayItem.Location.Country>(mockk(), mockk())
-        val selectedItem: RelayItemId = mockk()
-        every { mockCountries.filterOnSearchTerm(any(), selectedItem) } returns mockCountries
-        filteredRelayList.value = mockCountries
-        selectedRelayItemFlow.value = Constraint.Only(selectedItem)
+        filteredRelayList.value = testCountries
+        val selectedId = testCountries.first().id
+        selectedRelayItemFlow.value = Constraint.Only(selectedId)
 
         // Act, Assert
         viewModel.uiState.test {
             val actualState = awaitItem()
             assertIs<SelectLocationUiState.Content>(actualState)
-            assertLists(mockCountries, actualState.countries)
-            assertEquals(selectedItem, actualState.selectedItem)
+            assertLists(
+                testCountries.map { it.id },
+                actualState.relayListItems.mapNotNull { it.relayItemId() }
+            )
+            assertTrue(
+                actualState.relayListItems
+                    .filterIsInstance<RelayListItem.SelectableItem>()
+                    .first { it.relayItemId() == selectedId }
+                    .isSelected
+            )
         }
     }
 
     @Test
-    fun `given relayListWithSelection emits update with no selections selectedItem should be null`() =
-        runTest {
-            // Arrange
-            val mockCountries = listOf<RelayItem.Location.Country>(mockk(), mockk())
-            val selectedItem: RelayItemId? = null
-            every { mockCountries.filterOnSearchTerm(any(), selectedItem) } returns mockCountries
-            filteredRelayList.value = mockCountries
-            selectedRelayItemFlow.value = Constraint.Any
+    fun `given relay is selected all relay items should not be selected`() = runTest {
+        // Arrange
+        filteredRelayList.value = testCountries
+        selectedRelayItemFlow.value = Constraint.Any
 
-            // Act, Assert
-            viewModel.uiState.test {
-                val actualState = awaitItem()
-                assertIs<SelectLocationUiState.Content>(actualState)
-                assertLists(mockCountries, actualState.countries)
-                assertEquals(selectedItem, actualState.selectedItem)
-            }
+        // Act, Assert
+        viewModel.uiState.test {
+            val actualState = awaitItem()
+            assertIs<SelectLocationUiState.Content>(actualState)
+            assertLists(
+                testCountries.map { it.id },
+                actualState.relayListItems.mapNotNull { it.relayItemId() }
+            )
+            assertTrue(
+                actualState.relayListItems.filterIsInstance<RelayListItem.SelectableItem>().all {
+                    !it.isSelected
+                }
+            )
         }
+    }
 
     @Test
     fun `on selectRelay call uiSideEffect should emit CloseScreen and connect`() = runTest {
@@ -153,15 +173,8 @@ class SelectLocationViewModelTest {
     @Test
     fun `on onSearchTermInput call uiState should emit with filtered countries`() = runTest {
         // Arrange
-        val mockCustomList = listOf<RelayItem.CustomList>(mockk(relaxed = true))
-        val mockCountries = listOf<RelayItem.Location.Country>(mockk(), mockk())
-        val selectedItem: RelayItemId? = null
-        val mockRelayList: List<RelayItem.Location.Country> = mockk(relaxed = true)
-        val mockSearchString = "SEARCH"
-        every { mockRelayList.filterOnSearchTerm(mockSearchString, selectedItem) } returns
-            mockCountries
-        every { mockCustomList.filterOnSearchTerm(mockSearchString) } returns mockCustomList
-        filteredRelayList.value = mockRelayList
+        val mockSearchString = "got"
+        filteredRelayList.value = testCountries
         selectedRelayItemFlow.value = Constraint.Any
 
         // Act, Assert
@@ -172,25 +185,26 @@ class SelectLocationViewModelTest {
             // Update search string
             viewModel.onSearchTermInput(mockSearchString)
 
-            // Assert
+            // We get some unnecessary emissions for now
+            awaitItem()
+            awaitItem()
+            awaitItem()
+
             val actualState = awaitItem()
             assertIs<SelectLocationUiState.Content>(actualState)
-            assertLists(mockCountries, actualState.countries)
-            assertEquals(selectedItem, actualState.selectedItem)
+            assertTrue(
+                actualState.relayListItems.filterIsInstance<RelayListItem.GeoLocationItem>().any {
+                    it.item is RelayItem.Location.City && it.item.name == "Gothenburg"
+                }
+            )
         }
     }
 
     @Test
     fun `when onSearchTermInput returns empty result uiState should return empty list`() = runTest {
         // Arrange
-        val mockCustomList = listOf<RelayItem.CustomList>(mockk(relaxed = true))
-        val mockCountries = emptyList<RelayItem.Location.Country>()
-        val selectedItem: RelayItemId? = null
-        val mockRelayList: List<RelayItem.Location.Country> = mockk(relaxed = true)
+        filteredRelayList.value = testCountries
         val mockSearchString = "SEARCH"
-        every { mockRelayList.filterOnSearchTerm(mockSearchString, selectedItem) } returns
-            mockCountries
-        every { mockCustomList.filterOnSearchTerm(mockSearchString) } returns mockCustomList
 
         // Act, Assert
         viewModel.uiState.test {
@@ -200,10 +214,17 @@ class SelectLocationViewModelTest {
             // Update search string
             viewModel.onSearchTermInput(mockSearchString)
 
+            // We get some unnecessary emissions for now
+            awaitItem()
+            awaitItem()
+
             // Assert
             val actualState = awaitItem()
             assertIs<SelectLocationUiState.Content>(actualState)
-            assertEquals(mockSearchString, actualState.searchTerm)
+            assertEquals(
+                listOf(RelayListItem.LocationsEmptyText(mockSearchString)),
+                actualState.relayListItems
+            )
         }
     }
 
@@ -259,10 +280,13 @@ class SelectLocationViewModelTest {
         }
         val customList =
             RelayItem.CustomList(
-                id = CustomListId("1"),
-                customListName = CustomListName.fromString("custom"),
+                customList =
+                    CustomList(
+                        id = CustomListId("1"),
+                        name = CustomListName.fromString("custom"),
+                        locations = emptyList()
+                    ),
                 locations = emptyList(),
-                expanded = false
             )
         coEvery { mockCustomListActionUseCase(any<CustomListAction.UpdateLocations>()) } returns
             expectedResult.right()
@@ -276,6 +300,17 @@ class SelectLocationViewModelTest {
         }
     }
 
+    fun RelayListItem.relayItemId() =
+        when (this) {
+            is RelayListItem.CustomListFooter -> null
+            RelayListItem.CustomListHeader -> null
+            RelayListItem.LocationHeader -> null
+            is RelayListItem.LocationsEmptyText -> null
+            is RelayListItem.CustomListEntryItem -> item.id
+            is RelayListItem.CustomListItem -> item.id
+            is RelayListItem.GeoLocationItem -> item.id
+        }
+
     companion object {
         private const val RELAY_LIST_EXTENSIONS =
             "net.mullvad.mullvadvpn.relaylist.RelayListExtensionsKt"
@@ -283,5 +318,21 @@ class SelectLocationViewModelTest {
             "net.mullvad.mullvadvpn.relaylist.RelayItemExtensionsKt"
         private const val CUSTOM_LIST_EXTENSIONS =
             "net.mullvad.mullvadvpn.relaylist.CustomListExtensionsKt"
+
+        private val testCountries =
+            listOf<RelayItem.Location.Country>(
+                RelayItem.Location.Country(
+                    id = GeoLocationId.Country("se"),
+                    "Sweden",
+                    listOf(
+                        RelayItem.Location.City(
+                            id = GeoLocationId.City(GeoLocationId.Country("se"), "got"),
+                            "Gothenburg",
+                            emptyList()
+                        )
+                    )
+                ),
+                RelayItem.Location.Country(id = GeoLocationId.Country("no"), "Norway", emptyList())
+            )
     }
 }