Better support for View.findViewTreeOnBackPressedDispatcherOwner.

`fun View.findViewTreeOnBackPressedDispatcherOwner()` is AndroidX's preferred entry point to the exciting new world of `OnBackPressedDispatcherOwner`. With this PR we support it explicitly, in particular taking care to ensure that it can be called by newly constructed views before they have been attached to a parent. That is, we take care to make eager calls `View.setViewTreeSavedStateRegistryOwner` on every view we build.

We accomplish this mainly by riding the rails previously laid down via `WorkflowLifecycleOwner.installOn`, which now requires an `OnBackPressedDispatcherOwner` parameter. (This is the method used by `WorkflowViewStub` _et al_ to ensure that we're managing the JetPack Lifecycle correctly, and `OnBackPressedDispatcherOwner` is just another piece of that puzzle.)

In aid of that, we introduce a new `ViewEnvironmentKey`, `OnBackPressedDispatcherOwnerKey`. It is initialized by `WorkflowLayout`, our new `ComponentDialog.setContent` extension, and `@Composable fun WorkflowRendering()`. This key is not intended for use by feature code, it's more of an implementation detail that has to stay public to allow custom containers to be built. It ensures that `WorkflowViewStub` and friends will have access to the correct `OnBackPressedDispatcherOwner` before they have access to a parent view.

TODO: add tests of `@Composable fun BackHandler()`, in both activity and dialog windows.

workflow-ui/compose/build.gradle.kts

@@ -47,6 +47,7 @@ dependencies {
   api(project(":workflow-ui:core-common"))
 
   implementation(composeBom)
+  implementation(libs.androidx.activity.compose)
   implementation(libs.androidx.compose.foundation.layout)
   implementation(libs.androidx.compose.runtime.saveable)
   implementation(libs.androidx.compose.ui)

workflow-ui/compose/src/main/java/com/squareup/workflow1/ui/compose/WorkflowRendering.kt

@@ -3,6 +3,9 @@
 package com.squareup.workflow1.ui.compose
 
 import android.view.View
+import androidx.activity.OnBackPressedDispatcherOwner
+import androidx.activity.compose.LocalOnBackPressedDispatcherOwner
+import androidx.activity.setViewTreeOnBackPressedDispatcherOwner
 import androidx.compose.foundation.layout.Box
 import androidx.compose.runtime.Composable
 import androidx.compose.runtime.CompositionLocalProvider
@@ -20,6 +23,7 @@ import androidx.lifecycle.LifecycleOwner
 import androidx.lifecycle.LifecycleRegistry
 import androidx.lifecycle.ViewTreeLifecycleOwner
 import com.squareup.workflow1.ui.Compatible
+import com.squareup.workflow1.ui.OnBackPressedDispatcherOwnerKey
 import com.squareup.workflow1.ui.Screen
 import com.squareup.workflow1.ui.ScreenViewFactory
 import com.squareup.workflow1.ui.ScreenViewFactoryFinder
@@ -170,6 +174,10 @@ private fun <ScreenT : Screen> ScreenViewFactory<ScreenT>.asComposeViewFactory()
      *    we already have the correct one.
      *  - Propagate the current [LifecycleOwner] from [LocalLifecycleOwner] by setting it as the
      *    [ViewTreeLifecycleOwner] on the view.
+     *  - Propagate the current [OnBackPressedDispatcherOwner] from either
+     *    [LocalOnBackPressedDispatcherOwner] or the [viewEnvironment],
+     *    both on the [AndroidView] via [setViewTreeOnBackPressedDispatcherOwner],
+     *    and in the [ViewEnvironment] for use by any nested [WorkflowViewStub]
      *
      * Like `WorkflowViewStub`, this function uses the [originalFactory] to create and memoize a
      * [View] to display the [rendering], keeps it updated with the latest [rendering] and
@@ -181,17 +189,35 @@ private fun <ScreenT : Screen> ScreenViewFactory<ScreenT>.asComposeViewFactory()
     ) {
       val lifecycleOwner = LocalLifecycleOwner.current
 
+      // Make sure any nested WorkflowViewStub will be able to propagate the
+      // OnBackPressedDispatcherOwner, if we found one. No need to fail fast here.
+      // It's only an issue if someone tries to use it, and the error message
+      // at those call sites should be clear enough.
+      val onBackOrNull = LocalOnBackPressedDispatcherOwner.current
+        ?: viewEnvironment.map[OnBackPressedDispatcherOwnerKey] as? OnBackPressedDispatcherOwner
+
+      val envWithOnBack = onBackOrNull
+        ?.let { viewEnvironment + (OnBackPressedDispatcherOwnerKey to it) }
+        ?: viewEnvironment
+
       AndroidView(
         factory = { context ->
+
           // We pass in a null container because the container isn't a View, it's a composable. The
           // compose machinery will generate an intermediate view that it ends up adding this to but
           // we don't have access to that.
-          originalFactory.startShowing(rendering, viewEnvironment, context, container = null)
+          originalFactory
+            .startShowing(rendering, envWithOnBack, context, container = null)
             .let { viewHolder ->
               // Put the viewHolder in a tag so that we can find it in the update lambda, below.
               viewHolder.view.setTag(R.id.workflow_screen_view_holder, viewHolder)
-              // Unfortunately AndroidView doesn't propagate this itself.
+
+              // Unfortunately AndroidView doesn't propagate these itself.
               ViewTreeLifecycleOwner.set(viewHolder.view, lifecycleOwner)
+              onBackOrNull?.let {
+                viewHolder.view.setViewTreeOnBackPressedDispatcherOwner(it)
+              }
+
               // We don't propagate the (non-compose) SavedStateRegistryOwner, or the (compose)
               // SaveableStateRegistry, because currently all our navigation is implemented as
               // Android views, which ensures there is always an Android view between any state
@@ -206,7 +232,7 @@ private fun <ScreenT : Screen> ScreenViewFactory<ScreenT>.asComposeViewFactory()
           @Suppress("UNCHECKED_CAST")
           val viewHolder =
             view.getTag(R.id.workflow_screen_view_holder) as ScreenViewHolder<ScreenT>
-          viewHolder.show(rendering, viewEnvironment)
+          viewHolder.show(rendering, envWithOnBack)
         }
       )
     }

workflow-ui/container-android/src/main/java/com/squareup/workflow1/ui/modal/ModalContainer.kt

@@ -13,13 +13,15 @@ import android.view.View
 import android.view.ViewGroup
 import android.view.ViewGroup.LayoutParams.MATCH_PARENT
 import android.widget.FrameLayout
+import androidx.activity.OnBackPressedDispatcherOwner
 import androidx.lifecycle.DefaultLifecycleObserver
 import androidx.lifecycle.Lifecycle
 import androidx.lifecycle.LifecycleOwner
 import com.squareup.workflow1.ui.Compatible
 import com.squareup.workflow1.ui.ViewEnvironment
 import com.squareup.workflow1.ui.WorkflowUiExperimentalApi
 import com.squareup.workflow1.ui.WorkflowViewStub
+import com.squareup.workflow1.ui.androidx.WorkflowAndroidXSupport.onBackPressedDispatcherOwner
 import com.squareup.workflow1.ui.androidx.WorkflowAndroidXSupport.stateRegistryOwnerFromViewTreeOrContext
 import com.squareup.workflow1.ui.androidx.WorkflowLifecycleOwner
 import com.squareup.workflow1.ui.androidx.WorkflowSavedStateRegistryAggregator
@@ -88,6 +90,8 @@ public abstract class ModalContainer<ModalRenderingT : Any> @JvmOverloads constr
             // any, and so we can use our lifecycle to destroy-on-detach the dialog hierarchy.
             WorkflowLifecycleOwner.installOn(
               dialogView,
+              (ref.dialog as? OnBackPressedDispatcherOwner)
+                ?: viewEnvironment.onBackPressedDispatcherOwner(this),
               findParentLifecycle = { parentLifecycleOwner.lifecycle }
             )
             // Ensure that each dialog has its own SavedStateRegistryOwner,

workflow-ui/container-android/src/main/java/com/squareup/workflow1/ui/modal/ModalViewContainer.kt

@@ -18,11 +18,11 @@ import com.squareup.workflow1.ui.ScreenViewHolder
 import com.squareup.workflow1.ui.ViewEnvironment
 import com.squareup.workflow1.ui.ViewRegistry
 import com.squareup.workflow1.ui.WorkflowUiExperimentalApi
+import com.squareup.workflow1.ui.androidx.WorkflowAndroidXSupport.onBackPressedDispatcherOwnerOrNull
 import com.squareup.workflow1.ui.asScreen
 import com.squareup.workflow1.ui.bindShowRendering
 import com.squareup.workflow1.ui.container.BackButtonScreen
 import com.squareup.workflow1.ui.modal.ModalViewContainer.Companion.binding
-import com.squareup.workflow1.ui.onBackPressedDispatcherOwnerOrNull
 import com.squareup.workflow1.ui.show
 import com.squareup.workflow1.ui.startShowing
 import com.squareup.workflow1.ui.toViewFactory
@@ -93,7 +93,7 @@ public open class ModalViewContainer @JvmOverloads constructor(
 
         setOnKeyListener { _, keyCode, keyEvent ->
           if (keyCode == KeyEvent.KEYCODE_BACK && keyEvent.action == ACTION_UP) {
-            viewHolder.view.context.onBackPressedDispatcherOwnerOrNull()
+            viewHolder.view.onBackPressedDispatcherOwnerOrNull()
               ?.onBackPressedDispatcher
               ?.let {
                 if (it.hasEnabledCallbacks()) it.onBackPressed()

workflow-ui/core-android/api/core-android.api

@@ -33,7 +33,6 @@ public final class com/squareup/workflow1/ui/BackButtonScreen : com/squareup/wor
 
 public final class com/squareup/workflow1/ui/BackPressHandlerKt {
 	public static final fun getBackPressedHandler (Landroid/view/View;)Lkotlin/jvm/functions/Function0;
-	public static final fun onBackPressedDispatcherOwnerOrNull (Landroid/content/Context;)Landroidx/activity/OnBackPressedDispatcherOwner;
 	public static final fun setBackPressedHandler (Landroid/view/View;Lkotlin/jvm/functions/Function0;)V
 }
 
@@ -77,6 +76,12 @@ public final class com/squareup/workflow1/ui/LayoutScreenViewFactory : com/squar
 	public fun getType ()Lkotlin/reflect/KClass;
 }
 
+public final class com/squareup/workflow1/ui/OnBackPressedDispatcherOwnerKey : com/squareup/workflow1/ui/ViewEnvironmentKey {
+	public static final field INSTANCE Lcom/squareup/workflow1/ui/OnBackPressedDispatcherOwnerKey;
+	public fun getDefault ()Landroidx/activity/OnBackPressedDispatcherOwner;
+	public synthetic fun getDefault ()Ljava/lang/Object;
+}
+
 public final class com/squareup/workflow1/ui/ParcelableTextController : android/os/Parcelable, com/squareup/workflow1/ui/TextController {
 	public static final field CREATOR Lcom/squareup/workflow1/ui/ParcelableTextController$CREATOR;
 	public synthetic fun <init> (Landroid/os/Parcel;Lkotlin/jvm/internal/DefaultConstructorMarker;)V
@@ -275,6 +280,8 @@ public final class com/squareup/workflow1/ui/androidx/WorkflowAndroidXSupport {
 	public static final field INSTANCE Lcom/squareup/workflow1/ui/androidx/WorkflowAndroidXSupport;
 	public final fun lifecycleOwnerFromContext (Landroid/content/Context;)Landroidx/lifecycle/LifecycleOwner;
 	public final fun lifecycleOwnerFromViewTreeOrContextOrNull (Landroid/view/View;)Landroidx/lifecycle/LifecycleOwner;
+	public final fun onBackPressedDispatcherOwner (Lcom/squareup/workflow1/ui/ViewEnvironment;Landroid/view/View;)Landroidx/activity/OnBackPressedDispatcherOwner;
+	public final fun onBackPressedDispatcherOwnerOrNull (Landroid/view/View;)Landroidx/activity/OnBackPressedDispatcherOwner;
 	public final fun stateRegistryOwnerFromViewTreeOrContext (Landroid/view/View;)Landroidx/savedstate/SavedStateRegistryOwner;
 }
 
@@ -285,8 +292,8 @@ public abstract interface class com/squareup/workflow1/ui/androidx/WorkflowLifec
 
 public final class com/squareup/workflow1/ui/androidx/WorkflowLifecycleOwner$Companion {
 	public final fun get (Landroid/view/View;)Lcom/squareup/workflow1/ui/androidx/WorkflowLifecycleOwner;
-	public final fun installOn (Landroid/view/View;Lkotlin/jvm/functions/Function1;)V
-	public static synthetic fun installOn$default (Lcom/squareup/workflow1/ui/androidx/WorkflowLifecycleOwner$Companion;Landroid/view/View;Lkotlin/jvm/functions/Function1;ILjava/lang/Object;)V
+	public final fun installOn (Landroid/view/View;Landroidx/activity/OnBackPressedDispatcherOwner;Lkotlin/jvm/functions/Function1;)V
+	public static synthetic fun installOn$default (Lcom/squareup/workflow1/ui/androidx/WorkflowLifecycleOwner$Companion;Landroid/view/View;Landroidx/activity/OnBackPressedDispatcherOwner;Lkotlin/jvm/functions/Function1;ILjava/lang/Object;)V
 }
 
 public final class com/squareup/workflow1/ui/androidx/WorkflowSavedStateRegistryAggregator {

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

@@ -5,6 +5,9 @@ import android.os.Build.VERSION_CODES
 import android.os.Parcel
 import android.os.Parcelable
 import android.util.SparseArray
+import androidx.activity.OnBackPressedDispatcher
+import androidx.activity.OnBackPressedDispatcherOwner
+import androidx.lifecycle.Lifecycle
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.platform.app.InstrumentationRegistry
 import com.google.common.truth.Truth.assertThat
@@ -32,6 +35,11 @@ internal class ViewStateCacheTest {
   private val instrumentation = InstrumentationRegistry.getInstrumentation()
   private val viewEnvironment = EMPTY
 
+  private val fakeOnBack = object : OnBackPressedDispatcherOwner {
+    override fun getLifecycle(): Lifecycle = error("")
+    override fun getOnBackPressedDispatcher(): OnBackPressedDispatcher = error("")
+  }
+
   private object AScreen : Screen
 
   @Test fun saves_and_restores_self() {
@@ -59,9 +67,7 @@ internal class ViewStateCacheTest {
         @Suppress("DEPRECATION")
         parcel.readParcelable(ViewStateCache.Saved::class.java.classLoader)!!
       }
-      ).let { restoredState ->
-      ViewStateCache().apply { restore(restoredState) }
-    }
+      ).let { restoredState -> ViewStateCache().apply { restore(restoredState) } }
 
     assertThat(restoredCache.equalsForTest(cache)).isTrue()
   }
@@ -125,7 +131,7 @@ internal class ViewStateCacheTest {
     // "Navigate" back to the first screen, restoring state.
     val firstViewRestored = ViewStateTestView(instrumentation.context).apply {
       id = 2
-      WorkflowLifecycleOwner.installOn(this)
+      WorkflowLifecycleOwner.installOn(this, fakeOnBack)
     }
     val firstHolderRestored =
       ScreenViewHolder<NamedScreen<*>>(EMPTY, firstViewRestored) { _, _ -> }.also {
@@ -194,7 +200,7 @@ internal class ViewStateCacheTest {
   ): ScreenViewHolder<NamedScreen<*>> {
     val view = ViewStateTestView(instrumentation.context).also { view ->
       id?.let { view.id = id }
-      WorkflowLifecycleOwner.installOn(view)
+      WorkflowLifecycleOwner.installOn(view, fakeOnBack)
     }
     return ScreenViewHolder<NamedScreen<*>>(EMPTY, view) { _, _ -> }.also {
       it.show(firstRendering, viewEnvironment)

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

@@ -1,15 +1,13 @@
 package com.squareup.workflow1.ui
 
-import android.content.Context
-import android.content.ContextWrapper
 import android.view.View
 import android.view.View.OnAttachStateChangeListener
 import androidx.activity.OnBackPressedCallback
-import androidx.activity.OnBackPressedDispatcherOwner
 import androidx.lifecycle.DefaultLifecycleObserver
 import androidx.lifecycle.Lifecycle
 import androidx.lifecycle.LifecycleOwner
 import androidx.lifecycle.ViewTreeLifecycleOwner
+import com.squareup.workflow1.ui.androidx.WorkflowAndroidXSupport.onBackPressedDispatcherOwnerOrNull
 
 /**
  * A function passed to [View.backPressedHandler], to be called if the back
@@ -87,7 +85,7 @@ private class AttachStateAndLifecycleObserver(
   private var lifecycleOrNull: Lifecycle? = null
 
   fun start() {
-    view.context.onBackPressedDispatcherOwnerOrNull()
+    view.onBackPressedDispatcherOwnerOrNull()
       ?.let { owner ->
         owner.onBackPressedDispatcher.addCallback(owner, onBackPressedCallback)
         view.addOnAttachStateChangeListener(this)
@@ -138,10 +136,3 @@ internal class NullableOnBackPressedCallback : OnBackPressedCallback(false) {
     handlerOrNull?.invoke()
   }
 }
-
-@WorkflowUiExperimentalApi
-public tailrec fun Context.onBackPressedDispatcherOwnerOrNull(): OnBackPressedDispatcherOwner? =
-  when (this) {
-    is OnBackPressedDispatcherOwner -> this
-    else -> (this as? ContextWrapper)?.baseContext?.onBackPressedDispatcherOwnerOrNull()
-  }

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

@@ -0,0 +1,15 @@
+package com.squareup.workflow1.ui
+
+import androidx.activity.OnBackPressedDispatcherOwner
+
+/**
+ * Used by container classes to ensure that
+ * [View.findViewTreeOnBackPressedDispatcherOwner][androidx.activity.findViewTreeOnBackPressedDispatcherOwner]
+ * works before new views are attached to their parents. Not intended for use by
+ * feature code.
+ */
+@WorkflowUiExperimentalApi
+public object OnBackPressedDispatcherOwnerKey :
+  ViewEnvironmentKey<OnBackPressedDispatcherOwner>() {
+  override val default: OnBackPressedDispatcherOwner get() = error("Unset")
+}

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

@@ -2,7 +2,6 @@ package com.squareup.workflow1.ui
 
 import android.content.Context
 import android.os.Build.VERSION
-import android.os.Build.VERSION_CODES
 import android.os.Parcel
 import android.os.Parcelable
 import android.os.Parcelable.Creator
@@ -17,6 +16,7 @@ import androidx.lifecycle.Lifecycle.State
 import androidx.lifecycle.Lifecycle.State.STARTED
 import androidx.lifecycle.coroutineScope
 import androidx.lifecycle.repeatOnLifecycle
+import com.squareup.workflow1.ui.androidx.WorkflowAndroidXSupport.onBackPressedDispatcherOwnerOrNull
 import kotlinx.coroutines.CoroutineDispatcher
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.Dispatchers
@@ -68,7 +68,7 @@ public class WorkflowLayout(
     rootScreen: Screen,
     environment: ViewEnvironment = ViewEnvironment.EMPTY
   ) {
-    showing.show(rootScreen, environment)
+    showing.show(rootScreen, environment.withOnBackDispatcher())
     restoredChildState?.let { restoredState ->
       restoredChildState = null
       showing.actual.restoreHierarchyState(restoredState)
@@ -127,7 +127,7 @@ public class WorkflowLayout(
     environment: ViewEnvironment
   ) {
     @Suppress("DEPRECATION")
-    showing.update(newRendering, environment)
+    showing.update(newRendering, environment.withOnBackDispatcher())
     restoredChildState?.let { restoredState ->
       restoredChildState = null
       showing.actual.restoreHierarchyState(restoredState)
@@ -242,6 +242,22 @@ public class WorkflowLayout(
     // Make a no-op call.
   }
 
+  /**
+   * Attempts to seed the [ViewEnvironment] with an [OnBackPressedDispatcherOwnerKey]
+   * value if one wasn't set already. We're priming the pump that our
+   * `ViewEnvironment.onBackPressedDispatcherOwner` call relies on.
+   */
+  private fun ViewEnvironment.withOnBackDispatcher(): ViewEnvironment {
+    val envWithOnBack = if (map.containsKey(OnBackPressedDispatcherOwnerKey)) {
+      this
+    } else {
+      this@WorkflowLayout.onBackPressedDispatcherOwnerOrNull()
+        ?.let { this@withOnBackDispatcher + (OnBackPressedDispatcherOwnerKey to it) }
+        ?: this
+    }
+    return envWithOnBack
+  }
+
   private class SavedState : BaseSavedState {
     constructor(
       superState: Parcelable?,
@@ -251,7 +267,7 @@ public class WorkflowLayout(
     }
 
     constructor(source: Parcel) : super(source) {
-      this.childState = if (VERSION.SDK_INT >= VERSION_CODES.TIRAMISU) {
+      this.childState = if (VERSION.SDK_INT >= 33) {
         source.readSparseArray(SavedState::class.java.classLoader, Parcelable::class.java)!!
       } else {
         @Suppress("DEPRECATION")

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

@@ -9,6 +9,7 @@ import androidx.annotation.IdRes
 import androidx.savedstate.SavedStateRegistryOwner
 import androidx.savedstate.findViewTreeSavedStateRegistryOwner
 import androidx.savedstate.setViewTreeSavedStateRegistryOwner
+import com.squareup.workflow1.ui.androidx.WorkflowAndroidXSupport.onBackPressedDispatcherOwner
 import com.squareup.workflow1.ui.androidx.WorkflowLifecycleOwner
 
 /**
@@ -231,10 +232,10 @@ public class WorkflowViewStub @JvmOverloads constructor(
 
     holder = rendering.toViewFactory(viewEnvironment)
       .startShowing(rendering, viewEnvironment, parent.context, parent) { view, doStart ->
-        WorkflowLifecycleOwner.installOn(view)
+        WorkflowLifecycleOwner.installOn(view, viewEnvironment.onBackPressedDispatcherOwner(parent))
         doStart()
-      }.also {
-        val newView = it.view
+      }.apply {
+        val newView = view
 
         if (inflatedId != NO_ID) newView.id = inflatedId
         if (updatesVisibility) newView.visibility = visibility