diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarKeyguardController.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarKeyguardController.java
index edcd4c8..03d08eb 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarKeyguardController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarKeyguardController.java
@@ -1,17 +1,17 @@
 package com.android.launcher3.taskbar;
 
 import static com.android.launcher3.AbstractFloatingView.TYPE_ALL;
-import static com.android.systemui.shared.system.QuickStepContract.SCREEN_STATE_OFF;
+import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_AWAKE;
 import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_BACK_DISABLED;
 import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_BOUNCER_SHOWING;
 import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_DEVICE_DOZING;
 import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_DEVICE_DREAMING;
 import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_HOME_DISABLED;
 import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_OVERVIEW_DISABLED;
-import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_SCREEN_ON;
-import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_SCREEN_STATE_MASK;
 import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_STATUS_BAR_KEYGUARD_SHOWING;
 import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_STATUS_BAR_KEYGUARD_SHOWING_OCCLUDED;
+import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_WAKEFULNESS_MASK;
+import static com.android.systemui.shared.system.QuickStepContract.WAKEFULNESS_ASLEEP;
 
 import android.app.KeyguardManager;
 
@@ -29,7 +29,7 @@
             | SYSUI_STATE_STATUS_BAR_KEYGUARD_SHOWING | SYSUI_STATE_DEVICE_DOZING
             | SYSUI_STATE_OVERVIEW_DISABLED | SYSUI_STATE_HOME_DISABLED
             | SYSUI_STATE_BACK_DISABLED | SYSUI_STATE_STATUS_BAR_KEYGUARD_SHOWING_OCCLUDED
-            | SYSUI_STATE_SCREEN_STATE_MASK;
+            | SYSUI_STATE_WAKEFULNESS_MASK;
 
     // If any of these SysUi flags (via QuickstepContract) is set, the device to be considered
     // locked.
@@ -75,13 +75,13 @@
                 keyguardOccluded);
         updateIconsForBouncer();
 
-        boolean screenOffOrTransitioningOff = (systemUiStateFlags & SYSUI_STATE_SCREEN_ON) == 0;
-        boolean closeFloatingViews = keyguardShowing || screenOffOrTransitioningOff;
+        boolean asleepOrGoingToSleep = (systemUiStateFlags & SYSUI_STATE_AWAKE) == 0;
+        boolean closeFloatingViews = keyguardShowing || asleepOrGoingToSleep;
 
         if (closeFloatingViews) {
-            // animate the closing of the views, unless the screen is already fully turned off.
+            // animate the closing of the views, unless the screen is already asleep.
             boolean animateViewClosing =
-                    (systemUiStateFlags & SYSUI_STATE_SCREEN_STATE_MASK) != SCREEN_STATE_OFF;
+                    (systemUiStateFlags & SYSUI_STATE_WAKEFULNESS_MASK) != WAKEFULNESS_ASLEEP;
             AbstractFloatingView.closeOpenViews(mContext, animateViewClosing, TYPE_ALL);
         }
     }
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarLauncherStateController.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarLauncherStateController.java
index cf8148e..b98ea81 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarLauncherStateController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarLauncherStateController.java
@@ -22,7 +22,7 @@
 import static com.android.launcher3.util.FlagDebugUtils.appendFlag;
 import static com.android.launcher3.util.FlagDebugUtils.formatFlagChange;
 import static com.android.systemui.animation.Interpolators.EMPHASIZED;
-import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_SCREEN_ON;
+import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_AWAKE;
 
 import android.animation.Animator;
 import android.animation.AnimatorListenerAdapter;
@@ -81,17 +81,17 @@
      *
      * This is cleared as soon as the screen begins to transition off.
      */
-    private static final int FLAG_SCREEN_ON = 1 << 3;
+    private static final int FLAG_AWAKE = 1 << 3;
 
     /**
-     * Captures whether the launcher was active at the time the FLAG_SCREEN_ON was cleared.
-     * Always cleared when FLAG_SCREEN_ON is set.
+     * Captures whether the launcher was active at the time the FLAG_AWAKE was cleared.
+     * Always cleared when FLAG_AWAKE is set.
      * <p>
-     * FLAG_RESUMED will be cleared when the screen is off, since all apps get paused at this point.
-     * Thus, this flag indicates whether the launcher will be shown when the screen gets turned on
+     * FLAG_RESUMED will be cleared when the device is asleep, since all apps get paused at this
+     * point. Thus, this flag indicates whether the launcher will be shown when the device wakes up
      * again.
      */
