Introduces SavedStateHandle.removeWorkflowState()

If you're juggling multiple workflow runtimes per `Activity` (if you're the kind of person who needs to cal `Job.cancel` on the `Job` returned from `WorkflowLayout.take`) then you might also need to throw away the state captured by an AndroidX `SavedStateHandler` passed to the Android flavor of `renderWorkflowIn()` to prevent memory leaks.

Author: rjrjr

workflow-ui/core-android/build.gradle.kts

@@ -11,7 +11,10 @@ android {
 }
 
 dependencies {
+  androidTestImplementation(libs.androidx.activity.ktx)
   androidTestImplementation(libs.androidx.appcompat)
+  androidTestImplementation(libs.androidx.lifecycle.viewmodel.ktx)
+  androidTestImplementation(libs.androidx.lifecycle.viewmodel.savedstate)
   androidTestImplementation(libs.truth)
 
   api(libs.androidx.lifecycle.common)

workflow-ui/core-android/src/androidTest/java/com/squareup/workflow1/ui/AndroidRenderWorkflowInTest.kt

@@ -0,0 +1,82 @@
+package com.squareup.workflow1.ui
+
+import android.widget.FrameLayout
+import androidx.activity.ComponentActivity
+import androidx.lifecycle.SavedStateHandle
+import androidx.lifecycle.ViewModel
+import androidx.test.ext.junit.rules.ActivityScenarioRule
+import com.squareup.workflow1.StatelessWorkflow
+import com.squareup.workflow1.ui.internal.test.IdlingDispatcherRule
+import leakcanary.DetectLeaksAfterTestSuccess
+import org.junit.Rule
+import org.junit.Test
+import org.junit.rules.RuleChain
+import androidx.activity.viewModels
+import androidx.lifecycle.Lifecycle.State.CREATED
+import kotlinx.coroutines.flow.StateFlow
+import androidx.lifecycle.viewModelScope
+import com.google.common.truth.Truth.assertThat
+import kotlinx.coroutines.Job
+
+@OptIn(WorkflowUiExperimentalApi::class)
+internal class AndroidRenderWorkflowInTest {
+  @get:Rule val scenarioRule = ActivityScenarioRule(ComponentActivity::class.java)
+  private val scenario get() = scenarioRule.scenario
+
+  @get:Rule val rules: RuleChain = RuleChain.outerRule(DetectLeaksAfterTestSuccess())
+    .around(scenarioRule)
+    .around(IdlingDispatcherRule)
+
+  @Test fun removeWorkflowStateDoesWhatItSaysOnTheTin() {
+    var job: Job? = null
+    scenario.onActivity { activity ->
+      val model: SomeViewModel by activity.viewModels()
+      val renderings: StateFlow<Screen> = renderWorkflowIn(
+        workflow = SomeWorkflow,
+        scope = model.viewModelScope,
+        savedStateHandle = model.savedStateHandle
+      )
+
+      val layout = WorkflowLayout(activity)
+      activity.setContentView(layout)
+
+      assertThat(model.savedStateHandle.contains(KEY)).isFalse()
+
+      job = layout.take(activity.lifecycle, renderings)
+      assertThat(model.savedStateHandle.contains(KEY)).isFalse()
+    }
+
+    scenario.moveToState(CREATED)
+    scenario.onActivity { activity ->
+      val model: SomeViewModel by activity.viewModels()
+      assertThat(model.savedStateHandle.contains(KEY)).isTrue()
+
+      job?.cancel()
+      assertThat(model.savedStateHandle.contains(KEY)).isTrue()
+
+      model.savedStateHandle.removeWorkflowState()
+      assertThat(model.savedStateHandle.contains(KEY)).isFalse()
+    }
+  }
+
+  object SomeScreen : AndroidScreen<SomeScreen> {
+    override val viewFactory: ScreenViewFactory<SomeScreen> =
+      ScreenViewFactory.fromCode { _, initialEnvironment, context, _ ->
+        ScreenViewHolder(
+          initialEnvironment,
+          FrameLayout(context)
+        ) { _, _ -> }
+      }
+  }
+
+  object SomeWorkflow : StatelessWorkflow<Unit, Nothing, Screen>() {
+    override fun render(
+      renderProps: Unit,
+      context: RenderContext
+    ): Screen {
+      return SomeScreen
+    }
+  }
+
+  class SomeViewModel(val savedStateHandle: SavedStateHandle) : ViewModel()
+}

workflow-ui/core-android/src/main/java/com/squareup/workflow1/ui/AndroidRenderWorkflow.kt

@@ -1,5 +1,6 @@
 package com.squareup.workflow1.ui
 
+import androidx.annotation.VisibleForTesting
 import androidx.lifecycle.SavedStateHandle
 import com.squareup.workflow1.RuntimeConfig
 import com.squareup.workflow1.RuntimeConfigOptions
@@ -290,4 +291,15 @@ public fun <PropsT, OutputT, RenderingT> renderWorkflowIn(
     .stateIn(scope, Eagerly, renderingsAndSnapshots.value.rendering)
 }
 
-private const val KEY = "com.squareup.workflow1.ui.renderWorkflowIn-snapshot"
+/**
+ * Removes state added to the `savedStateHandle` argument of the Android-specific
+ * overload of [renderWorkflowIn]. For use in obscure cases like swapping between
+ * different Workflow runtimes in an app. Most apps will not use this function.
+ */
+@WorkflowUiExperimentalApi
+public fun SavedStateHandle.removeWorkflowState() {
+  remove<Any>(KEY)
+}
+
+@VisibleForTesting
+internal const val KEY = "com.squareup.workflow1.ui.renderWorkflowIn-snapshot"