Remove one-to-one relationship between provider and ownership

12 files changed, 127 insertions, 98 deletions

diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/cell/CheckboxCell.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/cell/CheckboxCell.kt
index 13c237d6f4..e1b0c1a5d5 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/cell/CheckboxCell.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/cell/CheckboxCell.kt
@@ -31,6 +31,7 @@ internal fun CheckboxCell(
     modifier: Modifier = Modifier,
     title: String,
     checked: Boolean,
+    enabled: Boolean = true,
     onCheckedChange: (Boolean) -> Unit,
     background: Color = MaterialTheme.colorScheme.surfaceContainerLow,
     startPadding: Dp = Dimens.mediumPadding,
@@ -41,7 +42,7 @@ internal fun CheckboxCell(
         verticalAlignment = Alignment.CenterVertically,
         modifier =
             modifier
-                .clickable { onCheckedChange(!checked) }
+                .clickable(enabled) { onCheckedChange(!checked) }
                 .defaultMinSize(minHeight = minHeight)
                 .fillMaxWidth()
                 .background(background)
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/preview/FilterUiStatePreviewParameterProvider.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/preview/FilterUiStatePreviewParameterProvider.kt
index 8728143067..0aa3be8158 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/preview/FilterUiStatePreviewParameterProvider.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/preview/FilterUiStatePreviewParameterProvider.kt
@@ -3,19 +3,17 @@ package net.mullvad.mullvadvpn.compose.preview
 import androidx.compose.ui.tooling.preview.PreviewParameterProvider
 import net.mullvad.mullvadvpn.compose.state.RelayFilterUiState
 import net.mullvad.mullvadvpn.lib.model.Ownership
-import net.mullvad.mullvadvpn.lib.model.Provider
 import net.mullvad.mullvadvpn.lib.model.ProviderId
 
-private val PROVIDER =
-    Provider(providerId = ProviderId("provider1"), ownership = Ownership.MullvadOwned)
+private val PROVIDER_TO_OWNERSHIPS = mapOf(ProviderId("provider1") to setOf(Ownership.MullvadOwned))
 
 class FilterUiStatePreviewParameterProvider : PreviewParameterProvider<RelayFilterUiState> {
     override val values =
         sequenceOf(
             RelayFilterUiState(
+                providerToOwnerships = PROVIDER_TO_OWNERSHIPS,
                 selectedOwnership = Ownership.MullvadOwned,
-                allProviders = listOf(PROVIDER),
-                selectedProviders = listOf(PROVIDER),
+                selectedProviders = PROVIDER_TO_OWNERSHIPS.keys.toList(),
             )
         )
 }
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 0af4199f40..20f69e3bcd 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
@@ -2,7 +2,6 @@ package net.mullvad.mullvadvpn.compose.preview
 
 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.ProviderId
 import net.mullvad.mullvadvpn.lib.model.RelayItem
 
@@ -56,7 +55,8 @@ private fun generateRelayItemRelay(
     RelayItem.Location.Relay(
         id = GeoLocationId.Hostname(city = cityCode, code = hostName),
         active = active,
-        provider = Provider(ProviderId("Provider"), Ownership.MullvadOwned),
+        provider = ProviderId("Provider"),
+        ownership = Ownership.MullvadOwned,
         daita = daita,
     )
 
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/preview/SelectLocationsUiStatePreviewParameterProvider.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/preview/SelectLocationsUiStatePreviewParameterProvider.kt
index 6893397f9f..965999c7f1 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/preview/SelectLocationsUiStatePreviewParameterProvider.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/preview/SelectLocationsUiStatePreviewParameterProvider.kt
@@ -3,8 +3,6 @@ package net.mullvad.mullvadvpn.compose.preview
 import androidx.compose.ui.tooling.preview.PreviewParameterProvider
 import net.mullvad.mullvadvpn.compose.state.RelayListType
 import net.mullvad.mullvadvpn.compose.state.SelectLocationUiState
-import net.mullvad.mullvadvpn.lib.model.Ownership
-import net.mullvad.mullvadvpn.lib.model.Provider
 import net.mullvad.mullvadvpn.usecase.FilterChip
 import net.mullvad.mullvadvpn.usecase.ModelOwnership
 
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/screen/FilterScreen.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/screen/FilterScreen.kt
index 60788bffd5..3d3f937c73 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/screen/FilterScreen.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/screen/FilterScreen.kt
@@ -46,7 +46,7 @@ import net.mullvad.mullvadvpn.compose.state.RelayFilterUiState
 import net.mullvad.mullvadvpn.compose.transitions.SlideInFromRightTransition
 import net.mullvad.mullvadvpn.compose.util.CollectSideEffectWithLifecycle
 import net.mullvad.mullvadvpn.lib.model.Ownership
-import net.mullvad.mullvadvpn.lib.model.Provider
+import net.mullvad.mullvadvpn.lib.model.ProviderId
 import net.mullvad.mullvadvpn.lib.theme.AppTheme
 import net.mullvad.mullvadvpn.lib.theme.Dimens
 import net.mullvad.mullvadvpn.viewmodel.FilterScreenSideEffect
@@ -98,7 +98,7 @@ fun FilterScreen(
     onApplyClick: () -> Unit,
     onSelectedOwnership: (ownership: Ownership?) -> Unit,
     onAllProviderCheckChange: (isChecked: Boolean) -> Unit,
-    onSelectedProvider: (checked: Boolean, provider: Provider) -> Unit,
+    onSelectedProvider: (checked: Boolean, provider: ProviderId) -> Unit,
 ) {
     var providerExpanded by rememberSaveable { mutableStateOf(false) }
     var ownershipExpanded by rememberSaveable { mutableStateOf(false) }
@@ -126,7 +126,7 @@ fun FilterScreen(
                 itemsWithDivider(
                     key = { it.name },
                     contentType = { ContentType.ITEM },
-                    items = state.filteredOwnershipByProviders,
+                    items = state.selectableOwnerships,
                 ) { ownership ->
                     Ownership(ownership, state, onSelectedOwnership)
                 }
@@ -139,9 +139,17 @@ fun FilterScreen(
                     AllProviders(state, onAllProviderCheckChange)
                 }
                 itemsWithDivider(
-                    key = { it.providerId.value },
+                    key = { it.value },
                     contentType = { ContentType.ITEM },
-                    items = state.filteredProvidersByOwnership,
+                    items = state.removedProviders,
+                ) { provider ->
+                    RemovedProvider(provider, state, onSelectedProvider)
+                }
+
+                itemsWithDivider(
+                    key = { it.value },
+                    contentType = { ContentType.ITEM },
+                    items = state.selectableProviders,
                 ) { provider ->
                     Provider(provider, state, onSelectedProvider)
                 }
@@ -216,14 +224,30 @@ private fun LazyItemScope.AllProviders(
 
 @Composable
 private fun LazyItemScope.Provider(
-    provider: Provider,
+    providerId: ProviderId,
+    state: RelayFilterUiState,
+    onSelectedProvider: (checked: Boolean, providerId: ProviderId) -> Unit,
+) {
+    CheckboxCell(
+        title = providerId.value,
+        checked = providerId in state.selectedProviders,
+        onCheckedChange = { checked -> onSelectedProvider(checked, providerId) },
+        modifier = Modifier.animateItem(),
+    )
+}
+
+@Composable
+private fun LazyItemScope.RemovedProvider(
+    providerId: ProviderId,
     state: RelayFilterUiState,
-    onSelectedProvider: (checked: Boolean, provider: Provider) -> Unit,
+    onSelectedProvider: (checked: Boolean, providerId: ProviderId) -> Unit,
 ) {
+    val checked = providerId in state.selectedProviders
     CheckboxCell(
-        title = provider.providerId.value,
-        checked = provider in state.selectedProviders,
-        onCheckedChange = { checked -> onSelectedProvider(checked, provider) },
+        title = stringResource(R.string.removed_provider, providerId.value),
+        checked = checked,
+        enabled = checked,
+        onCheckedChange = { checked -> onSelectedProvider(checked, providerId) },
         modifier = Modifier.animateItem(),
     )
 }
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 01fe84f76c..2ddf35fad5 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
@@ -2,7 +2,7 @@ package net.mullvad.mullvadvpn.compose.state
 
 import net.mullvad.mullvadvpn.lib.model.Constraint
 import net.mullvad.mullvadvpn.lib.model.Ownership
-import net.mullvad.mullvadvpn.lib.model.Provider
+import net.mullvad.mullvadvpn.lib.model.ProviderId
 import net.mullvad.mullvadvpn.lib.model.Providers
 
 fun Ownership?.toOwnershipConstraint(): Constraint<Ownership> =
@@ -11,18 +11,15 @@ fun Ownership?.toOwnershipConstraint(): Constraint<Ownership> =
         else -> Constraint.Only(this)
     }
 
-fun Constraint<Providers>.toSelectedProviders(allProviders: List<Provider>): List<Provider> =
+fun Constraint<Providers>.toSelectedProviders(allProviders: List<ProviderId>): List<ProviderId> =
     when (this) {
         Constraint.Any -> allProviders
-        is Constraint.Only ->
-            value.providers.toList().mapNotNull { provider ->
-                allProviders.firstOrNull { it.providerId == provider }
-            }
+        is Constraint.Only -> value.providers.toList()
     }
 
-fun List<Provider>.toConstraintProviders(allProviders: List<Provider>): Constraint<Providers> =
+fun List<ProviderId>.toConstraintProviders(allProviders: List<ProviderId>): Constraint<Providers> =
     if (size == allProviders.size) {
         Constraint.Any
     } else {
-        Constraint.Only(Providers(map { provider -> provider.providerId }.toHashSet()))
+        Constraint.Only(Providers(toHashSet()))
     }
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/state/RelayFilterUiState.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/state/RelayFilterUiState.kt
index ffc26b0bb7..0c81ac0aa5 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/state/RelayFilterUiState.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/state/RelayFilterUiState.kt
@@ -1,26 +1,37 @@
 package net.mullvad.mullvadvpn.compose.state
 
 import net.mullvad.mullvadvpn.lib.model.Ownership
-import net.mullvad.mullvadvpn.lib.model.Provider
+import net.mullvad.mullvadvpn.lib.model.ProviderId
 
 data class RelayFilterUiState(
+    private val providerToOwnerships: Map<ProviderId, Set<Ownership>> = emptyMap(),
     val selectedOwnership: Ownership? = null,
-    val allProviders: List<Provider> = emptyList(),
-    val selectedProviders: List<Provider> = allProviders,
+    val selectedProviders: List<ProviderId> = emptyList(),
 ) {
-    val isApplyButtonEnabled = selectedProviders.isNotEmpty()
+    val allProviders: List<ProviderId> = providerToOwnerships.keys.toList().sorted()
 
-    val filteredOwnershipByProviders =
+    val selectableOwnerships: List<Ownership> =
         if (selectedProviders.isEmpty()) {
-            Ownership.entries
-        } else {
-            Ownership.entries.filter { ownership ->
-                selectedProviders.any { provider -> provider.ownership == ownership }
+                Ownership.entries
+            } else {
+                providerToOwnerships
+                    .filterKeys { it in selectedProviders }
+                    .values
+                    .flatten()
+                    .distinct()
             }
+            .sorted()
+
+    val selectableProviders: List<ProviderId> =
+        if (selectedOwnership != null) {
+            providerToOwnerships.filterValues { selectedOwnership in it }.keys.toList().sorted()
+        } else {
+            allProviders
         }
-    val filteredProvidersByOwnership =
-        if (selectedOwnership == null) allProviders
-        else allProviders.filter { provider -> provider.ownership == selectedOwnership }
+
+    val removedProviders: List<ProviderId> = selectedProviders - allProviders
+
+    val isApplyButtonEnabled = selectedProviders.isNotEmpty()
 
     val isAllProvidersChecked = allProviders.size == selectedProviders.size
 }
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 650ee67eaa..4acf52c7b0 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
@@ -37,7 +37,6 @@ import net.mullvad.mullvadvpn.ui.MainActivity
 import net.mullvad.mullvadvpn.ui.serviceconnection.AppVersionInfoRepository
 import net.mullvad.mullvadvpn.ui.serviceconnection.ServiceConnectionManager
 import net.mullvad.mullvadvpn.usecase.AccountExpiryInAppNotificationUseCase
-import net.mullvad.mullvadvpn.usecase.AvailableProvidersUseCase
 import net.mullvad.mullvadvpn.usecase.EmptyPaymentUseCase
 import net.mullvad.mullvadvpn.usecase.FilterChipUseCase
 import net.mullvad.mullvadvpn.usecase.FilteredRelayListUseCase
@@ -47,6 +46,7 @@ import net.mullvad.mullvadvpn.usecase.NewDeviceNotificationUseCase
 import net.mullvad.mullvadvpn.usecase.OutOfTimeUseCase
 import net.mullvad.mullvadvpn.usecase.PaymentUseCase
 import net.mullvad.mullvadvpn.usecase.PlayPaymentUseCase
+import net.mullvad.mullvadvpn.usecase.ProviderToOwnershipsUseCase
 import net.mullvad.mullvadvpn.usecase.SelectedLocationTitleUseCase
 import net.mullvad.mullvadvpn.usecase.SelectedLocationUseCase
 import net.mullvad.mullvadvpn.usecase.SystemVpnSettingsAvailableUseCase
@@ -156,7 +156,7 @@ val uiModule = module {
     single { SystemVpnSettingsAvailableUseCase(androidContext()) }
     single { CustomListActionUseCase(get(), get()) }
     single { SelectedLocationTitleUseCase(get(), get()) }
-    single { AvailableProvidersUseCase(get()) }
+    single { ProviderToOwnershipsUseCase(get()) }
     single { FilterCustomListsRelayItemUseCase(get(), get(), get(), get()) }
     single { CustomListsRelayItemUseCase(get(), get()) }
     single { CustomListRelayItemsUseCase(get(), get()) }
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
index 5584d8e991..21d3294de4 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/relaylist/RelayItemExtensions.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/relaylist/RelayItemExtensions.kt
@@ -35,7 +35,7 @@ private fun RelayItem.Location.hasOwnership(ownershipConstraint: Constraint<Owne
         when (this) {
             is RelayItem.Location.Country -> cities.any { it.hasOwnership(ownershipConstraint) }
             is RelayItem.Location.City -> relays.any { it.hasOwnership(ownershipConstraint) }
-            is RelayItem.Location.Relay -> this.provider.ownership == ownershipConstraint.value
+            is RelayItem.Location.Relay -> ownershipConstraint.value == ownership
         }
     } else {
         true
@@ -46,8 +46,7 @@ private fun RelayItem.Location.hasProvider(providersConstraint: Constraint<Provi
         when (this) {
             is RelayItem.Location.Country -> cities.any { it.hasProvider(providersConstraint) }
             is RelayItem.Location.City -> relays.any { it.hasProvider(providersConstraint) }
-            is RelayItem.Location.Relay ->
-                providersConstraint.value.providers.contains(provider.providerId)
+            is RelayItem.Location.Relay -> provider in providersConstraint.value.providers
         }
     } else {
         true
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/usecase/FilterChipUseCase.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/usecase/FilterChipUseCase.kt
index 37e5f71ecc..361e29944e 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/usecase/FilterChipUseCase.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/usecase/FilterChipUseCase.kt
@@ -3,10 +3,9 @@ package net.mullvad.mullvadvpn.usecase
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.combine
 import net.mullvad.mullvadvpn.compose.state.RelayListType
-import net.mullvad.mullvadvpn.compose.state.toSelectedProviders
 import net.mullvad.mullvadvpn.lib.model.Constraint
 import net.mullvad.mullvadvpn.lib.model.Ownership
-import net.mullvad.mullvadvpn.lib.model.Provider
+import net.mullvad.mullvadvpn.lib.model.ProviderId
 import net.mullvad.mullvadvpn.lib.model.Providers
 import net.mullvad.mullvadvpn.lib.model.Settings
 import net.mullvad.mullvadvpn.repository.RelayListFilterRepository
@@ -18,7 +17,7 @@ typealias ModelOwnership = Ownership
 
 class FilterChipUseCase(
     private val relayListFilterRepository: RelayListFilterRepository,
-    private val availableProvidersUseCase: AvailableProvidersUseCase,
+    private val providerToOwnershipsUseCase: ProviderToOwnershipsUseCase,
     private val settingsRepository: SettingsRepository,
     private val wireguardConstraintsRepository: WireguardConstraintsRepository,
 ) {
@@ -26,19 +25,19 @@ class FilterChipUseCase(
         combine(
             relayListFilterRepository.selectedOwnership,
             relayListFilterRepository.selectedProviders,
-            availableProvidersUseCase(),
+            providerToOwnershipsUseCase(),
             settingsRepository.settingsUpdates,
             wireguardConstraintsRepository.wireguardConstraints,
         ) {
             selectedOwnership,
             selectedConstraintProviders,
-            allProviders,
+            providerOwnership,
             settings,
             wireguardConstraints ->
             filterChips(
                 selectedOwnership = selectedOwnership,
                 selectedConstraintProviders = selectedConstraintProviders,
-                allProviders = allProviders,
+                providerToOwnerships = providerOwnership,
                 daitaDirectOnly = settings?.daitaAndDirectOnly() == true,
                 isMultihopEnabled = wireguardConstraints?.isMultihopEnabled == true,
                 relayListType = relayListType,
@@ -48,7 +47,7 @@ class FilterChipUseCase(
     private fun filterChips(
         selectedOwnership: Constraint<Ownership>,
         selectedConstraintProviders: Constraint<Providers>,
-        allProviders: List<Provider>,
+        providerToOwnerships: Map<ProviderId, Set<Ownership>>,
         daitaDirectOnly: Boolean,
         isMultihopEnabled: Boolean,
         relayListType: RelayListType,
@@ -58,10 +57,22 @@ class FilterChipUseCase(
             when (selectedConstraintProviders) {
                 is Constraint.Any -> null
                 is Constraint.Only ->
-                    filterSelectedProvidersByOwnership(
-                            selectedConstraintProviders.toSelectedProviders(allProviders),
-                            ownershipFilter,
-                        )
+                    selectedConstraintProviders.value.providers
+                        .filter { providerId ->
+                            if (ownershipFilter == null) {
+                                true
+                            } else {
+                                val providerOwnerships = providerToOwnerships[providerId]
+                                // If the provider has been removed from the relay list we add it
+                                // so it is visible for the user, because we won't know what
+                                // ownerships it had.
+                                if (providerOwnerships == null) {
+                                    true
+                                } else {
+                                    providerOwnerships.contains(ownershipFilter)
+                                }
+                            }
+                        }
                         .size
             }
         return buildList {
@@ -83,13 +94,6 @@ class FilterChipUseCase(
         }
     }
 
-    private fun filterSelectedProvidersByOwnership(
-        selectedProviders: List<Provider>,
-        selectedOwnership: Ownership?,
-    ): List<Provider> =
-        if (selectedOwnership == null) selectedProviders
-        else selectedProviders.filter { it.ownership == selectedOwnership }
-
     private fun Settings.daitaAndDirectOnly() =
         tunnelOptions.wireguard.daitaSettings.enabled &&
             tunnelOptions.wireguard.daitaSettings.directOnly
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/usecase/AvailableProvidersUseCase.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/usecase/ProviderToOwnershipsUseCase.kt
index 14aa3824cf..42f5a94e29 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/usecase/AvailableProvidersUseCase.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/usecase/ProviderToOwnershipsUseCase.kt
@@ -2,18 +2,18 @@ package net.mullvad.mullvadvpn.usecase
 
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.map
-import net.mullvad.mullvadvpn.lib.model.Provider
+import net.mullvad.mullvadvpn.lib.model.Ownership
+import net.mullvad.mullvadvpn.lib.model.ProviderId
 import net.mullvad.mullvadvpn.lib.model.RelayItem
 import net.mullvad.mullvadvpn.repository.RelayListRepository
 
-class AvailableProvidersUseCase(private val relayListRepository: RelayListRepository) {
-
-    operator fun invoke(): Flow<List<Provider>> =
+class ProviderToOwnershipsUseCase(private val relayListRepository: RelayListRepository) {
+    operator fun invoke(): Flow<Map<ProviderId, Set<Ownership>>> =
         relayListRepository.relayList.map { relayList ->
             relayList
                 .flatMap(RelayItem.Location.Country::cities)
                 .flatMap(RelayItem.Location.City::relays)
-                .map(RelayItem.Location.Relay::provider)
-                .distinct()
+                .groupBy({ it.provider }, { it.ownership })
+                .mapValues { (_, ownerships) -> ownerships.toSet() }
         }
 }
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 548d6b70a7..c05515126e 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
@@ -16,27 +16,30 @@ import net.mullvad.mullvadvpn.compose.state.toConstraintProviders
 import net.mullvad.mullvadvpn.compose.state.toOwnershipConstraint
 import net.mullvad.mullvadvpn.compose.state.toSelectedProviders
 import net.mullvad.mullvadvpn.lib.model.Ownership
-import net.mullvad.mullvadvpn.lib.model.Provider
+import net.mullvad.mullvadvpn.lib.model.ProviderId
 import net.mullvad.mullvadvpn.repository.RelayListFilterRepository
-import net.mullvad.mullvadvpn.usecase.AvailableProvidersUseCase
+import net.mullvad.mullvadvpn.usecase.ProviderToOwnershipsUseCase
 
 class FilterViewModel(
-    private val availableProvidersUseCase: AvailableProvidersUseCase,
+    private val providerToOwnershipsUseCase: ProviderToOwnershipsUseCase,
     private val relayListFilterRepository: RelayListFilterRepository,
 ) : ViewModel() {
     private val _uiSideEffect = Channel<FilterScreenSideEffect>()
     val uiSideEffect = _uiSideEffect.receiveAsFlow()
 
     private val selectedOwnership = MutableStateFlow<Ownership?>(null)
-    private val selectedProviders = MutableStateFlow<List<Provider>>(emptyList())
+    private val selectedProviders = MutableStateFlow<List<ProviderId>>(emptyList())
 
     init {
         viewModelScope.launch {
             selectedProviders.value =
-                combine(availableProvidersUseCase(), relayListFilterRepository.selectedProviders) {
-                        allProviders,
-                        selectedConstraintProviders ->
-                        selectedConstraintProviders.toSelectedProviders(allProviders)
+                combine(
+                        providerToOwnershipsUseCase(),
+                        relayListFilterRepository.selectedProviders,
+                    ) { providerToOwnerships, selectedConstraintProviders ->
+                        selectedConstraintProviders.toSelectedProviders(
+                            providerToOwnerships.keys.toList()
+                        )
                     }
                     .first()
 
@@ -46,31 +49,25 @@ class FilterViewModel(
     }
 
     val uiState: StateFlow<RelayFilterUiState> =
-        combine(selectedOwnership, availableProvidersUseCase(), selectedProviders) {
-                selectedOwnership,
-                allProviders,
-                selectedProviders ->
-                RelayFilterUiState(
-                    selectedOwnership = selectedOwnership,
-                    allProviders = allProviders,
-                    selectedProviders = selectedProviders,
-                )
-            }
-            .stateIn(
-                viewModelScope,
-                SharingStarted.WhileSubscribed(),
-                RelayFilterUiState(
-                    allProviders = emptyList(),
-                    selectedOwnership = null,
-                    selectedProviders = emptyList(),
-                ),
-            )
+        combine(providerToOwnershipsUseCase(), selectedOwnership, selectedProviders, ::createState)
+            .stateIn(viewModelScope, SharingStarted.WhileSubscribed(), RelayFilterUiState())
+
+    private fun createState(
+        providerToOwnerships: Map<ProviderId, Set<Ownership>>,
+        selectedOwnership: Ownership?,
+        selectedProviders: List<ProviderId>,
+    ): RelayFilterUiState =
+        RelayFilterUiState(
+            providerToOwnerships = providerToOwnerships,
+            selectedOwnership = selectedOwnership,
+            selectedProviders = selectedProviders,
+        )
 
     fun setSelectedOwnership(ownership: Ownership?) {
         selectedOwnership.value = ownership
     }
 
-    fun setSelectedProvider(checked: Boolean, provider: Provider) {
+    fun setSelectedProvider(checked: Boolean, provider: ProviderId) {
         selectedProviders.value =
             if (checked) {
                 selectedProviders.value + provider
@@ -83,7 +80,7 @@ class FilterViewModel(
         viewModelScope.launch {
             selectedProviders.value =
                 if (isChecked) {
-                    availableProvidersUseCase().first()
+                    providerToOwnershipsUseCase().first().keys.toList()
                 } else {
                     emptyList()
                 }