Make UI state have constraints

7 files changed, 95 insertions, 75 deletions

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 0aa3be8158..2a5188e8ee 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
@@ -2,6 +2,7 @@ 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.Constraint
 import net.mullvad.mullvadvpn.lib.model.Ownership
 import net.mullvad.mullvadvpn.lib.model.ProviderId
 
@@ -12,8 +13,8 @@ class FilterUiStatePreviewParameterProvider : PreviewParameterProvider<RelayFilt
         sequenceOf(
             RelayFilterUiState(
                 providerToOwnerships = PROVIDER_TO_OWNERSHIPS,
-                selectedOwnership = Ownership.MullvadOwned,
-                selectedProviders = PROVIDER_TO_OWNERSHIPS.keys.toList(),
+                selectedOwnership = Constraint.Only(Ownership.MullvadOwned),
+                selectedProviders = Constraint.Only(PROVIDER_TO_OWNERSHIPS.keys),
             )
         )
 }
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 3d3f937c73..00d1e3bb9a 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
@@ -45,8 +45,10 @@ import net.mullvad.mullvadvpn.compose.preview.FilterUiStatePreviewParameterProvi
 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.Constraint
 import net.mullvad.mullvadvpn.lib.model.Ownership
 import net.mullvad.mullvadvpn.lib.model.ProviderId
+import net.mullvad.mullvadvpn.lib.model.Providers
 import net.mullvad.mullvadvpn.lib.theme.AppTheme
 import net.mullvad.mullvadvpn.lib.theme.Dimens
 import net.mullvad.mullvadvpn.viewmodel.FilterScreenSideEffect
