Convert select location into flat LazyColumn

22 files changed, 856 insertions, 720 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/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/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/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/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/lib/daemon-grpc/src/main/kotlin/net/mullvad/mullvadvpn/lib/daemon/grpc/mapper/FromDomain.kt b/android/lib/daemon-grpc/src/main/kotlin/net/mullvad/mullvadvpn/lib/daemon/grpc/mapper/FromDomain.kt
index 014bafb85b..4359ddff93 100644
--- a/android/lib/daemon-grpc/src/main/kotlin/net/mullvad/mullvadvpn/lib/daemon/grpc/mapper/FromDomain.kt
+++ b/android/lib/daemon-grpc/src/main/kotlin/net/mullvad/mullvadvpn/lib/daemon/grpc/mapper/FromDomain.kt
@@ -106,12 +106,10 @@ internal fun GeoLocationId.fromDomain(): ManagementInterface.GeographicLocationC
     ManagementInterface.GeographicLocationConstraint.newBuilder()
         .apply {
             when (val id = this@fromDomain) {
-                is GeoLocationId.Country -> setCountry(id.countryCode)
-                is GeoLocationId.City -> setCountry(id.countryCode.countryCode).setCity(id.cityCode)
+                is GeoLocationId.Country -> setCountry(id.code)
+                is GeoLocationId.City -> setCountry(id.country.code).setCity(id.code)
                 is GeoLocationId.Hostname ->
-                    setCountry(id.country.countryCode)
-                        .setCity(id.city.cityCode)
-                        .setHostname(id.hostname)
+                    setCountry(id.country.code).setCity(id.city.code).setHostname(id.code)
             }
         }
         .build()
