Update the relaylist model

Use sealed interface for RelayItem
Add CustomList as a type of RelayItem

11 files changed, 142 insertions, 126 deletions

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
new file mode 100644
index 0000000000..9ad0c220e9
--- /dev/null
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/relaylist/CustomListExtensions.kt
@@ -0,0 +1,19 @@
+package net.mullvad.mullvadvpn.relaylist
+
+import net.mullvad.mullvadvpn.model.CustomList
+
+fun CustomList.toRelayItemCustomList(
+    relayCountries: List<RelayItem.Country>
+): RelayItem.CustomList =
+    RelayItem.CustomList(
+        code = this.id,
+        id = this.id,
+        name = this.name,
+        expanded = false,
+        locations =
+            this.locations.mapNotNull { relayCountries.findItemForGeographicLocationConstraint(it) }
+    )
+
+fun List<CustomList>.toRelayItemLists(
+    relayCountries: List<RelayItem.Country>
+): List<RelayItem.CustomList> = this.map { it.toRelayItemCustomList(relayCountries) }
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/relaylist/Relay.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/relaylist/Relay.kt
deleted file mode 100644
index 9bdb59168b..0000000000
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/relaylist/Relay.kt
+++ /dev/null
@@ -1,16 +0,0 @@
-package net.mullvad.mullvadvpn.relaylist
-
-import net.mullvad.mullvadvpn.model.GeographicLocationConstraint
-
-data class Relay(
-    override val name: String,
-    override val location: GeographicLocationConstraint,
-    override val locationName: String,
-    override val active: Boolean
-) : RelayItem {
-    override val code = name
-    override val type = RelayItemType.Relay
-    override val hasChildren = false
-
-    override val expanded = false
-}
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/relaylist/RelayCity.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/relaylist/RelayCity.kt
deleted file mode 100644
index e3693963f8..0000000000
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/relaylist/RelayCity.kt
+++ /dev/null
@@ -1,19 +0,0 @@
-package net.mullvad.mullvadvpn.relaylist
-
-import net.mullvad.mullvadvpn.model.GeographicLocationConstraint
-
-data class RelayCity(
-    override val name: String,
-    override val code: String,
-    override val location: GeographicLocationConstraint,
-    override val expanded: Boolean,
-    val relays: List<Relay>
-) : RelayItem {
-    override val type = RelayItemType.City
-
-    override val active
-        get() = relays.any { relay -> relay.active }
-
-    override val hasChildren
-        get() = relays.isNotEmpty()
-}
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/relaylist/RelayCountry.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/relaylist/RelayCountry.kt
deleted file mode 100644
index 3ad1d5962f..0000000000
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/relaylist/RelayCountry.kt
+++ /dev/null
@@ -1,19 +0,0 @@
-package net.mullvad.mullvadvpn.relaylist
-
-import net.mullvad.mullvadvpn.model.GeographicLocationConstraint
-
-data class RelayCountry(
-    override val name: String,
-    override val code: String,
-    override val expanded: Boolean,
-    val cities: List<RelayCity>
-) : RelayItem {
-    override val type = RelayItemType.Country
-    override val location = GeographicLocationConstraint.Country(code)
-
-    override val active
-        get() = cities.any { city -> city.active }
-
-    override val hasChildren
-        get() = cities.isNotEmpty()
-}
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/relaylist/RelayItem.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/relaylist/RelayItem.kt
index 3d38d6b068..03ba7a02f3 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/relaylist/RelayItem.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/relaylist/RelayItem.kt
@@ -1,12 +1,11 @@
 package net.mullvad.mullvadvpn.relaylist
 
+import net.mullvad.mullvadvpn.model.GeoIpLocation
 import net.mullvad.mullvadvpn.model.GeographicLocationConstraint
 
