diff --git a/quickstep/src/com/android/quickstep/util/SplitAnimationController.kt b/quickstep/src/com/android/quickstep/util/SplitAnimationController.kt
index 9b4c772..48f257b 100644
--- a/quickstep/src/com/android/quickstep/util/SplitAnimationController.kt
+++ b/quickstep/src/com/android/quickstep/util/SplitAnimationController.kt
@@ -65,6 +65,7 @@
 import com.android.launcher3.util.SplitConfigurationOptions.SplitSelectSource
 import com.android.launcher3.views.BaseDragLayer
 import com.android.quickstep.TaskViewUtils
+import com.android.quickstep.util.SplitScreenUtils.Companion.extractTopParentAndChildren
 import com.android.quickstep.views.FloatingAppPairView
 import com.android.quickstep.views.FloatingTaskView
 import com.android.quickstep.views.GroupedTaskView
@@ -986,35 +987,9 @@
         progressUpdater.setDuration(QuickstepTransitionManager.APP_LAUNCH_DURATION)
         progressUpdater.interpolator = Interpolators.EMPHASIZED
 
-        var rootCandidate: Change? = null
-
-        for (change in transitionInfo.changes) {
-            val taskInfo: RunningTaskInfo = change.taskInfo ?: continue
-
-            // TODO (b/316490565): Replace this logic when SplitBounds is available to
-            //  startAnimation() and we can know the precise taskIds of launching tasks.
-            if (
-                taskInfo.windowingMode == windowingMode &&
-                    (change.mode == TRANSIT_OPEN || change.mode == TRANSIT_TO_FRONT)
-            ) {
-                // Found one!
-                rootCandidate = change
-                break
-            }
-        }
-
-        // If we could not find a proper root candidate, something went wrong.
-        check(rootCandidate != null) { "Could not find a split root candidate" }
-
-        // Recurse up the tree until parent is null, then we've found our root.
-        var parentToken: WindowContainerToken? = rootCandidate.parent
-        while (parentToken != null) {
-            rootCandidate = transitionInfo.getChange(parentToken) ?: break
-            parentToken = rootCandidate.parent
-        }
-
-        // Make sure nothing weird happened, like getChange() returning null.
-        check(rootCandidate != null) { "Failed to find a root leash" }
+        val splitTree: Pair<Change, List<Change>>? = extractTopParentAndChildren(transitionInfo)
+        check(splitTree != null) { "Could not find a split root candidate" }
+        val rootCandidate = splitTree.first
 
         // Starting position is a 34% size tile centered in the middle of the screen.
         // Ending position is the full device screen.
diff --git a/quickstep/src/com/android/quickstep/util/SplitScreenUtils.kt b/quickstep/src/com/android/quickstep/util/SplitScreenUtils.kt
index 4005c5a..7787e30 100644
--- a/quickstep/src/com/android/quickstep/util/SplitScreenUtils.kt
+++ b/quickstep/src/com/android/quickstep/util/SplitScreenUtils.kt
@@ -46,8 +46,8 @@
          * Given a TransitionInfo, generates the tree structure for those changes and extracts out
          * the top most root and it's two immediate children. Changes can be provided in any order.
          *
-         * @return a [Pair] where first -> top most split root, second -> [List] of 2,
-         *   leftTop/bottomRight stage roots
+         * @return null if no root is found, otherwise a [Pair] where first -> top most split root,
+         *         second -> [List] of 2, leftTop/bottomRight stage roots
          */
         fun extractTopParentAndChildren(
             transitionInfo: TransitionInfo