Refactor: Simplify location state management by introducing LocationsUIState and updating related ViewModel and UI logic

src/main/kotlin/org/jetbrains/plugins/template/weatherApp/services/WeatherAppViewModel.kt

@@ -2,7 +2,9 @@ package org.jetbrains.plugins.template.weatherApp.services
 
 import com.intellij.openapi.Disposable
 import kotlinx.coroutines.*
-import kotlinx.coroutines.flow.*
+import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.MutableStateFlow
+import kotlinx.coroutines.flow.asStateFlow
 import org.jetbrains.plugins.template.weatherApp.model.Location
 import org.jetbrains.plugins.template.weatherApp.model.SelectableLocation
 import org.jetbrains.plugins.template.weatherApp.model.WeatherForecastData
@@ -22,7 +24,7 @@ interface MyLocationsViewModelApi : Disposable {
 
     fun onLocationSelected(selectedLocationIndex: Int)
 
-    val myLocationsFlow: Flow<List<SelectableLocation>>
+    val myLocationsUIStateFlow: Flow<LocationsUIState>
 }
 
 /**
@@ -102,6 +104,8 @@ class LocationsUIState private constructor(
     val selectedLocation: Location?
         get() = locations.getOrNull(selectedIndex)
 
+    val isEmpty: Boolean get() = locations.isEmpty()
+
     /**
      * Convert to UI representation with selection state
      */
