Scroll to initial country.

feature/registration/src/main/java/org/signal/registration/RegistrationNavigation.kt

@@ -88,7 +88,7 @@ sealed interface RegistrationRoute : NavKey, Parcelable {
   data object PhoneNumberEntry : RegistrationRoute
 
   @Serializable
-  data object CountryCodePicker : RegistrationRoute
+  data class CountryCodePicker(val country: Country? = null) : RegistrationRoute
 
   @Serializable
   data object VerificationCodeEntry : RegistrationRoute
@@ -235,7 +235,7 @@ fun RegistrationNavHost(
     onBack = { viewModel.onEvent(RegistrationFlowEvent.NavigateBack) },
     modifier = modifier,
     transitionSpec = {
-      if (targetState.key == RegistrationRoute.CountryCodePicker.toString()) {
+      if (targetState.key is RegistrationRoute.CountryCodePicker) {
         TransitionSpecs.VerticalSlide.transitionSpec.invoke(this)
       } else {
         TransitionSpecs.HorizontalSlide.transitionSpec.invoke(this)
@@ -243,7 +243,7 @@ fun RegistrationNavHost(
     },
     popTransitionSpec = {
       when {
-        initialState.key == RegistrationRoute.CountryCodePicker.toString() -> {
+        initialState.key is RegistrationRoute.CountryCodePicker -> {
           TransitionSpecs.VerticalSlide.popTransitionSpec.invoke(this)
         }
         initialState.key == RegistrationRoute.EnterAepScreen.toString() -> {
@@ -258,9 +258,10 @@ fun RegistrationNavHost(
       }
     },
     predictivePopTransitionSpec = {
-      when (initialState.key) {
-        RegistrationRoute.CountryCodePicker.toString() -> TransitionSpecs.VerticalSlide.predictivePopTransitionSpec.invoke(this, it)
-        else -> TransitionSpecs.HorizontalSlide.predictivePopTransitionSpec.invoke(this, it)
+      if (initialState.key is RegistrationRoute.CountryCodePicker) {
+        TransitionSpecs.VerticalSlide.predictivePopTransitionSpec.invoke(this, it)
+      } else {
+        TransitionSpecs.HorizontalSlide.predictivePopTransitionSpec.invoke(this, it)
       }
     }
   )
@@ -335,13 +336,14 @@ private fun EntryProviderScope<NavKey>.navigationEntries(
   }
 
   // -- Country Code Picker
-  entry<RegistrationRoute.CountryCodePicker> {
+  entry<RegistrationRoute.CountryCodePicker> { key ->
     val viewModel: CountryCodePickerViewModel = viewModel(
       factory = CountryCodePickerViewModel.Factory(
         repository = CountryCodePickerRepository(),
         parentEventEmitter = parentEventEmitter,
         resultBus = registrationViewModel.resultBus,
-        resultKey = COUNTRY_CODE_RESULT
+        resultKey = COUNTRY_CODE_RESULT,
+        initialCountry = key.country
       )
     )
     val state by viewModel.state.collectAsStateWithLifecycle()

feature/registration/src/main/java/org/signal/registration/screens/countrycode/Country.kt

@@ -7,10 +7,12 @@ package org.signal.registration.screens.countrycode
 
 import android.os.Parcelable
 import kotlinx.parcelize.Parcelize
+import kotlinx.serialization.Serializable
 
 /**
  * Data class string describing useful characteristics of countries when selecting one. Used in the [CountryCodeState]
  * An example is: Country(emoji=🇺🇸, name = "United States", countryCode = 1, regionCode= "US")
  */
+@Serializable
 @Parcelize
 data class Country(val emoji: String, val name: String, val countryCode: Int, val regionCode: String) : Parcelable

feature/registration/src/main/java/org/signal/registration/screens/phonenumber/PhoneNumberEntryViewModel.kt

@@ -31,6 +31,7 @@ import org.signal.registration.RegistrationFlowState
 import org.signal.registration.RegistrationRepository
 import org.signal.registration.RegistrationRoute
 import org.signal.registration.screens.EventDrivenViewModel
+import org.signal.registration.screens.countrycode.Country
 import org.signal.registration.screens.countrycode.CountryUtils
 import org.signal.registration.screens.localbackuprestore.LocalBackupRestoreResult
 import org.signal.registration.screens.phonenumber.PhoneNumberEntryState.OneTimeEvent
@@ -109,7 +110,13 @@ class PhoneNumberEntryViewModel(
         stateEmitter(localState.copy(showSpinner = false))
       }
       is PhoneNumberEntryScreenEvents.CountryPicker -> {
-        state.also { parentEventEmitter.navigateTo(RegistrationRoute.CountryCodePicker) }
+        state.also {
+          parentEventEmitter.navigateTo(
+            RegistrationRoute.CountryCodePicker(
+              Country(state.countryEmoji, state.countryName, state.countryCode.toIntOrNull() ?: 0, state.regionCode).takeIf { state.countryName.isNotEmpty() }
+            )
+          )
+        }
       }
       is PhoneNumberEntryScreenEvents.CaptchaCompleted -> {
         stateEmitter(applyCaptchaCompleted(state, event.token, parentEventEmitter))

feature/registration/src/test/java/org/signal/registration/screens/phonenumber/PhoneNumberEntryViewModelTest.kt

@@ -243,7 +243,7 @@ class PhoneNumberEntryViewModelTest {
 
     assertThat(emittedEvents).hasSize(1)
     assertThat(emittedEvents.first()).isEqualTo(
-      RegistrationFlowEvent.NavigateToScreen(RegistrationRoute.CountryCodePicker)
+      RegistrationFlowEvent.NavigateToScreen(RegistrationRoute.CountryCodePicker())
     )
   }