diff --git a/android/lib/daemon-grpc/src/main/kotlin/net/mullvad/mullvadvpn/lib/daemon/grpc/mapper/ToDomain.kt b/android/lib/daemon-grpc/src/main/kotlin/net/mullvad/mullvadvpn/lib/daemon/grpc/mapper/ToDomain.kt
index 13ebe74350..50599ad7f4 100644
--- a/android/lib/daemon-grpc/src/main/kotlin/net/mullvad/mullvadvpn/lib/daemon/grpc/mapper/ToDomain.kt
+++ b/android/lib/daemon-grpc/src/main/kotlin/net/mullvad/mullvadvpn/lib/daemon/grpc/mapper/ToDomain.kt
@@ -455,7 +455,6 @@ internal fun ManagementInterface.RelayListCountry.toDomain(): RelayItem.Location
     return RelayItem.Location.Country(
         countryCode,
         name,
-        false,
         citiesList
             .map { city -> city.toDomain(countryCode) }
             .filter { it.relays.isNotEmpty() }
@@ -470,7 +469,6 @@ internal fun ManagementInterface.RelayListCity.toDomain(
     return RelayItem.Location.City(
         name = name,
         id = cityCode,
-        expanded = false,
         relays =
             relaysList
                 .filter { it.endpointType == ManagementInterface.Relay.RelayType.WIREGUARD }
diff --git a/android/lib/model/src/main/kotlin/net/mullvad/mullvadvpn/lib/model/RelayItem.kt b/android/lib/model/src/main/kotlin/net/mullvad/mullvadvpn/lib/model/RelayItem.kt
index a31a6f67df..17bc563a8d 100644
--- a/android/lib/model/src/main/kotlin/net/mullvad/mullvadvpn/lib/model/RelayItem.kt
+++ b/android/lib/model/src/main/kotlin/net/mullvad/mullvadvpn/lib/model/RelayItem.kt
@@ -2,28 +2,26 @@ package net.mullvad.mullvadvpn.lib.model
 
 import arrow.optics.optics
 
+typealias DomainCustomList = CustomList
+
 @optics
 sealed interface RelayItem {
     val id: RelayItemId
     val name: String
+
     val active: Boolean
     val hasChildren: Boolean
-    val expanded: Boolean
 
     @optics
     data class CustomList(
-        override val id: CustomListId,
-        val customListName: CustomListName,
+        val customList: DomainCustomList,
         val locations: List<Location>,
-        override val expanded: Boolean,
     ) : RelayItem {
-        override val name: String = customListName.value
+        override val name: String = customList.name.value
+        override val id = customList.id
 
-        override val active
-            get() = locations.any { location -> location.active }
-
-        override val hasChildren
-            get() = locations.isNotEmpty()
+        override val active = locations.any { it.active }
+        override val hasChildren: Boolean = locations.isNotEmpty()
 
         companion object
     }
@@ -36,16 +34,11 @@ sealed interface RelayItem {
         data class Country(
             override val id: GeoLocationId.Country,
             override val name: String,
-            override val expanded: Boolean,
             val cities: List<City>
         ) : Location {
             val relays = cities.flatMap { city -> city.relays }
-
-            override val active
-                get() = cities.any { city -> city.active }
-
-            override val hasChildren
-                get() = cities.isNotEmpty()
+            override val active = cities.any { it.active }
+            override val hasChildren: Boolean = cities.isNotEmpty()
 
             companion object
         }
@@ -54,15 +47,10 @@ sealed interface RelayItem {
         data class City(
             override val id: GeoLocationId.City,
             override val name: String,
-            override val expanded: Boolean,
             val relays: List<Relay>
         ) : Location {
-
-            override val active
-                get() = relays.any { relay -> relay.active }
-
-            override val hasChildren
-                get() = relays.isNotEmpty()
+            override val active = relays.any { it.active }
+            override val hasChildren: Boolean = relays.isNotEmpty()
 
             companion object
         }
@@ -73,10 +61,8 @@ sealed interface RelayItem {
             val provider: Provider,
             override val active: Boolean,
         ) : Location {
-            override val name: String = id.hostname
-
-            override val hasChildren = false
-            override val expanded = false
+            override val name: String = id.code
+            override val hasChildren: Boolean = false
 
             companion object
         }
diff --git a/android/lib/model/src/main/kotlin/net/mullvad/mullvadvpn/lib/model/RelayItemId.kt b/android/lib/model/src/main/kotlin/net/mullvad/mullvadvpn/lib/model/RelayItemId.kt
index c560fab49c..85326d8d2a 100644
--- a/android/lib/model/src/main/kotlin/net/mullvad/mullvadvpn/lib/model/RelayItemId.kt
+++ b/android/lib/model/src/main/kotlin/net/mullvad/mullvadvpn/lib/model/RelayItemId.kt
@@ -21,28 +21,29 @@ value class CustomListId(val value: String) : RelayItemId, Parcelable {
 sealed interface GeoLocationId : RelayItemId, Parcelable {
     @optics
     @Parcelize
-    data class Country(val countryCode: String) : GeoLocationId {
+    data class Country(override val code: String) : GeoLocationId {
         companion object
     }
 
     @optics
     @Parcelize
-    data class City(val countryCode: Country, val cityCode: String) : GeoLocationId {
+    data class City(override val country: Country, override val code: String) : GeoLocationId {
         companion object
     }
 
     @optics
     @Parcelize
-    data class Hostname(val city: City, val hostname: String) : GeoLocationId {
+    data class Hostname(val city: City, override val code: String) : GeoLocationId {
         companion object
     }
 
+    val code: String
     val country: Country
         get() =
             when (this) {
                 is Country -> this
-                is City -> this.countryCode
-                is Hostname -> this.city.countryCode
+                is City -> this.country
+                is Hostname -> this.city.country
             }
 
     companion object
diff --git a/android/lib/theme/src/main/kotlin/net/mullvad/mullvadvpn/lib/theme/dimensions/Dimensions.kt b/android/lib/theme/src/main/kotlin/net/mullvad/mullvadvpn/lib/theme/dimensions/Dimensions.kt
index d4e5e4803d..922e97073b 100644
--- a/android/lib/theme/src/main/kotlin/net/mullvad/mullvadvpn/lib/theme/dimensions/Dimensions.kt
+++ b/android/lib/theme/src/main/kotlin/net/mullvad/mullvadvpn/lib/theme/dimensions/Dimensions.kt
@@ -18,7 +18,7 @@ data class Dimensions(
     val cellHeight: Dp = 56.dp,
     val cellHeightTwoRows: Dp = 72.dp,
     val cellLabelVerticalPadding: Dp = 14.dp,
-    val cellStartPadding: Dp = 22.dp,
+    val cellStartPadding: Dp = 14.dp,
     val cellStartPaddingInteractive: Dp = 14.dp,
     val cellTopPadding: Dp = 6.dp,
     val cellVerticalSpacing: Dp = 14.dp,
@@ -41,7 +41,7 @@ data class Dimensions(
     val dropdownMenuVerticalPadding: Dp = 8.dp, // Used to remove padding from dropdown menu
     val dropdownMenuBorder: Dp = 1.dp,
     val expandableCellChevronSize: Dp = 30.dp,
-    val filterTittlePadding: Dp = 4.dp,
+    val filterTitlePadding: Dp = 4.dp,
     val formTextFieldMinHeight: Dp = 72.dp,
     val iconFailSuccessTopMargin: Dp = 30.dp,
     val iconHeight: Dp = 44.dp,