diff --git a/quickstep/src/com/android/quickstep/TopTaskTracker.java b/quickstep/src/com/android/quickstep/TopTaskTracker.java
index 6f502d0ece..24cf72c718 100644
--- a/quickstep/src/com/android/quickstep/TopTaskTracker.java
+++ b/quickstep/src/com/android/quickstep/TopTaskTracker.java
@@ -127,23 +127,16 @@ public class TopTaskTracker extends ISplitScreenListener.Stub implements TaskSta
 
     @Override
     public void onTaskStageChanged(int taskId, @StageType int stage, boolean visible) {
-        // If task is not visible but we are tracking it, stop tracking it
-        if (!visible) {
+        // If a task is not visible anymore or has been moved to undefined, stop tracking it.
+        if (!visible || stage == SplitConfigurationOptions.STAGE_TYPE_UNDEFINED) {
             if (mMainStagePosition.taskId == taskId) {
-                resetTaskId(mMainStagePosition);
+                mMainStagePosition.taskId = INVALID_TASK_ID;
             } else if (mSideStagePosition.taskId == taskId) {
-                resetTaskId(mSideStagePosition);
+                mSideStagePosition.taskId = INVALID_TASK_ID;
             } // else it's an un-tracked child
             return;
         }
 
-        // If stage has moved to undefined, stop tracking the task
-        if (stage == SplitConfigurationOptions.STAGE_TYPE_UNDEFINED) {
-            resetTaskId(taskId == mMainStagePosition.taskId
-                    ? mMainStagePosition : mSideStagePosition);
-            return;
-        }
-
         if (stage == SplitConfigurationOptions.STAGE_TYPE_MAIN) {
             mMainStagePosition.taskId = taskId;
         } else {
@@ -161,10 +154,6 @@ public class TopTaskTracker extends ISplitScreenListener.Stub implements TaskSta
         mPinnedTaskId = INVALID_TASK_ID;
     }
 
-    private void resetTaskId(SplitStageInfo taskPosition) {
-        taskPosition.taskId = -1;
-    }
-
     /**
      * @return index 0 will be task in left/top position, index 1 in right/bottom position.
      *         Will return empty array if device is not in staged split