-    private static final int FLAG_LAUNCHER_ACTIVE_AT_SCREEN_OFF = 1 << 4;
+    private static final int FLAG_LAUNCHER_WAS_ACTIVE_WHILE_AWAKE = 1 << 4;
 
     /** Whether the device is currently locked. */
     private static final int FLAG_DEVICE_LOCKED = 1 << 5;
@@ -265,15 +265,15 @@
 
     /** SysUI flags updated, see QuickStepContract.SYSUI_STATE_* values. */
     public void updateStateForSysuiFlags(int systemUiStateFlags, boolean skipAnim) {
-        final boolean prevScreenIsOn = hasAnyFlag(FLAG_SCREEN_ON);
-        final boolean currScreenIsOn = hasAnyFlag(systemUiStateFlags, SYSUI_STATE_SCREEN_ON);
+        final boolean prevIsAwake = hasAnyFlag(FLAG_AWAKE);
+        final boolean currIsAwake = hasAnyFlag(systemUiStateFlags, SYSUI_STATE_AWAKE);
 
-        updateStateForFlag(FLAG_SCREEN_ON, currScreenIsOn);
-        if (prevScreenIsOn != currScreenIsOn) {
+        updateStateForFlag(FLAG_AWAKE, currIsAwake);
+        if (prevIsAwake != currIsAwake) {
             // The screen is switching between on/off. When turning off, capture whether the
             // launcher is active and memoize this state.
-            updateStateForFlag(FLAG_LAUNCHER_ACTIVE_AT_SCREEN_OFF,
-                    prevScreenIsOn && hasAnyFlag(FLAGS_LAUNCHER_ACTIVE));
+            updateStateForFlag(FLAG_LAUNCHER_WAS_ACTIVE_WHILE_AWAKE,
+                    prevIsAwake && hasAnyFlag(FLAGS_LAUNCHER_ACTIVE));
         }
 
         boolean isDeviceLocked = hasAnyFlag(systemUiStateFlags, MASK_ANY_SYSUI_LOCKED);
@@ -377,7 +377,7 @@
             }
         }
 
