uiState doesn't update as expected

Whenever I trigger the "FavoritePost" or "DeletePost" and the uiState has TAB_TWO (Favorites) selected it either swaps to TAB_ONE, stays in TAB_TWO but displays all the Posts instead of just the favorite ones or just works properly.

Here is the ViewModel of that screen:

@HiltViewModel
class MainViewModel @Inject constructor(
    private val repo: MainRepository
) : ViewModelTemplate<MainState, MainScreenEvents>(MainState()) {

    init {
        initState()
    }

    override fun initState() {
        d("MainViewModel", "Initialized posts")
        viewModelScope.launch {
            repo.getPosts()
                .distinctUntilChanged()
                .collectLatest {
                    _uiState.emit(
                        _uiState.value.copy(posts = it)
                    )
                }
        }
    }

    override fun onEvent(event: MainScreenEvents) {
        when (event) {
            is ChangeTab -> {
                viewModelScope.launch {
                    d("MainViewModel", "Changing to ${event.tabNumber}")
                    when (event.tabNumber) {
                        TAB_ONE -> repo.getPosts()
                        TAB_TWO -> repo.getFavoritedPosts()
                    }.distinctUntilChanged()
                        .collect {
                        _uiState.emit(
                            _uiState.value.copy(
                                selectedTab = event.tabNumber,
                                posts = it
                            )
                        )
                    }
                }
            }

            is DeletePost -> {
                viewModelScope.launch {
                    repo.setDeletedPost(event.post)
                    if (event.post.isFavorited)
                        repo.setFavoritePost(event.post)
                }
            }

            is FavoritePost -> {
                viewModelScope.launch {
                    repo.setFavoritePost(event.post)
                }
            }
        }
    }
}

ViewModelTemplate:

abstract class ViewModelTemplate<T : States, in R : Events>(initialState: T) : ViewModel() {

    protected val _uiState: MutableStateFlow<T> = MutableStateFlow(initialState)
    val uiState: StateFlow<T> = _uiState

    protected open fun initState() {
        d("viewmodel", "I did nothing!")
    }

    abstract fun onEvent(event: R)
}

Events class:

sealed class Events {
    sealed class MainScreenEvents : Events() {
        data class FavoritePost(val post: PostsEntity) : MainScreenEvents()
        data class DeletePost(val post: PostsEntity) : MainScreenEvents()
        data class ChangeTab(val tabNumber: MainScreenTabId) : MainScreenEvents()
    }

    sealed class DeletedPostsScreenEvents : Events() {
        data class CheckPost(val post: DeletedPosts) : DeletedPostsScreenEvents()
        data object RestorePosts : DeletedPostsScreenEvents()
        data object RestoreAllPosts : DeletedPostsScreenEvents()
    }

    sealed class PostScreenEvents : Events() {
        data class PostEntity(val title: String, val body: String) : PostScreenEvents()
    }
}

Tabs:

enum class MainScreenTabId(val text: String, val icon: ImageVector) {
    TAB_ONE("Home", Default.Home),
    TAB_TWO("Favorite", Default.Favorite)
}

States:

sealed class States {
    data class MainState(
        val posts: List<PostsEntity> = emptyList(),
        val selectedTab: MainScreenTabId = TAB_ONE
    ) : States()

    data class DeletedPostsState(
        val posts: List<DeletedPosts> = emptyList()
    ) : States()

    data object EmptyState : States()
}

Repository:

class MainRepository @Inject constructor(
    private val postsService: PostsServiceImplementation,
    private val postsDao: PostsDao
) {
    init {
        CoroutineScope(IO).launch {
            postsDao.updateDbFromClient(
                postsService.getPosts().map {
                    it.toEntity()
                })
        }
    }

    fun getPosts(): Flow<List<PostsEntity>> {
        return postsDao.getAvailablePosts()
    }

    fun getFavoritedPosts(): Flow<List<PostsEntity>> {
        return postsDao.getFavoritedPosts()
    }

    fun getDeletedPosts(): Flow<List<PostsEntity>> {
        return postsDao.getDeletedPosts()
    }

    suspend fun setFavoritePost(post: PostsEntity) {
        postsDao.updatePost(post.copy(isFavorited = !post.isFavorited))
    }

    suspend fun setDeletedPost(post: PostsEntity) {
        postsDao.updatePost(post.copy(isDeleted = !post.isDeleted))
    }

    suspend fun createPosts(postRequest: PostRequest): PostResponse? {
        return postsService.createPosts(postRequest)
    }
}

I use RoomDB if that helps but I genuinely think that the problem is how I handle Flows inside the ViewModel.

Also if you have something to add about my code in general, please do, I want to learn!

I can guarantee the problem is not how I handle the events in the composable.

Tried centralizing the StateFlow handling inside a single function (didn't change anything)

Cheers!

You shouldn't collect flows in the view model, they should only be transformed and converted to StateFlows which are then collected in your composables.

Remove _uiState and uiState from ViewModelTemplate and replace the init block and initState in MainViewModel by this:

private val currentTab: MutableStateFlow<MainScreenTabId> = MutableStateFlow(TAB_ONE)

val uiState: StateFlow<MainState> = currentTab
    .flatMapLatest {
        when (it) {
            TAB_ONE -> repo.getPosts()
            TAB_TWO -> repo.getFavoritedPosts()
        }
    }
    .combine(currentTab, ::MainState)
    .stateIn(
        scope = viewModelScope,
        started = SharingStarted.WhileSubscribed(5_000),
        initialValue = MainState(),
    )

The base is the current tab which is wrapped in a flow so other flow operations can be based on it. The first thing that happens is that the flow is switched to the correct repository flow depending on the current tab. Then the content of the flow, together with the current tab is converted to a MainState object. Then the flow is converted to a StateFlow so it can be easily collected by the UI.

When switching tabs the only thing to do is to set currentTab and the StateFlow (and the UI) is updated accordingly. You just need to simplify the ChangeTab branch in onEvent to this:

is ChangeTab -> {
    currentTab.value = event.tabNumber
}