Merge pull request #1027 from square/ray/2-install-on-back-dispatcher

2/4 Better support for `View.findViewTreeOnBackPressedDispatcherOwner`.

workflow-ui/compose-tooling/dependencies/releaseRuntimeClasspath.txt

@@ -3,6 +3,7 @@
 :workflow-ui:compose
 :workflow-ui:core-android
 :workflow-ui:core-common
+androidx.activity:activity-compose:1.6.1
 androidx.activity:activity-ktx:1.6.1
 androidx.activity:activity:1.6.1
 androidx.annotation:annotation-experimental:1.1.0

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/dependencies/releaseRuntimeClasspath.txt

@@ -2,6 +2,7 @@
 :workflow-runtime
 :workflow-ui:core-android
 :workflow-ui:core-common
+androidx.activity:activity-compose:1.6.1
 androidx.activity:activity-ktx:1.6.1
 androidx.activity:activity:1.6.1
 androidx.annotation:annotation-experimental:1.1.0

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