Have each Fragment register its own back pressed callback

By splitting the OnBackPressedCallback apart at
a FragmentManager level, we can ensure that the
appropriate Lifecycle is used for each callback,
allowing developers to interleave the Fragment
back stack management with their own custom
callbacks.

Test: updated PrimaryNavFragmentTest tests
BUG: 111598096
Change-Id: I22cc38794c1ee20da370f1843d667b34a9529c22
diff --git a/activity/api/1.0.0-alpha07.txt b/activity/api/1.0.0-alpha07.txt
index 7b3c69f..ba5e733 100644
--- a/activity/api/1.0.0-alpha07.txt
+++ b/activity/api/1.0.0-alpha07.txt
@@ -1,7 +1,7 @@
 // Signature format: 3.0
 package androidx.activity {
 
-  public class ComponentActivity extends androidx.core.app.ComponentActivity implements androidx.lifecycle.LifecycleOwner androidx.savedstate.SavedStateRegistryOwner androidx.lifecycle.ViewModelStoreOwner {
+  public class ComponentActivity extends androidx.core.app.ComponentActivity implements androidx.lifecycle.LifecycleOwner androidx.activity.OnBackPressedDispatcherOwner androidx.savedstate.SavedStateRegistryOwner androidx.lifecycle.ViewModelStoreOwner {
     ctor public ComponentActivity();
     ctor @ContentView public ComponentActivity(@LayoutRes int);
     method @Deprecated public void addOnBackPressedCallback(androidx.activity.OnBackPressedCallback);
@@ -31,5 +31,9 @@
     method @MainThread public void onBackPressed();
   }
 
+  public interface OnBackPressedDispatcherOwner extends androidx.lifecycle.LifecycleOwner {
+    method public androidx.activity.OnBackPressedDispatcher getOnBackPressedDispatcher();
+  }
+
 }
 
diff --git a/activity/api/current.txt b/activity/api/current.txt
index 7b3c69f..ba5e733 100644
--- a/activity/api/current.txt
+++ b/activity/api/current.txt
@@ -1,7 +1,7 @@
 // Signature format: 3.0
 package androidx.activity {
 
-  public class ComponentActivity extends androidx.core.app.ComponentActivity implements androidx.lifecycle.LifecycleOwner androidx.savedstate.SavedStateRegistryOwner androidx.lifecycle.ViewModelStoreOwner {
+  public class ComponentActivity extends androidx.core.app.ComponentActivity implements androidx.lifecycle.LifecycleOwner androidx.activity.OnBackPressedDispatcherOwner androidx.savedstate.SavedStateRegistryOwner androidx.lifecycle.ViewModelStoreOwner {
     ctor public ComponentActivity();
     ctor @ContentView public ComponentActivity(@LayoutRes int);
     method @Deprecated public void addOnBackPressedCallback(androidx.activity.OnBackPressedCallback);
@@ -31,5 +31,9 @@
     method @MainThread public void onBackPressed();
   }
 
+  public interface OnBackPressedDispatcherOwner extends androidx.lifecycle.LifecycleOwner {
+    method public androidx.activity.OnBackPressedDispatcher getOnBackPressedDispatcher();
+  }
+
 }
 
diff --git a/activity/src/main/java/androidx/activity/ComponentActivity.java b/activity/src/main/java/androidx/activity/ComponentActivity.java
index 46ae98f..fde474e 100644
--- a/activity/src/main/java/androidx/activity/ComponentActivity.java
+++ b/activity/src/main/java/androidx/activity/ComponentActivity.java
@@ -49,7 +49,8 @@
 public class ComponentActivity extends androidx.core.app.ComponentActivity implements
         LifecycleOwner,
         ViewModelStoreOwner,
-        SavedStateRegistryOwner {
+        SavedStateRegistryOwner,
+        OnBackPressedDispatcherOwner {
 
     static final class NonConfigurationInstances {
         Object custom;
@@ -287,6 +288,7 @@
      * @return The {@link OnBackPressedDispatcher} associated with this ComponentActivity.
      */
     @NonNull
+    @Override
     public final OnBackPressedDispatcher getOnBackPressedDispatcher() {
         return mOnBackPressedDispatcher;
     }
diff --git a/activity/src/main/java/androidx/activity/OnBackPressedDispatcherOwner.java b/activity/src/main/java/androidx/activity/OnBackPressedDispatcherOwner.java
new file mode 100644
index 0000000..e07b3f9
--- /dev/null
+++ b/activity/src/main/java/androidx/activity/OnBackPressedDispatcherOwner.java
@@ -0,0 +1,40 @@
+/*
+ * Copyright 2019 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package androidx.activity;
+
+import androidx.annotation.NonNull;
+import androidx.lifecycle.LifecycleOwner;
+
+/**
+ * A class that has an {@link OnBackPressedDispatcher} that allows you to register a
+ * {@link OnBackPressedCallback} for handling the system back button.
+ * <p>
+ * It is expected that classes that implement this interface route the system back button
+ * to the dispatcher
+ *
+ * @see OnBackPressedDispatcher
+ */
+public interface OnBackPressedDispatcherOwner extends LifecycleOwner {
+
+    /**
+     * Retrieve the {@link OnBackPressedDispatcher} that should handle the system back button.
+     *
+     * @return The {@link OnBackPressedDispatcher}.
+     */
+    @NonNull
+    OnBackPressedDispatcher getOnBackPressedDispatcher();
+}
diff --git a/fragment/src/androidTest/java/androidx/fragment/app/PrimaryNavFragmentTest.kt b/fragment/src/androidTest/java/androidx/fragment/app/PrimaryNavFragmentTest.kt
index 2b06bf5..b1ed36a 100644
--- a/fragment/src/androidTest/java/androidx/fragment/app/PrimaryNavFragmentTest.kt
+++ b/fragment/src/androidTest/java/androidx/fragment/app/PrimaryNavFragmentTest.kt
@@ -16,6 +16,7 @@
 
 package androidx.fragment.app
 
+import android.app.Activity
 import androidx.fragment.app.test.EmptyFragmentTestActivity
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
@@ -58,10 +59,7 @@
             .that(cfm.backStackEntryCount)
             .isEqualTo(1)
 
-        // Should execute the pop for the child fragmentmanager
-        assertWithMessage("popBackStackImmediate returned no action performed")
-            .that(popBackStackImmediate(fm))
-            .isTrue()
+        activityRule.onBackPressed()
 
         assertWithMessage("child transaction still on back stack")
             .that(cfm.backStackEntryCount)
@@ -93,7 +91,7 @@
             .that(fm.primaryNavigationFragment)
             .isNull()
 
-        popBackStackImmediate(fm)
+        activityRule.onBackPressed()
 
         assertWithMessage("primary nav fragment was not restored on pop")
             .that(fm.primaryNavigationFragment)
@@ -112,7 +110,7 @@
             .that(fm.primaryNavigationFragment)
             .isSameAs(strictFragment2)
 
-        popBackStackImmediate(fm)
+        activityRule.onBackPressed()
 
         assertWithMessage("primary nav fragment not restored on pop")
             .that(fm.primaryNavigationFragment)
@@ -127,7 +125,7 @@
         assertWithMessage("primary nav fragment not retained when set again in new transaction")
             .that(fm.primaryNavigationFragment)
             .isSameAs(strictFragment1)
-        popBackStackImmediate(fm)
+        activityRule.onBackPressed()
 
         assertWithMessage(
             "same primary nav fragment not retained when set primary nav transaction popped")
@@ -162,7 +160,7 @@
             .that(fm.primaryNavigationFragment)
             .isNull()
 
-        popBackStackImmediate(fm)
+        activityRule.onBackPressed()
 
         assertWithMessage("primary nav fragment not restored after popping replace")
             .that(fm.primaryNavigationFragment)
@@ -188,7 +186,7 @@
             .that(fm.primaryNavigationFragment)
             .isSameAs(strictFragment2)
 
-        popBackStackImmediate(fm)
+        activityRule.onBackPressed()
 
         assertWithMessage("primary nav fragment not null after popping replace")
             .that(fm.primaryNavigationFragment)
@@ -199,9 +197,7 @@
         activityRule.runOnUiThread { fm.executePendingTransactions() }
     }
 
-    private fun popBackStackImmediate(fm: FragmentManager): Boolean {
-        var popped = false
-        activityRule.runOnUiThread { popped = fm.popBackStackImmediate() }
-        return popped
+    private fun ActivityTestRule<out Activity>.onBackPressed() = runOnUiThread {
+        activity.onBackPressed()
     }
 }
diff --git a/fragment/src/main/java/androidx/fragment/app/FragmentActivity.java b/fragment/src/main/java/androidx/fragment/app/FragmentActivity.java
index d70aaca..1718a2a 100644
--- a/fragment/src/main/java/androidx/fragment/app/FragmentActivity.java
+++ b/fragment/src/main/java/androidx/fragment/app/FragmentActivity.java
@@ -35,7 +35,8 @@
 import android.view.Window;
 
 import androidx.activity.ComponentActivity;
-import androidx.activity.OnBackPressedCallback;
+import androidx.activity.OnBackPressedDispatcher;
+import androidx.activity.OnBackPressedDispatcherOwner;
 import androidx.annotation.CallSuper;
 import androidx.annotation.ContentView;
 import androidx.annotation.LayoutRes;
@@ -116,7 +117,6 @@
      */
     public FragmentActivity() {
         super();
-        init();
     }
 
     /**
@@ -132,47 +132,6 @@
     @ContentView
     public FragmentActivity(@LayoutRes int contentLayoutId) {
         super(contentLayoutId);
-        init();
-    }
-
-    private void init() {
-        // Route onBackPressed() callbacks to the FragmentManager
-        getOnBackPressedDispatcher().addCallback(this, new OnBackPressedCallback(true) {
-            @Override
-            public boolean isEnabled() {
-                FragmentManager fragmentManager = mFragments.getSupportFragmentManager();
-                // Either this FragmentManager needs to have a back stack
-                // or its primary navigation Fragment needs to have one
-                return fragmentManager.getBackStackEntryCount() > 0
-                        || hasPrimaryNavigationBackStack(fragmentManager);
-            }
-
-            /**
-             * Recursively check down the FragmentManager hierarchy of primary
-             * navigation Fragments for a FragmentManager that has a back stack
-             * that can be popped.
-             */
-            private boolean hasPrimaryNavigationBackStack(
-                    @NonNull FragmentManager fragmentManager) {
-                Fragment primaryNavFragment = fragmentManager.getPrimaryNavigationFragment();
-                if (primaryNavFragment == null) {
-                    return false;
-                }
-                FragmentManager primaryNavFragmentManager =
-                        primaryNavFragment.peekChildFragmentManager();
-                if (primaryNavFragmentManager == null) {
-                    return false;
-                }
-                return primaryNavFragmentManager.getBackStackEntryCount() > 0
-                        || hasPrimaryNavigationBackStack(primaryNavFragmentManager);
-            }
-
-            @Override
-            public void handleOnBackPressed() {
-                FragmentManager fragmentManager = mFragments.getSupportFragmentManager();
-                fragmentManager.popBackStackImmediate();
-            }
-        });
     }
 
     // ------------------------------------------------------------------------
@@ -891,18 +850,31 @@
         }
     }
 