@@ -96,7 +98,7 @@ fun FilterScreen(
     state: RelayFilterUiState,
     onBackClick: () -> Unit,
     onApplyClick: () -> Unit,
-    onSelectedOwnership: (ownership: Ownership?) -> Unit,
+    onSelectedOwnership: (ownership: Constraint<Ownership>) -> Unit,
     onAllProviderCheckChange: (isChecked: Boolean) -> Unit,
     onSelectedProvider: (checked: Boolean, provider: ProviderId) -> Unit,
 ) {
@@ -121,14 +123,14 @@ fun FilterScreen(
             }
             if (ownershipExpanded) {
                 item(key = Keys.OWNERSHIP_ALL, contentType = ContentType.ITEM) {
-                    AnyOwnership(state, onSelectedOwnership)
+                    AnyOwnership(state, { onSelectedOwnership(Constraint.Any) })
                 }
                 itemsWithDivider(
                     key = { it.name },
                     contentType = { ContentType.ITEM },
                     items = state.selectableOwnerships,
                 ) { ownership ->
-                    Ownership(ownership, state, onSelectedOwnership)
+                    Ownership(ownership, state, { onSelectedOwnership(Constraint.Only(it)) })
                 }
             }
             itemWithDivider(key = Keys.PROVIDERS_TITLE, contentType = ContentType.HEADER) {
@@ -171,14 +173,11 @@ private fun LazyItemScope.OwnershipHeader(expanded: Boolean, onToggleExpanded: (
 }
 
 @Composable
-private fun LazyItemScope.AnyOwnership(
-    state: RelayFilterUiState,
-    onSelectedOwnership: (ownership: Ownership?) -> Unit,
-) {
+private fun LazyItemScope.AnyOwnership(state: RelayFilterUiState, onSelectedOwnership: () -> Unit) {
     SelectableCell(
         title = stringResource(id = R.string.any),
-        isSelected = state.selectedOwnership == null,
-        onCellClicked = { onSelectedOwnership(null) },
+        isSelected = state.selectedOwnership is Constraint.Any,
+        onCellClicked = { onSelectedOwnership() },
         modifier = Modifier.animateItem(),
     )
 }
@@ -187,11 +186,11 @@ private fun LazyItemScope.AnyOwnership(
 private fun LazyItemScope.Ownership(
     ownership: Ownership,
     state: RelayFilterUiState,
-    onSelectedOwnership: (ownership: Ownership?) -> Unit,
+    onSelectedOwnership: (ownership: Ownership) -> Unit,
 ) {
     SelectableCell(
         title = stringResource(id = ownership.stringResource()),
-        isSelected = ownership == state.selectedOwnership,
+        isSelected = ownership == state.selectedOwnership.getOrNull(),
         onCellClicked = { onSelectedOwnership(ownership) },
         modifier = Modifier.animateItem(),
     )
@@ -230,12 +229,18 @@ private fun LazyItemScope.Provider(
 ) {
     CheckboxCell(
         title = providerId.value,
-        checked = providerId in state.selectedProviders,
+        checked = providerId.isChecked(state.selectedProviders),
         onCheckedChange = { checked -> onSelectedProvider(checked, providerId) },
         modifier = Modifier.animateItem(),
     )
 }
 
+private fun ProviderId.isChecked(constraint: Constraint<Providers>) =
+    when (constraint) {
+        Constraint.Any -> true
+        is Constraint.Only -> this in constraint.value
+    }
+
 @Composable
 private fun LazyItemScope.RemovedProvider(
     providerId: ProviderId,
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 2ddf35fad5..b7531f15ec 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,6 @@ 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.ProviderId
 import net.mullvad.mullvadvpn.lib.model.Providers
 
 fun Ownership?.toOwnershipConstraint(): Constraint<Ownership> =
@@ -11,15 +10,9 @@ fun Ownership?.toOwnershipConstraint(): Constraint<Ownership> =
         else -> Constraint.Only(this)
     }
 
-fun Constraint<Providers>.toSelectedProviders(allProviders: List<ProviderId>): List<ProviderId> =
-    when (this) {
-        Constraint.Any -> allProviders
-        is Constraint.Only -> value.providers.toList()
-    }
-
-fun List<ProviderId>.toConstraintProviders(allProviders: List<ProviderId>): Constraint<Providers> =
+fun Providers.toConstraintProviders(allProviders: Providers): Constraint<Providers> =
     if (size == allProviders.size) {
         Constraint.Any
     } else {
-        Constraint.Only(Providers(toHashSet()))
+        Constraint.Only(this)
     }
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 0c81ac0aa5..d3e8a2b685 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,37 +1,41 @@
 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.ProviderId
+import net.mullvad.mullvadvpn.lib.model.Providers
 
 data class RelayFilterUiState(
     private val providerToOwnerships: Map<ProviderId, Set<Ownership>> = emptyMap(),
-    val selectedOwnership: Ownership? = null,
-    val selectedProviders: List<ProviderId> = emptyList(),
+    val selectedOwnership: Constraint<Ownership> = Constraint.Any,
+    val selectedProviders: Constraint<Providers> = Constraint.Any,
 ) {
-    val allProviders: List<ProviderId> = providerToOwnerships.keys.toList().sorted()
+    val allProviders: Providers = providerToOwnerships.keys
 
     val selectableOwnerships: List<Ownership> =
-        if (selectedProviders.isEmpty()) {
-                Ownership.entries
-            } else {
-                providerToOwnerships
-                    .filterKeys { it in selectedProviders }
-                    .values
-                    .flatten()
-                    .distinct()
-            }
-            .sorted()
+        when (selectedProviders) {
+            Constraint.Any -> Ownership.entries
+            is Constraint.Only ->
+                if (selectedProviders.value.isEmpty()) {
+                    Ownership.entries
+                } else {
+                    providerToOwnerships
+                        .filterKeys { it in selectedProviders.value }
+                        .values
+                        .flatten()
+                        .distinct()
+                }
+        }.sorted()
 
     val selectableProviders: List<ProviderId> =
-        if (selectedOwnership != null) {
-            providerToOwnerships.filterValues { selectedOwnership in it }.keys.toList().sorted()
-        } else {
-            allProviders
-        }
+        when (selectedOwnership) {
+            Constraint.Any -> allProviders.toList()
+            is Constraint.Only ->
+                providerToOwnerships.filterValues { selectedOwnership.value in it }.keys
+        }.sorted()
 
+    val isApplyButtonEnabled = selectedProviders.getOrNull()?.isNotEmpty() != false
     val removedProviders: List<ProviderId> = selectedProviders - allProviders
 
-    val isApplyButtonEnabled = selectedProviders.isNotEmpty()
-
-    val isAllProvidersChecked = allProviders.size == selectedProviders.size
+    val isAllProvidersChecked = selectedProviders is Constraint.Any
 }
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 21d3294de4..d58da5bc9a 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
@@ -46,7 +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 -> provider in providersConstraint.value.providers
+            is RelayItem.Location.Relay -> provider in providersConstraint.value
         }
     } 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 361e29944e..8c542202ff 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
@@ -57,7 +57,7 @@ class FilterChipUseCase(
             when (selectedConstraintProviders) {
                 is Constraint.Any -> null
                 is Constraint.Only ->
-                    selectedConstraintProviders.value.providers
+                    selectedConstraintProviders.value
                         .filter { providerId ->
                             if (ownershipFilter == null) {
                                 true
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 c05515126e..41d98b40fb 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
@@ -2,6 +2,7 @@ package net.mullvad.mullvadvpn.viewmodel
 
 import androidx.lifecycle.ViewModel
 import androidx.lifecycle.viewModelScope
+import kotlin.collections.plus
 import kotlinx.coroutines.channels.Channel
 import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.flow.SharingStarted
@@ -10,13 +11,13 @@ import kotlinx.coroutines.flow.combine
 import kotlinx.coroutines.flow.first
 import kotlinx.coroutines.flow.receiveAsFlow
 import kotlinx.coroutines.flow.stateIn
+import kotlinx.coroutines.flow.update
 import kotlinx.coroutines.launch
 import net.mullvad.mullvadvpn.compose.state.RelayFilterUiState
-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.Constraint
 import net.mullvad.mullvadvpn.lib.model.Ownership
 import net.mullvad.mullvadvpn.lib.model.ProviderId
+import net.mullvad.mullvadvpn.lib.model.Providers
 import net.mullvad.mullvadvpn.repository.RelayListFilterRepository
 import net.mullvad.mullvadvpn.usecase.ProviderToOwnershipsUseCase
 
@@ -27,24 +28,13 @@ class FilterViewModel(
     private val _uiSideEffect = Channel<FilterScreenSideEffect>()
     val uiSideEffect = _uiSideEffect.receiveAsFlow()
 
-    private val selectedOwnership = MutableStateFlow<Ownership?>(null)
-    private val selectedProviders = MutableStateFlow<List<ProviderId>>(emptyList())
+    private val selectedOwnership = MutableStateFlow<Constraint<Ownership>>(Constraint.Any)
+    private val selectedProviders = MutableStateFlow<Constraint<Providers>>(Constraint.Any)
 
     init {
         viewModelScope.launch {
-            selectedProviders.value =
-                combine(
-                        providerToOwnershipsUseCase(),
-                        relayListFilterRepository.selectedProviders,
-                    ) { providerToOwnerships, selectedConstraintProviders ->
-                        selectedConstraintProviders.toSelectedProviders(
-                            providerToOwnerships.keys.toList()
-                        )
-                    }
-                    .first()
-
-            val ownershipConstraint = relayListFilterRepository.selectedOwnership.first()
-            selectedOwnership.value = ownershipConstraint.getOrNull()
+            selectedProviders.value = relayListFilterRepository.selectedProviders.first()
+            selectedOwnership.value = relayListFilterRepository.selectedOwnership.first()
         }
     }
 
@@ -54,8 +44,8 @@ class FilterViewModel(
 
     private fun createState(
         providerToOwnerships: Map<ProviderId, Set<Ownership>>,
-        selectedOwnership: Ownership?,
-        selectedProviders: List<ProviderId>,
+        selectedOwnership: Constraint<Ownership>,
+        selectedProviders: Constraint<Providers>,
     ): RelayFilterUiState =
         RelayFilterUiState(
             providerToOwnerships = providerToOwnerships,
@@ -63,34 +53,61 @@ class FilterViewModel(
             selectedProviders = selectedProviders,
         )
 
-    fun setSelectedOwnership(ownership: Ownership?) {
+    fun setSelectedOwnership(ownership: Constraint<Ownership>) {
         selectedOwnership.value = ownership
     }
 
     fun setSelectedProvider(checked: Boolean, provider: ProviderId) {
-        selectedProviders.value =
+        selectedProviders.update {
             if (checked) {
-                selectedProviders.value + provider
+                it.check(provider, uiState.value.allProviders)
             } else {
-                selectedProviders.value - provider
+                it.uncheck(provider, uiState.value.allProviders)
+            }
+        }
+    }
+
+    private fun Constraint<Providers>.check(
+        provider: ProviderId,
+        allProviders: Providers,
+    ): Constraint<Providers> {
+        return when (this) {
+            is Constraint.Any -> Constraint.Any
+            is Constraint.Only -> {
+                val newProviderList = value + provider
+                if (allProviders.size == newProviderList.size) {
+                    Constraint.Any
+                } else {
+                    Constraint.Only(newProviderList)
+                }
             }
+        }
+    }
+
+    private fun Constraint<Providers>.uncheck(
+        provider: ProviderId,
+        allProviders: Providers,
+    ): Constraint<Providers> {
+        return when (this) {
+            is Constraint.Any -> Constraint.Only(allProviders - provider)
+            is Constraint.Only -> Constraint.Only(value - provider)
+        }
     }
 
     fun setAllProviders(isChecked: Boolean) {
         viewModelScope.launch {
             selectedProviders.value =
                 if (isChecked) {
-                    providerToOwnershipsUseCase().first().keys.toList()
+                    Constraint.Any
                 } else {
-                    emptyList()
+                    Constraint.Only(emptySet())
                 }
         }
     }
 
     fun onApplyButtonClicked() {
-        val newSelectedOwnership = selectedOwnership.value.toOwnershipConstraint()
-        val newSelectedProviders =
-            selectedProviders.value.toConstraintProviders(uiState.value.allProviders)
+        val newSelectedOwnership = selectedOwnership.value
+        val newSelectedProviders = selectedProviders.value
 
         viewModelScope.launch {
             relayListFilterRepository.updateSelectedOwnershipAndProviderFilter(