Make UI state have constraints

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.toList()),
             )
         )
 }

android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/screen/FilterScreen.kt

@@ -45,6 +45,7 @@ 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.theme.AppTheme
@@ -96,7 +97,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 +122,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) {
@@ -163,14 +164,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(),
     )
 }
@@ -179,11 +177,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(),
     )
@@ -222,12 +220,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<List<ProviderId>>) =
+    when (constraint) {
+        Constraint.Any -> true
+        is Constraint.Only -> this in constraint.value
+    }
+
 @Composable
 private fun TopBar(onBackClick: () -> Unit) {
     Row(Modifier.fillMaxWidth(), verticalAlignment = Alignment.CenterVertically) {

android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/state/FilterConstrainExtensions.kt

@@ -11,15 +11,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> =
     if (size == allProviders.size) {
         Constraint.Any
     } else {
-        Constraint.Only(Providers(toHashSet()))
+        Constraint.Only(this)
     }

android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/state/RelayFilterUiState.kt

@@ -1,35 +1,43 @@
 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
 
 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<List<ProviderId>> = Constraint.Any,
 ) {
     val allProviders: List<ProviderId> = providerToOwnerships.keys.toList().sorted()
 
     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
+            is Constraint.Only ->
+                providerToOwnerships
+                    .filterValues { selectedOwnership.value in it }
+                    .keys
+                    .toList()
+                    .sorted()
         }
 
-    val isApplyButtonEnabled = selectedProviders.isNotEmpty()
+    val isApplyButtonEnabled = selectedProviders.getOrNull()?.isNotEmpty() != false
 
-    val isAllProvidersChecked = allProviders.size == selectedProviders.size
+    val isAllProvidersChecked = selectedProviders is Constraint.Any
 }

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

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

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,43 +44,70 @@ 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,
             selectedOwnership = selectedOwnership,
             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(emptyList())
                 }
         }
     }
 
     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(

android/app/src/test/kotlin/net/mullvad/mullvadvpn/repository/RelayListFilterRepositoryTest.kt

@@ -59,8 +59,7 @@ class RelayListFilterRepositoryTest {
     fun `when settings is updated selected providers should update`() = runTest {
         // Arrange
         val mockSettings: Settings = mockk()
-        val selectedProviders: Constraint<Providers> =
-            Constraint.Only(Providers(setOf(ProviderId("Prove"))))
+        val selectedProviders: Constraint<Providers> = Constraint.Only(listOf(ProviderId("Prove")))
         every { mockSettings.relaySettings.relayConstraints.providers } returns selectedProviders
 
         // Act, Assert
@@ -147,7 +146,7 @@ class RelayListFilterRepositoryTest {
     @Test
     fun `when successfully updating selected providers should return successful`() = runTest {
         // Arrange
-        val providers = Constraint.Only(Providers(setOf(ProviderId("Mopp"))))
+        val providers = Constraint.Only(listOf(ProviderId("Mopp")))
         coEvery { mockManagementService.setProviders(providers) } returns Unit.right()
 
         // Act
@@ -162,7 +161,7 @@ class RelayListFilterRepositoryTest {
     fun `when failing to update selected providers should return SetWireguardConstraintsError`() =
         runTest {
             // Arrange
-            val providers = Constraint.Only(Providers(setOf(ProviderId("Mopp"))))
+            val providers = Constraint.Only(listOf(ProviderId("Mopp")))
             val error = SetWireguardConstraintsError.Unknown(mockk())
             coEvery { mockManagementService.setProviders(providers) } returns error.left()
 

android/app/src/test/kotlin/net/mullvad/mullvadvpn/usecase/FilterChipUseCaseTest.kt

@@ -71,7 +71,7 @@ class FilterChipUseCaseTest {
     @Test
     fun `when provider filter is applied should return correct number of providers`() = runTest {
         // Arrange
-        val expectedProviders = Providers(providers = setOf(ProviderId("1"), ProviderId("2")))
+        val expectedProviders = listOf(ProviderId("1"), ProviderId("2"))
         selectedProviders.value = Constraint.Only(expectedProviders)
         providerToOwnerships.value =
             mapOf(
@@ -88,7 +88,7 @@ class FilterChipUseCaseTest {
     fun `when provider and ownership filter is applied should return correct filter chips`() =
         runTest {
             // Arrange
-            val expectedProviders = Providers(providers = setOf(ProviderId("1")))
+            val expectedProviders = listOf(ProviderId("1"))
             val expectedOwnership = Ownership.MullvadOwned
             selectedProviders.value = Constraint.Only(expectedProviders)
             selectedOwnership.value = Constraint.Only(expectedOwnership)