@@ -217,9 +221,19 @@ class WeatherAppViewModel(
 
     private var currentWeatherJob: Job? = null
 
-    private val myLocations = MutableStateFlow(myInitialLocations)
+    private val _myLocationsUIStateFlow = MutableStateFlow(LocationsUIState.initial(myInitialLocations))
 
-    private val selectedLocationIndex = MutableStateFlow(myLocations.value.lastIndex)
+    /**
+     * A state flow representing the current UI state of locations, including the list of locations
+     * and the selected location index. This is observed by the UI layer to render location data and
+     * selection state dynamically.
+     *
+     * This ensures that the state is safely encapsulated and modifications only occur through
+     * authorized ViewModel methods.
+     *
+     * @see LocationsUIState
+     */
+    override val myLocationsUIStateFlow: Flow<LocationsUIState> = _myLocationsUIStateFlow.asStateFlow()
 
     private val _weatherState = MutableStateFlow<WeatherForecastUIState>(WeatherForecastUIState.Empty)
 
@@ -235,24 +249,10 @@ class WeatherAppViewModel(
      */
     override val weatherForecastUIState: Flow<WeatherForecastUIState> = _weatherState.asStateFlow()
 
-    /**
-     * A [StateFlow] that emits a list of [SelectableLocation] objects representing the user's
-     * current locations along with the selection state of each location.
-     */
-    override val myLocationsFlow: StateFlow<List<SelectableLocation>> = myLocations
-        .combine(selectedLocationIndex) { locations, selectedIndex ->
-            locations.mapIndexed { index, location ->
-                SelectableLocation(location, index == selectedIndex)
-            }
-        }.stateIn(viewModelScope, SharingStarted.WhileSubscribed(), emptyList())
-
     override fun onAddLocation(locationToAdd: Location) {
-        if (myLocations.value.contains(locationToAdd)) {
+        val newState = _myLocationsUIStateFlow.value.withLocationAdded(locationToAdd)
-            selectedLocationIndex.value = myLocations.value.indexOf(locationToAdd)
+        updateLocationsUIStateWith(newState)
-        } else {
+
-            myLocations.value += locationToAdd
-            selectedLocationIndex.value = myLocations.value.lastIndex
-        }
 
         if (_weatherState.value.getLocationOrNull() != locationToAdd) {
             onReloadWeatherForecast()
@@ -260,27 +260,23 @@ class WeatherAppViewModel(
     }
 
     override fun onDeleteLocation(locationToDelete: Location) {
-        myLocations.value -= locationToDelete
+        val newState = _myLocationsUIStateFlow.value.withLocationDeleted(locationToDelete)
-
+        updateLocationsUIStateWith(newState)
-        val itemIndex = myLocations.value.indexOf(locationToDelete)
-        val currentSelectedIndex = selectedLocationIndex.value
-        if (itemIndex in 0..currentSelectedIndex) {
-            selectedLocationIndex.value = (currentSelectedIndex - 1).coerceAtLeast(0)
-        }
 
         onReloadWeatherForecast()
     }
 
     override fun onLocationSelected(selectedLocationIndex: Int) {
-        if (this.selectedLocationIndex.value == selectedLocationIndex) return
+        val newState = _myLocationsUIStateFlow.value.withItemAtIndexSelected(selectedLocationIndex)
+        updateLocationsUIStateWith(newState)
 
-        this.selectedLocationIndex.value = selectedLocationIndex
+        if (_weatherState.value.getLocationOrNull() != newState.selectedLocation) {
-
+            onReloadWeatherForecast()
-        onReloadWeatherForecast()
+        }
     }
 
     override fun onReloadWeatherForecast() {
-        myLocations.value.getOrNull(selectedLocationIndex.value)?.let { location ->
+        _myLocationsUIStateFlow.value.selectedLocation?.let { location ->
             onLoadWeatherForecast(location)
         }
     }
@@ -314,6 +310,10 @@ class WeatherAppViewModel(
         viewModelScope.cancel()
     }
 
+    private fun updateLocationsUIStateWith(newState: LocationsUIState) {
+        _myLocationsUIStateFlow.value = newState
+    }
+
     private fun errorStateFor(
         location: Location,
         error: Throwable

src/main/kotlin/org/jetbrains/plugins/template/weatherApp/ui/WeatherAppSample.kt

@@ -15,7 +15,7 @@ import androidx.compose.ui.unit.dp
 import org.jetbrains.jewel.bridge.retrieveColorOrUnspecified
 import org.jetbrains.jewel.foundation.lazy.SelectableLazyColumn
 import org.jetbrains.jewel.foundation.lazy.SelectionMode
-import org.jetbrains.jewel.foundation.lazy.items
+import org.jetbrains.jewel.foundation.lazy.itemsIndexed
 import org.jetbrains.jewel.foundation.lazy.rememberSelectableLazyListState
 import org.jetbrains.jewel.foundation.theme.JewelTheme
 import org.jetbrains.jewel.ui.component.*
@@ -23,11 +23,7 @@ import org.jetbrains.jewel.ui.icon.IconKey
 import org.jetbrains.jewel.ui.icons.AllIconsKeys
 import org.jetbrains.plugins.template.ComposeTemplateBundle
 import org.jetbrains.plugins.template.weatherApp.model.Location
-import org.jetbrains.plugins.template.weatherApp.model.SelectableLocation
+import org.jetbrains.plugins.template.weatherApp.services.*
-import org.jetbrains.plugins.template.weatherApp.services.MyLocationsViewModelApi
-import org.jetbrains.plugins.template.weatherApp.services.SearchAutoCompletionItemProvider
-import org.jetbrains.plugins.template.weatherApp.services.WeatherForecastUIState
-import org.jetbrains.plugins.template.weatherApp.services.WeatherViewModelApi
 import org.jetbrains.plugins.template.weatherApp.ui.components.SearchToolbarMenu
 import org.jetbrains.plugins.template.weatherApp.ui.components.WeatherDetailsCard
 
@@ -89,10 +85,11 @@ fun MyLocationsListWithEmptyListPlaceholder(
     modifier: Modifier = Modifier,
     myLocationsViewModelApi: MyLocationsViewModelApi
 ) {
-    val myLocations = myLocationsViewModelApi.myLocationsFlow.collectAsState(emptyList()).value
+    val myLocationsUIState =
+        myLocationsViewModelApi.myLocationsUIStateFlow.collectAsState(LocationsUIState.empty()).value
 
-    if (myLocations.isNotEmpty()) {
+    if (!myLocationsUIState.isEmpty) {
-        MyLocationList(myLocations, modifier, myLocationsViewModelApi)
+        MyLocationList(myLocationsUIState, modifier, myLocationsViewModelApi)
     } else {
         EmptyListPlaceholder(modifier)
     }
@@ -128,23 +125,23 @@ private fun EmptyListPlaceholder(
 
 @Composable
 private fun MyLocationList(
-    myLocations: List<SelectableLocation>,
+    myLocationsUIState: LocationsUIState,
     modifier: Modifier,
     myLocationsViewModelApi: MyLocationsViewModelApi
 ) {
     val listState = rememberSelectableLazyListState()
     // JEWEL-938 This will trigger on SelectableLazyColum's `onSelectedIndexesChange` callback
-    LaunchedEffect(myLocations) {
+    LaunchedEffect(myLocationsUIState) {
         var lastActiveItemIndex = -1
         val selectedItemKeys = mutableSetOf<String>()
-        myLocations.forEachIndexed { index, location ->
+        myLocationsUIState.locations.forEachIndexed { index, location ->
-            if (location.isSelected) {
+            if (index == myLocationsUIState.selectedIndex) {
                 if (lastActiveItemIndex == -1) {
                     // Only the first selected item should be active
                     lastActiveItemIndex = index
                 }
                 // Must match the key used in the `items()` call's `key` parameter to ensure correct item identity.
-                selectedItemKeys.add(location.location.label)
+                selectedItemKeys.add(location.label)
             }
         }
         // Sets the first selected item as an active item to avoid triggering on click event when user clocks on it
@@ -162,13 +159,13 @@ private fun MyLocationList(
             myLocationsViewModelApi.onLocationSelected(selectedLocationIndex)
         },
     ) {
-        items(
+        itemsIndexed(
-            items = myLocations,
+            items = myLocationsUIState.locations,
-            key = { item -> item.location.label },
+            key = { _, item -> item.label },
-        ) { item ->
+        ) { index, item ->
 
             ContentItemRow(
-                item = item.location, isSelected = item.isSelected, isActive = isActive
+                item = item, isSelected = myLocationsUIState.selectedIndex == index, isActive = isActive
             )
         }
     }
@@ -201,7 +198,9 @@ private fun RightColumn(
     searchAutoCompletionItemProvider: SearchAutoCompletionItemProvider<Location>,
     modifier: Modifier = Modifier,
 ) {
-    val weatherForecastData = weatherViewModelApi.weatherForecastUIState.collectAsState(WeatherForecastUIState.Empty).value
+    val weatherForecastData = weatherViewModelApi
+        .weatherForecastUIState
+        .collectAsState(WeatherForecastUIState.Empty).value
 
     Column(modifier) {
         SearchToolbarMenu(

src/test/kotlin/org/jetbrains/plugins/template/MyLocationListTest.kt

@@ -10,9 +10,9 @@ import androidx.compose.ui.test.onNodeWithText
 import androidx.compose.ui.test.performClick
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.MutableStateFlow
-import kotlinx.coroutines.flow.combine
+import kotlinx.coroutines.flow.asStateFlow
 import org.jetbrains.plugins.template.weatherApp.model.Location
-import org.jetbrains.plugins.template.weatherApp.model.SelectableLocation
+import org.jetbrains.plugins.template.weatherApp.services.LocationsUIState
 import org.jetbrains.plugins.template.weatherApp.services.MyLocationsViewModelApi
 import org.jetbrains.plugins.template.weatherApp.ui.MyLocationsListWithEmptyListPlaceholder
 import org.junit.Test
@@ -145,40 +145,27 @@ internal class MyLocationListTest : ComposeBasedTestCase() {
         locations: List<Location> = emptyList()
     ) : MyLocationsViewModelApi {
 
-        private val locationsFlow = MutableStateFlow(locations.toMutableList())
+        private val _myLocationsUIStateFlow: MutableStateFlow<LocationsUIState> =
+            MutableStateFlow(LocationsUIState.initial(locations))
 
-        private val selectedItemIndex = MutableStateFlow(if (locations.isNotEmpty()) 0 else -1)
-        private val _myLocations = locationsFlow
-            .combine(selectedItemIndex) { locations, selectedIndex ->
-                locations.mapIndexed { index, location ->
-                    SelectableLocation(location, index == selectedIndex)
-                }
-            }
-
-        override val myLocationsFlow: Flow<List<SelectableLocation>> = _myLocations
 
         override fun onAddLocation(locationToAdd: Location) {
-            val currentLocations = locationsFlow.value
+            _myLocationsUIStateFlow.value = _myLocationsUIStateFlow.value.withLocationAdded(locationToAdd)
-            currentLocations.add(locationToAdd)
-
-            locationsFlow.value = currentLocations
-            selectedItemIndex.value = currentLocations.lastIndex
         }
 
         override fun onDeleteLocation(locationToDelete: Location) {
-            val currentLocations = locationsFlow.value
+            _myLocationsUIStateFlow.value = _myLocationsUIStateFlow.value.withLocationDeleted(locationToDelete)
-            currentLocations.remove(locationToDelete)
 
-            locationsFlow.value = currentLocations
-            selectedItemIndex.value = currentLocations.lastIndex
         }
 
         override fun onLocationSelected(selectedLocationIndex: Int) {
-            selectedItemIndex.value = selectedLocationIndex
+            _myLocationsUIStateFlow.value = _myLocationsUIStateFlow.value.withItemAtIndexSelected(selectedLocationIndex)
         }
 
-        override fun dispose() {
+        override val myLocationsUIStateFlow: Flow<LocationsUIState>
+            get() = _myLocationsUIStateFlow.asStateFlow()
 
+        override fun dispose() {
         }
     }