-interface RelayItem {
-    val type: RelayItemType
+sealed interface RelayItem {
     val name: String
     val code: String
-    val location: GeographicLocationConstraint
     val active: Boolean
     val hasChildren: Boolean
 
@@ -14,4 +13,67 @@ interface RelayItem {
         get() = name
 
     val expanded: Boolean
+
+    data class CustomList(
+        override val name: String,
+        override val code: String,
+        override val expanded: Boolean,
+        val id: String,
+        val locations: List<RelayItem>,
+    ) : RelayItem {
+        override val active
+            get() = locations.any { location -> location.active }
+
+        override val hasChildren
+            get() = locations.isNotEmpty()
+    }
+
+    data class Country(
+        override val name: String,
+        override val code: String,
+        override val expanded: Boolean,
+        val cities: List<City>
+    ) : RelayItem {
+        val location = GeographicLocationConstraint.Country(code)
+        override val active
+            get() = cities.any { city -> city.active }
+
+        override val hasChildren
+            get() = cities.isNotEmpty()
+    }
+
+    data class City(
+        override val name: String,
+        override val code: String,
+        override val expanded: Boolean,
+        val location: GeographicLocationConstraint.City,
+        val relays: List<Relay>
+    ) : RelayItem {
+
+        override val active
+            get() = relays.any { relay -> relay.active }
+
+        override val hasChildren
+            get() = relays.isNotEmpty()
+    }
+
+    data class Relay(
+        override val name: String,
+        override val locationName: String,
+        override val active: Boolean,
+        val location: GeographicLocationConstraint.Hostname,
+    ) : RelayItem {
+        override val code = name
+        override val hasChildren = false
+        override val expanded = false
+    }
+
+    fun location(): GeoIpLocation? {
+        return when (this) {
+            is CustomList -> null
+            is Country -> location.location
+            is City -> location.location
+            is Relay -> location.location
+        }
+    }
 }
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/relaylist/RelayItemExtensions.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/relaylist/RelayItemExtensions.kt
new file mode 100644
index 0000000000..6808d707b2
--- /dev/null
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/relaylist/RelayItemExtensions.kt
@@ -0,0 +1,12 @@
+package net.mullvad.mullvadvpn.relaylist
+
+import net.mullvad.mullvadvpn.model.LocationConstraint
+
+fun RelayItem.toLocationConstraint(): LocationConstraint {
+    return when (this) {
+        is RelayItem.Country -> LocationConstraint.Location(location)
+        is RelayItem.City -> LocationConstraint.Location(location)
+        is RelayItem.Relay -> LocationConstraint.Location(location)
+        is RelayItem.CustomList -> LocationConstraint.CustomList(id)
+    }
+}
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/relaylist/RelayList.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/relaylist/RelayList.kt
index c74c74ba43..30e4146245 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/relaylist/RelayList.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/relaylist/RelayList.kt
@@ -1,3 +1,7 @@
 package net.mullvad.mullvadvpn.relaylist
 
-data class RelayList(val country: List<RelayCountry>, val selectedItem: RelayItem?)
+data class RelayList(
+    val customLists: List<RelayItem.CustomList>,
+    val country: List<RelayItem.Country>,
+    val selectedItem: RelayItem?,
+)
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 6462ed9f41..06e00a022a 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
@@ -15,17 +15,17 @@ import net.mullvad.mullvadvpn.model.RelayList
 fun RelayList.toRelayCountries(
     ownership: Constraint<Ownership>,
     providers: Constraint<Providers>
-): List<RelayCountry> {
+): List<RelayItem.Country> {
     val relayCountries =
         this.countries
             .map { country ->
-                val cities = mutableListOf<RelayCity>()
-                val relayCountry = RelayCountry(country.name, country.code, false, cities)
+                val cities = mutableListOf<RelayItem.City>()
+                val relayCountry = RelayItem.Country(country.name, country.code, false, cities)
 
                 for (city in country.cities) {
-                    val relays = mutableListOf<Relay>()
+                    val relays = mutableListOf<RelayItem.Relay>()
                     val relayCity =
-                        RelayCity(
+                        RelayItem.City(
                             name = city.name,
                             code = city.code,
                             location = GeographicLocationConstraint.City(country.code, city.code),
@@ -38,7 +38,7 @@ fun RelayList.toRelayCountries(
 
                     for (relay in validCityRelays) {
                         relays.add(
-                            Relay(
+                            RelayItem.Relay(
                                 name = relay.hostname,
                                 location =
                                     GeographicLocationConstraint.Hostname(
@@ -69,32 +69,26 @@ fun RelayList.toRelayCountries(
     return relayCountries.toList()
 }
 
-fun List<RelayCountry>.findItemForLocation(
-    constraint: Constraint<GeographicLocationConstraint>
-): RelayItem? {
-    return when (constraint) {
-        is Constraint.Any -> null
-        is Constraint.Only -> {
-            when (val location = constraint.value) {
-                is GeographicLocationConstraint.Country -> {
-                    this.find { country -> country.code == location.countryCode }
-                }
-                is GeographicLocationConstraint.City -> {
-                    val country = this.find { country -> country.code == location.countryCode }
+fun List<RelayItem.Country>.findItemForGeographicLocationConstraint(
+    constraint: GeographicLocationConstraint
+) =
+    when (constraint) {
+        is GeographicLocationConstraint.Country -> {
+            this.find { country -> country.code == constraint.countryCode }
+        }
+        is GeographicLocationConstraint.City -> {
+            val country = this.find { country -> country.code == constraint.countryCode }
 
-                    country?.cities?.find { city -> city.code == location.cityCode }
-                }
-                is GeographicLocationConstraint.Hostname -> {
-                    val country = this.find { country -> country.code == location.countryCode }
+            country?.cities?.find { city -> city.code == constraint.cityCode }
+        }
+        is GeographicLocationConstraint.Hostname -> {
+            val country = this.find { country -> country.code == constraint.countryCode }
 
-                    val city = country?.cities?.find { city -> city.code == location.cityCode }
+            val city = country?.cities?.find { city -> city.code == constraint.cityCode }
 
-                    city?.relays?.find { relay -> relay.name == location.hostname }
-                }
-            }
+            city?.relays?.find { relay -> relay.name == constraint.hostname }
         }
     }
-}
 
 /**
  * Filter and expand the list based on search terms If a country is matched, that country and all
@@ -102,14 +96,15 @@ fun List<RelayCountry>.findItemForLocation(
  * 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.
  */
-fun List<RelayCountry>.filterOnSearchTerm(
+@Suppress("NestedBlockDepth")
+fun List<RelayItem.Country>.filterOnSearchTerm(
     searchTerm: String,
     selectedItem: RelayItem?
-): List<RelayCountry> {
+): List<RelayItem.Country> {
     return if (searchTerm.length >= MIN_SEARCH_LENGTH) {
-        val filteredCountries = mutableMapOf<String, RelayCountry>()
+        val filteredCountries = mutableMapOf<String, RelayItem.Country>()
         this.forEach { relayCountry ->
-            val cities = mutableListOf<RelayCity>()
+            val cities = mutableListOf<RelayItem.City>()
 
             // 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
@@ -122,7 +117,7 @@ fun List<RelayCountry>.filterOnSearchTerm(
 
             // Go through and try to match the search term with every city
             relayCountry.cities.forEach { relayCity ->
-                val relays = mutableListOf<Relay>()
+                val relays = mutableListOf<RelayItem.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.
@@ -200,31 +195,31 @@ private fun List<DaemonRelay>.filterValidRelays(
         }
 
 /** Expand the parent(s), if any, for the current selected item */
-private fun List<RelayCountry>.expandItemForSelection(
+private fun List<RelayItem.Country>.expandItemForSelection(
     selectedItem: RelayItem?
-): List<RelayCountry> {
+): List<RelayItem.Country> {
     return selectedItem?.let {
-        when (val location = selectedItem.location) {
-            is GeographicLocationConstraint.Country -> {
+        when (selectedItem) {
+            is RelayItem.Country -> {
                 this
             }
-            is GeographicLocationConstraint.City -> {
+            is RelayItem.City -> {
                 this.map { country ->
-                    if (country.code == location.countryCode) {
+                    if (country.code == selectedItem.location.countryCode) {
                         country.copy(expanded = true)
                     } else {
                         country
                     }
                 }
             }
-            is GeographicLocationConstraint.Hostname -> {
+            is RelayItem.Relay -> {
                 this.map { country ->
-                    if (country.code == location.countryCode) {
+                    if (country.code == selectedItem.location.countryCode) {
                         country.copy(
                             expanded = true,
                             cities =
                                 country.cities.map { city ->
-                                    if (city.code == location.cityCode) {
+                                    if (city.code == selectedItem.location.cityCode) {
                                         city.copy(expanded = true)
                                     } else {
                                         city
@@ -236,6 +231,7 @@ private fun List<RelayCountry>.expandItemForSelection(
                     }
                 }
             }
+            is RelayItem.CustomList -> this
         }
     } ?: this
 }
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/relaylist/RelayNameComparator.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/relaylist/RelayNameComparator.kt
index 92a3c9c1d6..c062fd1466 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/relaylist/RelayNameComparator.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/relaylist/RelayNameComparator.kt
@@ -1,7 +1,7 @@
 package net.mullvad.mullvadvpn.relaylist
 
-internal object RelayNameComparator : Comparator<Relay> {
-    override fun compare(o1: Relay, o2: Relay): Int {
+internal object RelayNameComparator : Comparator<RelayItem.Relay> {
+    override fun compare(o1: RelayItem.Relay, o2: RelayItem.Relay): Int {
         val partitions1 = o1.name.split(regex)
         val partitions2 = o2.name.split(regex)
         return if (partitions1.size > partitions2.size) partitions1 compareWith partitions2
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/serviceconnection/RelayListListener.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/serviceconnection/RelayListListener.kt
index 2d6ce332f6..841c9e0c59 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/serviceconnection/RelayListListener.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/serviceconnection/RelayListListener.kt
@@ -13,7 +13,7 @@ import net.mullvad.mullvadvpn.lib.ipc.MessageHandler
 import net.mullvad.mullvadvpn.lib.ipc.Request
 import net.mullvad.mullvadvpn.lib.ipc.events
 import net.mullvad.mullvadvpn.model.Constraint
-import net.mullvad.mullvadvpn.model.GeographicLocationConstraint
+import net.mullvad.mullvadvpn.model.LocationConstraint
 import net.mullvad.mullvadvpn.model.Ownership
 import net.mullvad.mullvadvpn.model.Providers
 import net.mullvad.mullvadvpn.model.RelayList
@@ -38,7 +38,7 @@ class RelayListListener(
                 defaultRelayList()
             )
 
-    fun updateSelectedRelayLocation(value: GeographicLocationConstraint) {
+    fun updateSelectedRelayLocation(value: LocationConstraint) {
         messageHandler.trySendRequest(Request.SetRelayLocation(value))
     }
 
diff --git a/android/lib/common/src/main/kotlin/net/mullvad/mullvadvpn/lib/common/util/LocationConstraintExtensions.kt b/android/lib/common/src/main/kotlin/net/mullvad/mullvadvpn/lib/common/util/LocationConstraintExtensions.kt
deleted file mode 100644
index d845e3aba9..0000000000
--- a/android/lib/common/src/main/kotlin/net/mullvad/mullvadvpn/lib/common/util/LocationConstraintExtensions.kt
+++ /dev/null
@@ -1,23 +0,0 @@
-package net.mullvad.mullvadvpn.lib.common.util
-
-import net.mullvad.mullvadvpn.model.Constraint
-import net.mullvad.mullvadvpn.model.GeographicLocationConstraint
-import net.mullvad.mullvadvpn.model.LocationConstraint
-
-fun LocationConstraint.toGeographicLocationConstraint(): GeographicLocationConstraint? =
-    when (this) {
-        is LocationConstraint.Location -> this.location
-        is LocationConstraint.CustomList -> null
-    }
-
-fun Constraint<LocationConstraint>.toGeographicLocationConstraint():
-    Constraint<GeographicLocationConstraint> =
-    when (this) {
-        is Constraint.Only ->
-            when (value) {
-                is LocationConstraint.Location ->
-                    Constraint.Only((value as LocationConstraint.Location).location)
-                is LocationConstraint.CustomList -> Constraint.Any()
-            }
-        is Constraint.Any -> Constraint.Any()
-    }