-    class HostCallbacks extends FragmentHostCallback<FragmentActivity>
-            implements ViewModelStoreOwner {
+    class HostCallbacks extends FragmentHostCallback<FragmentActivity> implements
+            ViewModelStoreOwner,
+            OnBackPressedDispatcherOwner {
         public HostCallbacks() {
             super(FragmentActivity.this /*fragmentActivity*/);
         }
 
         @NonNull
         @Override
+        public Lifecycle getLifecycle() {
+            return FragmentActivity.this.getLifecycle();
+        }
+
+        @NonNull
+        @Override
         public ViewModelStore getViewModelStore() {
             return FragmentActivity.this.getViewModelStore();
         }
 
+        @NonNull
+        @Override
+        public OnBackPressedDispatcher getOnBackPressedDispatcher() {
+            return FragmentActivity.this.getOnBackPressedDispatcher();
+        }
+
         @Override
         public void onDump(@NonNull String prefix, @Nullable FileDescriptor fd,
                 @NonNull PrintWriter writer, @Nullable String[] args) {
diff --git a/fragment/src/main/java/androidx/fragment/app/FragmentManagerImpl.java b/fragment/src/main/java/androidx/fragment/app/FragmentManagerImpl.java
index 6213072..fed4bd3 100644
--- a/fragment/src/main/java/androidx/fragment/app/FragmentManagerImpl.java
+++ b/fragment/src/main/java/androidx/fragment/app/FragmentManagerImpl.java
@@ -45,6 +45,9 @@
 import android.view.animation.ScaleAnimation;
 import android.view.animation.Transformation;
 
+import androidx.activity.OnBackPressedCallback;
+import androidx.activity.OnBackPressedDispatcher;
+import androidx.activity.OnBackPressedDispatcherOwner;
 import androidx.annotation.NonNull;
 import androidx.annotation.Nullable;
 import androidx.collection.ArraySet;
@@ -52,6 +55,7 @@
 import androidx.core.util.LogWriter;
 import androidx.core.view.OneShotPreDrawListener;
 import androidx.lifecycle.Lifecycle;
+import androidx.lifecycle.LifecycleOwner;
 import androidx.lifecycle.ViewModelStore;
 import androidx.lifecycle.ViewModelStoreOwner;
 
@@ -2463,11 +2467,54 @@
     }
 
     public void attachController(@NonNull FragmentHostCallback host,
-                                 @NonNull FragmentContainer container, @Nullable Fragment parent) {
+            @NonNull FragmentContainer container, @Nullable final Fragment parent) {
         if (mHost != null) throw new IllegalStateException("Already attached");
         mHost = host;
         mContainer = container;
         mParent = parent;
+        // Set up the OnBackPressedCallback
+        if (host instanceof OnBackPressedDispatcherOwner) {
+            OnBackPressedDispatcherOwner dispatcherOwner = ((OnBackPressedDispatcherOwner) host);
+            OnBackPressedDispatcher dispatcher = dispatcherOwner.getOnBackPressedDispatcher();
+            LifecycleOwner owner = parent != null ? parent : dispatcherOwner;
+            dispatcher.addCallback(owner, new OnBackPressedCallback(true) {
+                @Override
+                public boolean isEnabled() {
+                    // This FragmentManager needs to have a back stack for this to be enabled
+                    // and the parent fragment, if it exists, needs to be the primary navigation
+                    // fragment.
+                    return getBackStackEntryCount() > 0 && isPrimaryNavigation(parent);
+                }
+
+                /**
+                 * Recursively check up the FragmentManager hierarchy of primary
+                 * navigation Fragments to ensure that all of the parent Fragments are the
+                 * primary navigation Fragment for their associated FragmentManager
+                 */
+                private boolean isPrimaryNavigation(@Nullable Fragment parent) {
+                    // If the parent is null, then we're at the root host
+                    // and we're always the primary navigation
+                    if (parent == null) {
+                        return true;
+                    }
+                    FragmentManagerImpl parentFragmentManager = parent.mFragmentManager;
+                    Fragment primaryNavigationFragment = parentFragmentManager
+                            .getPrimaryNavigationFragment();
+                    // The parent Fragment needs to be the primary navigation Fragment
+                    // and, if it has a parent itself, that parent also needs to be
+                    // the primary navigation fragment, recursively up the stack
+                    return parent == primaryNavigationFragment
+                            && isPrimaryNavigation(parentFragmentManager.mParent);
+                }
+
+                @Override
+                public void handleOnBackPressed() {
+                    popBackStackImmediate();
+                }
+            });
+        }
+
+        // Get the FragmentManagerViewModel
         if (parent != null) {
             mNonConfig = parent.mFragmentManager.getChildNonConfig(parent);
         } else if (host instanceof ViewModelStoreOwner) {