-        if (hasAnyFlag(changedFlags, FLAGS_LAUNCHER_ACTIVE | FLAG_SCREEN_ON)) {
+        if (hasAnyFlag(changedFlags, FLAGS_LAUNCHER_ACTIVE | FLAG_AWAKE)) {
             animatorSet.addListener(new AnimatorListenerAdapter() {
                 @Override
                 public void onAnimationStart(Animator animation) {
@@ -573,10 +573,10 @@
 
     /** Whether the launcher is considered active. */
     private boolean isInLauncher() {
-        if (hasAnyFlag(FLAG_SCREEN_ON)) {
+        if (hasAnyFlag(FLAG_AWAKE)) {
             return hasAnyFlag(FLAGS_LAUNCHER_ACTIVE);
         } else {
-            return hasAnyFlag(FLAG_LAUNCHER_ACTIVE_AT_SCREEN_OFF);
+            return hasAnyFlag(FLAG_LAUNCHER_WAS_ACTIVE_WHILE_AWAKE);
         }
     }
 
@@ -677,9 +677,9 @@
         appendFlag(result, flags, FLAG_TRANSITION_TO_RESUMED, "transition_to_resumed");
         appendFlag(result, flags, FLAG_LAUNCHER_IN_STATE_TRANSITION,
                 "launcher_in_state_transition");
-        appendFlag(result, flags, FLAG_SCREEN_ON, "screen_on");
-        appendFlag(result, flags, FLAG_LAUNCHER_ACTIVE_AT_SCREEN_OFF,
-                "launcher_active_at_screen_off");
+        appendFlag(result, flags, FLAG_AWAKE, "awake");
+        appendFlag(result, flags, FLAG_LAUNCHER_WAS_ACTIVE_WHILE_AWAKE,
+                "was_active_while_awake");
         appendFlag(result, flags, FLAG_DEVICE_LOCKED, "device_locked");
         return result.toString();
     }
diff --git a/quickstep/src/com/android/launcher3/uioverrides/QuickstepLauncher.java b/quickstep/src/com/android/launcher3/uioverrides/QuickstepLauncher.java
index 5e22703..7318298 100644
--- a/quickstep/src/com/android/launcher3/uioverrides/QuickstepLauncher.java
+++ b/quickstep/src/com/android/launcher3/uioverrides/QuickstepLauncher.java
@@ -639,7 +639,7 @@
         PendingAnimation anim = new PendingAnimation(TABLET_HOME_TO_SPLIT.getDuration());
         RectF startingTaskRect = new RectF();
         final FloatingTaskView floatingTaskView = FloatingTaskView.getFloatingTaskView(this,
-                source.view, null /* thumbnail */, source.drawable, startingTaskRect);
+                source.getView(), null /* thumbnail */, source.getDrawable(), startingTaskRect);
         floatingTaskView.setAlpha(1);
         floatingTaskView.addStagingAnimation(anim, startingTaskRect, tempRect,
                 false /* fadeWithThumbnail */, true /* isStagedTask */);
diff --git a/quickstep/src/com/android/quickstep/util/SplitAnimationController.kt b/quickstep/src/com/android/quickstep/util/SplitAnimationController.kt
index 6dd67de..b76fe5c 100644
--- a/quickstep/src/com/android/quickstep/util/SplitAnimationController.kt
+++ b/quickstep/src/com/android/quickstep/util/SplitAnimationController.kt
@@ -24,6 +24,7 @@
 import com.android.launcher3.DeviceProfile
 import com.android.launcher3.anim.PendingAnimation
 import com.android.launcher3.util.SplitConfigurationOptions.SplitSelectSource
+import com.android.quickstep.views.IconView
 import com.android.quickstep.views.TaskThumbnailView
 import com.android.quickstep.views.TaskView
 import com.android.quickstep.views.TaskView.TaskIdAttributeContainer
@@ -52,21 +53,22 @@
      * depending on the state of the surface from which the split was initiated
      */
     fun getFirstAnimInitViews(taskViewSupplier: Supplier<TaskView>,
-                              splitSelectSourceSupplier: Supplier<SplitSelectSource>)
+                              splitSelectSourceSupplier: Supplier<SplitSelectSource?>)
             : SplitAnimInitProps {
+        val splitSelectSource = splitSelectSourceSupplier.get()
         if (!splitSelectStateController.isAnimateCurrentTaskDismissal) {
             // Initiating from home
-            val splitSelectSource = splitSelectSourceSupplier.get()
-            return SplitAnimInitProps(splitSelectSource.view, originalBitmap = null,
+            return SplitAnimInitProps(splitSelectSource!!.view, originalBitmap = null,
                     splitSelectSource.drawable, fadeWithThumbnail = false, isStagedTask = true,
                     iconView = null)
         } else if (splitSelectStateController.isDismissingFromSplitPair) {
             // Initiating split from overview, but on a split pair
             val taskView = taskViewSupplier.get()
             for (container : TaskIdAttributeContainer in taskView.taskIdAttributeContainers) {
-                if (container.task.key.id == splitSelectStateController.initialTaskId) {
+                if (container.task.getKey().getId() == splitSelectStateController.initialTaskId) {
+                    val drawable = getDrawable(container.iconView, splitSelectSource)
                     return SplitAnimInitProps(container.thumbnailView,
-                            container.thumbnailView.thumbnail, container.iconView.drawable!!,
+                            container.thumbnailView.thumbnail, drawable!!,
                             fadeWithThumbnail = true, isStagedTask = true,
                             iconView = container.iconView
                     )
@@ -77,14 +79,28 @@
         } else {
             // Initiating split from overview on fullscreen task TaskView
             val taskView = taskViewSupplier.get()
+            val drawable = getDrawable(taskView.iconView, splitSelectSource)
             return SplitAnimInitProps(taskView.thumbnail, taskView.thumbnail.thumbnail,
-                    taskView.iconView.drawable!!, fadeWithThumbnail = true, isStagedTask = true,
+                    drawable!!, fadeWithThumbnail = true, isStagedTask = true,
                     taskView.iconView
             )
         }
     }
 
     /**
+     * Returns the drawable that's provided in iconView, however if that
+     * is null it falls back to the drawable that's in splitSelectSource.
+     * TaskView's icon drawable can be null if the TaskView is scrolled far enough off screen
+     * @return [Drawable]
+     */
+    fun getDrawable(iconView: IconView, splitSelectSource: SplitSelectSource?) : Drawable? {
+        if (iconView.drawable == null && splitSelectSource != null) {
+            return splitSelectSource.drawable
+        }
+        return iconView.drawable
+    }
+
+    /**
      * When selecting first app from split pair, second app's thumbnail remains. This animates
      * the second thumbnail by expanding it to take up the full taskViewWidth/Height and overlaying
      * it with [TaskThumbnailView]'s splashView. Adds animations to the provided builder.
diff --git a/quickstep/tests/src/com/android/quickstep/util/SplitAnimationControllerTest.kt b/quickstep/tests/src/com/android/quickstep/util/SplitAnimationControllerTest.kt
new file mode 100644
index 0000000..7e07b81
--- /dev/null
+++ b/quickstep/tests/src/com/android/quickstep/util/SplitAnimationControllerTest.kt
@@ -0,0 +1,165 @@
+/*
+ *  Copyright (C) 2023 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 com.android.quickstep.util
+
+import android.graphics.Bitmap
+import android.graphics.drawable.Drawable
+import android.view.View
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import com.android.launcher3.util.SplitConfigurationOptions
+import com.android.quickstep.views.GroupedTaskView
+import com.android.quickstep.views.IconView
+import com.android.quickstep.views.TaskThumbnailView
+import com.android.quickstep.views.TaskView
+import com.android.quickstep.views.TaskView.TaskIdAttributeContainer
+import com.android.systemui.shared.recents.model.Task
+import org.junit.Assert.assertEquals
+import org.junit.Before
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.mockito.Mock
+import org.mockito.MockitoAnnotations
+import org.mockito.Mockito.`when` as whenever
+
+@RunWith(AndroidJUnit4::class)
+class SplitAnimationControllerTest {
+
+    private val taskId = 9
+
+    @Mock lateinit var mockSplitSelectStateController: SplitSelectStateController
+    // TaskView
+    @Mock lateinit var mockTaskView: TaskView
+    @Mock lateinit var mockThumbnailView: TaskThumbnailView
+    @Mock lateinit var mockBitmap: Bitmap
+    @Mock lateinit var mockIconView: IconView
+    @Mock lateinit var mockTaskViewDrawable: Drawable
+    // GroupedTaskView
+    @Mock lateinit var mockGroupedTaskView: GroupedTaskView
+    @Mock lateinit var mockTask: Task
+    @Mock lateinit var mockTaskKey: Task.TaskKey
+    @Mock lateinit var mockTaskIdAttributeContainer: TaskIdAttributeContainer
+
+    // SplitSelectSource
+    @Mock lateinit var splitSelectSource: SplitConfigurationOptions.SplitSelectSource
+    @Mock lateinit var mockSplitSourceDrawable: Drawable
+    @Mock lateinit var mockSplitSourceView: View
+
+    lateinit var splitAnimationController: SplitAnimationController
+
+    @Before
+    fun setup() {
+        MockitoAnnotations.initMocks(this)
+
+        whenever(mockTaskView.thumbnail).thenReturn(mockThumbnailView)
+        whenever(mockThumbnailView.thumbnail).thenReturn(mockBitmap)
+        whenever(mockTaskView.iconView).thenReturn(mockIconView)
+        whenever(mockIconView.drawable).thenReturn(mockTaskViewDrawable)
+
+        whenever(splitSelectSource.drawable).thenReturn(mockSplitSourceDrawable)
+        whenever(splitSelectSource.view).thenReturn(mockSplitSourceView)
+
+        splitAnimationController = SplitAnimationController(mockSplitSelectStateController)
+    }
+
+    @Test
+    fun getFirstAnimInitViews_nullTaskViewIcon_useSplitSourceIcon() {
+        // Hit fullscreen task dismissal state
+        whenever(mockSplitSelectStateController.isAnimateCurrentTaskDismissal).thenReturn(true)
+        whenever(mockSplitSelectStateController.isDismissingFromSplitPair).thenReturn(false)
+
+        // Missing taskView icon
+        whenever(mockIconView.drawable).thenReturn(null)
+
+        val splitAnimInitProps : SplitAnimationController.Companion.SplitAnimInitProps =
+                splitAnimationController.getFirstAnimInitViews(
+                        { mockTaskView }, { splitSelectSource })
+
+        assertEquals("Did not fallback to use splitSource icon drawable",
+                mockSplitSourceDrawable, splitAnimInitProps.iconDrawable)
+    }
+
+    @Test
+    fun getFirstAnimInitViews_validTaskViewIcon_useTaskViewIcon() {
+        // Hit fullscreen task dismissal state
+        whenever(mockSplitSelectStateController.isAnimateCurrentTaskDismissal).thenReturn(true)
+        whenever(mockSplitSelectStateController.isDismissingFromSplitPair).thenReturn(false)
+
+        val splitAnimInitProps : SplitAnimationController.Companion.SplitAnimInitProps =
+                splitAnimationController.getFirstAnimInitViews(
+                        { mockTaskView }, { splitSelectSource })
+
+        assertEquals("Did not use taskView icon drawable", mockTaskViewDrawable,
+                splitAnimInitProps.iconDrawable)
+    }
+
+    @Test
+    fun getFirstAnimInitViews_validTaskViewNullSplitSource_useTaskViewIcon() {
+        // Hit fullscreen task dismissal state
+        whenever(mockSplitSelectStateController.isAnimateCurrentTaskDismissal).thenReturn(true)
+        whenever(mockSplitSelectStateController.isDismissingFromSplitPair).thenReturn(false)
+
+        // Set split source to null
+        whenever(splitSelectSource.drawable).thenReturn(null)
+
+        val splitAnimInitProps : SplitAnimationController.Companion.SplitAnimInitProps =
+                splitAnimationController.getFirstAnimInitViews(
+                        { mockTaskView }, { splitSelectSource })
+
+        assertEquals("Did not use taskView icon drawable", mockTaskViewDrawable,
+                splitAnimInitProps.iconDrawable)
+    }
+
+    @Test
+    fun getFirstAnimInitViews_nullTaskViewValidSplitSource_noTaskDismissal() {
+        // Hit initiating split from home
+        whenever(mockSplitSelectStateController.isAnimateCurrentTaskDismissal).thenReturn(false)
+        whenever(mockSplitSelectStateController.isDismissingFromSplitPair).thenReturn(false)
+
+        val splitAnimInitProps : SplitAnimationController.Companion.SplitAnimInitProps =
+                splitAnimationController.getFirstAnimInitViews(
+                        { mockTaskView }, { splitSelectSource })
+
+        assertEquals("Did not use splitSource icon drawable", mockSplitSourceDrawable,
+                splitAnimInitProps.iconDrawable)
+    }
+
+    @Test
+    fun getFirstAnimInitViews_nullTaskViewValidSplitSource_groupedTaskView() {
+        // Hit groupedTaskView dismissal
+        whenever(mockSplitSelectStateController.isAnimateCurrentTaskDismissal).thenReturn(true)
+        whenever(mockSplitSelectStateController.isDismissingFromSplitPair).thenReturn(true)
+
+        // Remove icon view from GroupedTaskView
+        whenever(mockIconView.drawable).thenReturn(null)
+
+        whenever(mockTaskIdAttributeContainer.task).thenReturn(mockTask)
+        whenever(mockTaskIdAttributeContainer.iconView).thenReturn(mockIconView)
+        whenever(mockTaskIdAttributeContainer.thumbnailView).thenReturn(mockThumbnailView)
+        whenever(mockTask.getKey()).thenReturn(mockTaskKey)
+        whenever(mockTaskKey.getId()).thenReturn(taskId)
+        whenever(mockSplitSelectStateController.initialTaskId).thenReturn(taskId)
+        whenever(mockGroupedTaskView.taskIdAttributeContainers)
+                .thenReturn(Array(1) { mockTaskIdAttributeContainer })
+        val splitAnimInitProps : SplitAnimationController.Companion.SplitAnimInitProps =
+                splitAnimationController.getFirstAnimInitViews(
+                        { mockGroupedTaskView }, { splitSelectSource })
+
+        assertEquals("Did not use splitSource icon drawable", mockSplitSourceDrawable,
+                splitAnimInitProps.iconDrawable)
+    }
+}
\ No newline at end of file
diff --git a/src/com/android/launcher3/util/SplitConfigurationOptions.java b/src/com/android/launcher3/util/SplitConfigurationOptions.java
index 8c5e782..1ae43d0 100644
--- a/src/com/android/launcher3/util/SplitConfigurationOptions.java
+++ b/src/com/android/launcher3/util/SplitConfigurationOptions.java
@@ -200,8 +200,8 @@
         /** Keep in sync w/ ActivityTaskManager#INVALID_TASK_ID (unreference-able) */
         private static final int INVALID_TASK_ID = -1;
 
-        public final View view;
-        public final Drawable drawable;
+        private View view;
+        private Drawable drawable;
         public final Intent intent;
         public final SplitPositionOption position;
         public final ItemInfo itemInfo;
@@ -224,5 +224,13 @@
             this.itemInfo = itemInfo;
             this.splitEvent = splitEvent;
         }
+
+        public Drawable getDrawable() {
+            return drawable;
+        }
+
+        public View getView() {
+            return view;
+        }
     }
 }