Reduce object allocations in WM in some frequently called methods

With the use of lambdas to get all windows in the window container
hierarchy, we need to be careful in frequently called code paths to
make sure the number of objects we allocate isn't crazy. This CL
converts some commonly called code paths that use lambda to use a
method reference for the lambda so we only need to allocate once vs.
each time the code path is executed.

Test: Perform some common operations on the phone and make sure
      the object allocations that show-up in Allocator Tracker for
      window manager seems reasonable

Change-Id: Ie0f245980de96ec68a4e62e76130db7d98c3f7d9
diff --git a/services/core/java/com/android/server/am/ActivityManagerService.java b/services/core/java/com/android/server/am/ActivityManagerService.java
index 6aa0dc9..a5b3020 100644
--- a/services/core/java/com/android/server/am/ActivityManagerService.java
+++ b/services/core/java/com/android/server/am/ActivityManagerService.java
@@ -1386,6 +1386,8 @@
 
     final long[] mTmpLong = new long[2];
 
+    private final ArraySet<BroadcastQueue> mTmpBroadcastQueue = new ArraySet();
+
     static final class ProcessChangeItem {
         static final int CHANGE_ACTIVITIES = 1<<0;
         static final int CHANGE_PROCESS_STATE = 1<<1;
@@ -19521,7 +19523,7 @@
         int schedGroup;
         int procState;
         boolean foregroundActivities = false;
-        final ArraySet<BroadcastQueue> queues = new ArraySet<BroadcastQueue>();
+        mTmpBroadcastQueue.clear();
         if (app == TOP_APP) {
             // The last app on the list is the foreground app.
             adj = ProcessList.FOREGROUND_APP_ADJ;
@@ -19535,13 +19537,13 @@
             schedGroup = ProcessList.SCHED_GROUP_DEFAULT;
             app.adjType = "instrumentation";
             procState = ActivityManager.PROCESS_STATE_FOREGROUND_SERVICE;
-        } else if (isReceivingBroadcastLocked(app, queues)) {
+        } else if (isReceivingBroadcastLocked(app, mTmpBroadcastQueue)) {
             // An app that is currently receiving a broadcast also
             // counts as being in the foreground for OOM killer purposes.
             // It's placed in a sched group based on the nature of the
             // broadcast as reflected by which queue it's active in.
             adj = ProcessList.FOREGROUND_APP_ADJ;
-            schedGroup = (queues.contains(mFgBroadcastQueue))
+            schedGroup = (mTmpBroadcastQueue.contains(mFgBroadcastQueue))
                     ? ProcessList.SCHED_GROUP_DEFAULT : ProcessList.SCHED_GROUP_BACKGROUND;
             app.adjType = "broadcast";
             procState = ActivityManager.PROCESS_STATE_RECEIVER;
diff --git a/services/core/java/com/android/server/wm/AppWindowToken.java b/services/core/java/com/android/server/wm/AppWindowToken.java
index e930818..00c37d2 100644
--- a/services/core/java/com/android/server/wm/AppWindowToken.java
+++ b/services/core/java/com/android/server/wm/AppWindowToken.java
@@ -1253,7 +1253,8 @@
             mService.mAnimator.mAppWindowAnimating = true;
         } else if (mAppAnimator.wasAnimating) {
             // stopped animating, do one more pass through the layout
-            setAppLayoutChanges(FINISH_LAYOUT_REDO_WALLPAPER, "appToken " + this + " done");
+            setAppLayoutChanges(FINISH_LAYOUT_REDO_WALLPAPER,
+                    DEBUG_LAYOUT_REPEATS ? "appToken " + this + " done" : null);
             if (DEBUG_ANIM) Slog.v(TAG, "updateWindowsApps...: done animating " + this);
         }
     }
diff --git a/services/core/java/com/android/server/wm/DisplayContent.java b/services/core/java/com/android/server/wm/DisplayContent.java
index 1df7c66..203137d 100644
--- a/services/core/java/com/android/server/wm/DisplayContent.java
+++ b/services/core/java/com/android/server/wm/DisplayContent.java
@@ -124,6 +124,8 @@
 import java.util.Iterator;
 import java.util.LinkedList;
 import java.util.List;
+import java.util.function.Consumer;
+import java.util.function.Predicate;
 
 /**
  * Utility class for keeping track of the WindowStates and other pertinent contents of a
@@ -156,6 +158,11 @@
             new NonAppWindowContainers("mImeWindowsContainers");
 
     private WindowState mTmpWindow;
+    private WindowState mTmpWindow2;
+    private WindowAnimator mTmpWindowAnimator;
+    private boolean mTmpRecoveringMemory;
+    private boolean mUpdateImeTarget;
+    private boolean mTmpInitial;
 
     // Mapping from a token IBinder to a WindowToken object on this display.
     private final HashMap<IBinder, WindowToken> mTokenMap = new HashMap();
@@ -178,7 +185,7 @@
     private boolean mLayoutNeeded;
     int pendingLayoutChanges;
     // TODO(multi-display): remove some of the usages.
-    final boolean isDefaultDisplay;
+    boolean isDefaultDisplay;
 
     /** Window tokens that are in the process of exiting, but still on screen for animations. */
     final ArrayList<WindowToken> mExitingTokens = new ArrayList<>();
@@ -200,7 +207,7 @@
     private final Matrix mTmpMatrix = new Matrix();
     private final Region mTmpRegion = new Region();
 
-    final WindowManagerService mService;
+    WindowManagerService mService;
 
     /** Remove this display when animation on it has completed. */
     private boolean mDeferredRemoval;
@@ -208,7 +215,7 @@
     final DockedStackDividerController mDividerControllerLocked;
     final PinnedStackController mPinnedStackControllerLocked;
 
-    final DimLayerController mDimLayerController;
+    DimLayerController mDimLayerController;
 
     final ArrayList<WindowState> mTapExcludedWindows = new ArrayList<>();
 
@@ -231,9 +238,390 @@
     private boolean mRemovingDisplay = false;
 
     private final WindowLayersController mLayersController;
-    final WallpaperController mWallpaperController;
+    WallpaperController mWallpaperController;
     int mInputMethodAnimLayerAdjustment;
 
+    private final Consumer<WindowState> mUpdateWindowsForAnimator = w -> {
+        WindowStateAnimator winAnimator = w.mWinAnimator;
+        if (winAnimator.hasSurface()) {
+            final boolean wasAnimating = winAnimator.mWasAnimating;
+            final boolean nowAnimating = winAnimator.stepAnimationLocked(
+                    mTmpWindowAnimator.mCurrentTime);
+            winAnimator.mWasAnimating = nowAnimating;
+            mTmpWindowAnimator.orAnimating(nowAnimating);
+
+            if (DEBUG_WALLPAPER) Slog.v(TAG,
+                    w + ": wasAnimating=" + wasAnimating + ", nowAnimating=" + nowAnimating);
+
+            if (wasAnimating && !winAnimator.mAnimating
+                    && mWallpaperController.isWallpaperTarget(w)) {
+                mTmpWindowAnimator.mBulkUpdateParams |= SET_WALLPAPER_MAY_CHANGE;
+                pendingLayoutChanges |= FINISH_LAYOUT_REDO_WALLPAPER;
+                if (DEBUG_LAYOUT_REPEATS) {
+                    mService.mWindowPlacerLocked.debugLayoutRepeats(
+                            "updateWindowsAndWallpaperLocked 2", pendingLayoutChanges);
+                }
+            }
+        }
+
+        final AppWindowToken atoken = w.mAppToken;
+        if (winAnimator.mDrawState == READY_TO_SHOW) {
+            if (atoken == null || atoken.allDrawn) {
+                if (w.performShowLocked()) {
+                    pendingLayoutChanges |= FINISH_LAYOUT_REDO_ANIM;
+                    if (DEBUG_LAYOUT_REPEATS) {
+                        mService.mWindowPlacerLocked.debugLayoutRepeats(
+                                "updateWindowsAndWallpaperLocked 5", pendingLayoutChanges);
+                    }
+                }
+            }
+        }
+        final AppWindowAnimator appAnimator = winAnimator.mAppAnimator;
+        if (appAnimator != null && appAnimator.thumbnail != null) {
+            if (appAnimator.thumbnailTransactionSeq
+                    != mTmpWindowAnimator.mAnimTransactionSequence) {
+                appAnimator.thumbnailTransactionSeq =
+                        mTmpWindowAnimator.mAnimTransactionSequence;
+                appAnimator.thumbnailLayer = 0;
+            }
+            if (appAnimator.thumbnailLayer < winAnimator.mAnimLayer) {
+                appAnimator.thumbnailLayer = winAnimator.mAnimLayer;
+            }
+        }
+    };
+
+    private final Consumer<WindowState> mUpdateWallpaperForAnimator = w -> {
+        final WindowStateAnimator winAnimator = w.mWinAnimator;
+        if (winAnimator.mSurfaceController == null || !winAnimator.hasSurface()) {
+            return;
+        }
+
+        final int flags = w.mAttrs.flags;
+
+        // If this window is animating, make a note that we have an animating window and take
+        // care of a request to run a detached wallpaper animation.
+        if (winAnimator.mAnimating) {
+            if (winAnimator.mAnimation != null) {
+                if ((flags & FLAG_SHOW_WALLPAPER) != 0
+                        && winAnimator.mAnimation.getDetachWallpaper()) {
+                    mTmpWindow = w;
+                }
+                final int color = winAnimator.mAnimation.getBackgroundColor();
+                if (color != 0) {
+                    final TaskStack stack = w.getStack();
+                    if (stack != null) {
+                        stack.setAnimationBackground(winAnimator, color);
+                    }
+                }
+            }
+            mTmpWindowAnimator.setAnimating(true);
+        }
+
+        // If this window's app token is running a detached wallpaper animation, make a note so
+        // we can ensure the wallpaper is displayed behind it.
+        final AppWindowAnimator appAnimator = winAnimator.mAppAnimator;
+        if (appAnimator != null && appAnimator.animation != null
+                && appAnimator.animating) {
+            if ((flags & FLAG_SHOW_WALLPAPER) != 0
+                    && appAnimator.animation.getDetachWallpaper()) {
+                mTmpWindow = w;
+            }
+
+            final int color = appAnimator.animation.getBackgroundColor();
+            if (color != 0) {
+                final TaskStack stack = w.getStack();
+                if (stack != null) {
+                    stack.setAnimationBackground(winAnimator, color);
+                }
+            }
+        }
+    };
+
+    private final Consumer<WindowState> mSetInputMethodAnimLayerAdjustment =
+            w -> w.adjustAnimLayer(mInputMethodAnimLayerAdjustment);
+
+    private final Consumer<WindowState> mScheduleToastTimeout = w -> {
+        final int lostFocusUid = mTmpWindow.mOwnerUid;
+        final Handler handler = mService.mH;
+        if (w.mAttrs.type == TYPE_TOAST && w.mOwnerUid == lostFocusUid) {
+            if (!handler.hasMessages(WINDOW_HIDE_TIMEOUT, w)) {
+                handler.sendMessageDelayed(handler.obtainMessage(WINDOW_HIDE_TIMEOUT, w),
+                        w.mAttrs.hideTimeoutMilliseconds);
+            }
+        }
+    };
+
+    private final ToBooleanFunction<WindowState> mFindFocusedWindow = w -> {
+        final AppWindowToken focusedApp = mService.mFocusedApp;
+        if (DEBUG_FOCUS) Slog.v(TAG_WM, "Looking for focus: " + w
+                + ", flags=" + w.mAttrs.flags + ", canReceive=" + w.canReceiveKeys());
+
+        if (!w.canReceiveKeys()) {
+            return false;
+        }
+
+        final AppWindowToken wtoken = w.mAppToken;
+
+        // If this window's application has been removed, just skip it.
+        if (wtoken != null && (wtoken.removed || wtoken.sendingToBottom)) {
+            if (DEBUG_FOCUS) Slog.v(TAG_WM, "Skipping " + wtoken + " because "
+                    + (wtoken.removed ? "removed" : "sendingToBottom"));
+            return false;
+        }
+
+        if (focusedApp == null) {
+            if (DEBUG_FOCUS_LIGHT) Slog.v(TAG_WM, "findFocusedWindow: focusedApp=null"
+                    + " using new focus @ " + w);
+            mTmpWindow = w;
+            return true;
+        }
+
+        if (!focusedApp.windowsAreFocusable()) {
+            // Current focused app windows aren't focusable...
+            if (DEBUG_FOCUS_LIGHT) Slog.v(TAG_WM, "findFocusedWindow: focusedApp windows not"
+                    + " focusable using new focus @ " + w);
+            mTmpWindow = w;
+            return true;
+        }
+
+        // Descend through all of the app tokens and find the first that either matches
+        // win.mAppToken (return win) or mFocusedApp (return null).
+        if (wtoken != null && w.mAttrs.type != TYPE_APPLICATION_STARTING) {
+            if (focusedApp.compareTo(wtoken) > 0) {
+                // App stack below focused app stack. No focus for you!!!
+                if (DEBUG_FOCUS_LIGHT) Slog.v(TAG_WM,
+                        "findFocusedWindow: Reached focused app=" + focusedApp);
+                mTmpWindow = null;
+                return true;
+            }
+        }
+
+        if (DEBUG_FOCUS_LIGHT) Slog.v(TAG_WM, "findFocusedWindow: Found new focus @ " + w);
+        mTmpWindow = w;
+        return true;
+    };
+
+    private final Consumer<WindowState> mPrepareWindowSurfaces =
+            w -> w.mWinAnimator.prepareSurfaceLocked(true);
+
+    private final Consumer<WindowState> mPerformLayout = w -> {
+        // Don't do layout of a window if it is not visible, or soon won't be visible, to avoid
+        // wasting time and funky changes while a window is animating away.
+        final boolean gone = (mTmpWindow != null && mService.mPolicy.canBeHiddenByKeyguardLw(w))
+                || w.isGoneForLayoutLw();
+
+        if (DEBUG_LAYOUT && !w.mLayoutAttached) {
+            Slog.v(TAG, "1ST PASS " + w + ": gone=" + gone + " mHaveFrame=" + w.mHaveFrame
+                    + " mLayoutAttached=" + w.mLayoutAttached
+                    + " screen changed=" + w.isConfigChanged());
+            final AppWindowToken atoken = w.mAppToken;
+            if (gone) Slog.v(TAG, "  GONE: mViewVisibility=" + w.mViewVisibility
+                    + " mRelayoutCalled=" + w.mRelayoutCalled + " hidden=" + w.mToken.hidden
+                    + " hiddenRequested=" + (atoken != null && atoken.hiddenRequested)
+                    + " parentHidden=" + w.isParentWindowHidden());
+            else Slog.v(TAG, "  VIS: mViewVisibility=" + w.mViewVisibility
+                    + " mRelayoutCalled=" + w.mRelayoutCalled + " hidden=" + w.mToken.hidden
+                    + " hiddenRequested=" + (atoken != null && atoken.hiddenRequested)
+                    + " parentHidden=" + w.isParentWindowHidden());
+        }
+
+        // If this view is GONE, then skip it -- keep the current frame, and let the caller know
+        // so they can ignore it if they want.  (We do the normal layout for INVISIBLE windows,
+        // since that means "perform layout as normal, just don't display").
+        if (!gone || !w.mHaveFrame || w.mLayoutNeeded
+                || ((w.isConfigChanged() || w.setReportResizeHints())
+                && !w.isGoneForLayoutLw() &&
+                ((w.mAttrs.privateFlags & PRIVATE_FLAG_KEYGUARD) != 0 ||
+                        (w.mHasSurface && w.mAppToken != null &&
+                                w.mAppToken.layoutConfigChanges)))) {
+            if (!w.mLayoutAttached) {
+                if (mTmpInitial) {
+                    //Slog.i(TAG, "Window " + this + " clearing mContentChanged - initial");
+                    w.mContentChanged = false;
+                }
+                if (w.mAttrs.type == TYPE_DREAM) {
+                    // Don't layout windows behind a dream, so that if it does stuff like hide
+                    // the status bar we won't get a bad transition when it goes away.
+                    mTmpWindow = w;
+                }
+                w.mLayoutNeeded = false;
+                w.prelayout();
+                mService.mPolicy.layoutWindowLw(w, null);
+                w.mLayoutSeq = mService.mLayoutSeq;
+
+                // Window frames may have changed. Update dim layer with the new bounds.
+                final Task task = w.getTask();
+                if (task != null) {
+                    mDimLayerController.updateDimLayer(task);
+                }
+
+                if (DEBUG_LAYOUT) Slog.v(TAG, "  LAYOUT: mFrame=" + w.mFrame
+                        + " mContainingFrame=" + w.mContainingFrame
+                        + " mDisplayFrame=" + w.mDisplayFrame);
+            }
+        }
+    };
+
+    private final Consumer<WindowState> mPerformLayoutAttached = w -> {
+        if (w.mLayoutAttached) {
+            if (DEBUG_LAYOUT) Slog.v(TAG, "2ND PASS " + w + " mHaveFrame=" + w.mHaveFrame
+                    + " mViewVisibility=" + w.mViewVisibility
+                    + " mRelayoutCalled=" + w.mRelayoutCalled);
+            // If this view is GONE, then skip it -- keep the current frame, and let the caller
+            // know so they can ignore it if they want.  (We do the normal layout for INVISIBLE
+            // windows, since that means "perform layout as normal, just don't display").
+            if (mTmpWindow != null && mService.mPolicy.canBeHiddenByKeyguardLw(w)) {
+                return;
+            }
+            if ((w.mViewVisibility != GONE && w.mRelayoutCalled) || !w.mHaveFrame
+                    || w.mLayoutNeeded) {
+                if (mTmpInitial) {
+                    //Slog.i(TAG, "Window " + this + " clearing mContentChanged - initial");
+                    w.mContentChanged = false;
+                }
+                w.mLayoutNeeded = false;
+                w.prelayout();
+                mService.mPolicy.layoutWindowLw(w, w.getParentWindow());
+                w.mLayoutSeq = mService.mLayoutSeq;
+                if (DEBUG_LAYOUT) Slog.v(TAG, " LAYOUT: mFrame=" + w.mFrame
+                        + " mContainingFrame=" + w.mContainingFrame
+                        + " mDisplayFrame=" + w.mDisplayFrame);
+            }
+        } else if (w.mAttrs.type == TYPE_DREAM) {
+            // Don't layout windows behind a dream, so that if it does stuff like hide the
+            // status bar we won't get a bad transition when it goes away.
+            mTmpWindow = mTmpWindow2;
+        }
+    };
+
+    private final Predicate<WindowState> mComputeImeTargetPredicate = w -> {
+        if (DEBUG_INPUT_METHOD && mUpdateImeTarget) Slog.i(TAG_WM, "Checking window @" + w
+                + " fl=0x" + Integer.toHexString(w.mAttrs.flags));
+        return w.canBeImeTarget();
+    };
+
+    private final Consumer<WindowState> mApplyPostLayoutPolicy =
+            w -> mService.mPolicy.applyPostLayoutPolicyLw(w, w.mAttrs, w.getParentWindow(),
+                    mService.mInputMethodTarget);
+
+    private final Consumer<WindowState> mApplySurfaceChangesTransaction = w -> {
+        final WindowSurfacePlacer surfacePlacer = mService.mWindowPlacerLocked;
+        final boolean obscuredChanged = w.mObscured !=
+                mTmpApplySurfaceChangesTransactionState.obscured;
+        final RootWindowContainer root = mService.mRoot;
+        // Only used if default window
+        final boolean someoneLosingFocus = !mService.mLosingFocus.isEmpty();
+
+        // Update effect.
+        w.mObscured = mTmpApplySurfaceChangesTransactionState.obscured;
+        if (!mTmpApplySurfaceChangesTransactionState.obscured) {
+            final boolean isDisplayed = w.isDisplayedLw();
+
+            if (isDisplayed && w.isObscuringDisplay()) {
+                // This window completely covers everything behind it, so we want to leave all
+                // of them as undimmed (for performance reasons).
+                root.mObscuringWindow = w;
+                mTmpApplySurfaceChangesTransactionState.obscured = true;
+            }
+
+            mTmpApplySurfaceChangesTransactionState.displayHasContent |=
+                    root.handleNotObscuredLocked(w,
+                            mTmpApplySurfaceChangesTransactionState.obscured,
+                            mTmpApplySurfaceChangesTransactionState.syswin);
+
+            if (w.mHasSurface && isDisplayed) {
+                final int type = w.mAttrs.type;
+                if (type == TYPE_SYSTEM_DIALOG || type == TYPE_SYSTEM_ERROR
+                        || (w.mAttrs.privateFlags & PRIVATE_FLAG_KEYGUARD) != 0) {
+                    mTmpApplySurfaceChangesTransactionState.syswin = true;
+                }
+                if (mTmpApplySurfaceChangesTransactionState.preferredRefreshRate == 0
+                        && w.mAttrs.preferredRefreshRate != 0) {
+                    mTmpApplySurfaceChangesTransactionState.preferredRefreshRate
+                            = w.mAttrs.preferredRefreshRate;
+                }
+                if (mTmpApplySurfaceChangesTransactionState.preferredModeId == 0
+                        && w.mAttrs.preferredDisplayModeId != 0) {
+                    mTmpApplySurfaceChangesTransactionState.preferredModeId
+                            = w.mAttrs.preferredDisplayModeId;
+                }
+            }
+        }
+
+        w.applyDimLayerIfNeeded();
+
+        if (isDefaultDisplay && obscuredChanged && w.isVisibleLw()
+                && mWallpaperController.isWallpaperTarget(w)) {
+            // This is the wallpaper target and its obscured state changed... make sure the
+            // current wallpaper's visibility has been updated accordingly.
+            mWallpaperController.updateWallpaperVisibility();
+        }
+
+        w.handleWindowMovedIfNeeded();
+
+        final WindowStateAnimator winAnimator = w.mWinAnimator;
+
+        //Slog.i(TAG, "Window " + this + " clearing mContentChanged - done placing");
+        w.mContentChanged = false;
+
+        // Moved from updateWindowsAndWallpaperLocked().
+        if (w.mHasSurface) {
+            // Take care of the window being ready to display.
+            final boolean committed = winAnimator.commitFinishDrawingLocked();
+            if (isDefaultDisplay && committed) {
+                if (w.mAttrs.type == TYPE_DREAM) {
+                    // HACK: When a dream is shown, it may at that point hide the lock screen.
+                    // So we need to redo the layout to let the phone window manager make this
+                    // happen.
+                    pendingLayoutChanges |= FINISH_LAYOUT_REDO_LAYOUT;
+                    if (DEBUG_LAYOUT_REPEATS) {
+                        surfacePlacer.debugLayoutRepeats(
+                                "dream and commitFinishDrawingLocked true",
+                                pendingLayoutChanges);
+                    }
+                }
+                if ((w.mAttrs.flags & FLAG_SHOW_WALLPAPER) != 0) {
+                    if (DEBUG_WALLPAPER_LIGHT) Slog.v(TAG,
+                            "First draw done in potential wallpaper target " + w);
+                    root.mWallpaperMayChange = true;
+                    pendingLayoutChanges |= FINISH_LAYOUT_REDO_WALLPAPER;
+                    if (DEBUG_LAYOUT_REPEATS) {
+                        surfacePlacer.debugLayoutRepeats(
+                                "wallpaper and commitFinishDrawingLocked true",
+                                pendingLayoutChanges);
+                    }
+                }
+            }
+            if (!winAnimator.isAnimationStarting() && !winAnimator.isWaitingForOpening()) {
+                // Updates the shown frame before we set up the surface. This is needed
+                // because the resizing could change the top-left position (in addition to
+                // size) of the window. setSurfaceBoundariesLocked uses mShownPosition to
+                // position the surface.
+                //
+                // If an animation is being started, we can't call this method because the
+                // animation hasn't processed its initial transformation yet, but in general
+                // we do want to update the position if the window is animating.
+                winAnimator.computeShownFrameLocked();
+            }
+            winAnimator.setSurfaceBoundariesLocked(mTmpRecoveringMemory /* recoveringMemory */);
+        }
+
+        final AppWindowToken atoken = w.mAppToken;
+        if (atoken != null) {
+            final boolean updateAllDrawn = atoken.updateDrawnWindowStates(w);
+            if (updateAllDrawn && !mTmpUpdateAllDrawn.contains(atoken)) {
+                mTmpUpdateAllDrawn.add(atoken);
+            }
+        }
+
+        if (isDefaultDisplay && someoneLosingFocus && w == mService.mCurrentFocus
+                && w.isDisplayedLw()) {
+            mTmpApplySurfaceChangesTransactionState.focusDisplayed = true;
+        }
+
+        w.updateResizingWindowIfNeeded();
+    };
+
     /**
      * @param display May not be null.
      * @param service You know.
@@ -905,9 +1293,8 @@
     void setInputMethodAnimLayerAdjustment(int adj) {
         if (DEBUG_LAYERS) Slog.v(TAG_WM, "Setting im layer adj to " + adj);
         mInputMethodAnimLayerAdjustment = adj;
-        mImeWindowsContainers.forAllWindows(w -> {
-            w.adjustAnimLayer(adj);
-        }, true /* traverseTopToBottom */);
+        mImeWindowsContainers.forAllWindows(mSetInputMethodAnimLayerAdjustment,
+                true /* traverseTopToBottom */);
     }
 
     /**
@@ -1105,71 +1492,17 @@
         if (oldFocus == null || (newFocus != null && newFocus.mOwnerUid == oldFocus.mOwnerUid)) {
             return;
         }
-        final int lostFocusUid = oldFocus.mOwnerUid;
-        final Handler handler = mService.mH;
 
-        forAllWindows(w -> {
-            if (w.mAttrs.type == TYPE_TOAST && w.mOwnerUid == lostFocusUid) {
-                if (!handler.hasMessages(WINDOW_HIDE_TIMEOUT, w)) {
-                    handler.sendMessageDelayed(handler.obtainMessage(WINDOW_HIDE_TIMEOUT, w),
-                            w.mAttrs.hideTimeoutMilliseconds);
-                }
-            }
-        }, false /* traverseTopToBottom */);
+        // Used to communicate the old focus to the callback method.
+        mTmpWindow = oldFocus;
+
+        forAllWindows(mScheduleToastTimeout, false /* traverseTopToBottom */);
     }
 
     WindowState findFocusedWindow() {
-        final AppWindowToken focusedApp = mService.mFocusedApp;
         mTmpWindow = null;
 
-        forAllWindows(w -> {
-            if (DEBUG_FOCUS) Slog.v(TAG_WM, "Looking for focus: " + w
-                    + ", flags=" + w.mAttrs.flags + ", canReceive=" + w.canReceiveKeys());
-
-            if (!w.canReceiveKeys()) {
-                return false;
-            }
-
-            final AppWindowToken wtoken = w.mAppToken;
-
-            // If this window's application has been removed, just skip it.
-            if (wtoken != null && (wtoken.removed || wtoken.sendingToBottom)) {
-                if (DEBUG_FOCUS) Slog.v(TAG_WM, "Skipping " + wtoken + " because "
-                        + (wtoken.removed ? "removed" : "sendingToBottom"));
-                return false;
-            }
-
-            if (focusedApp == null) {
-                if (DEBUG_FOCUS_LIGHT) Slog.v(TAG_WM, "findFocusedWindow: focusedApp=null"
-                        + " using new focus @ " + w);
-                mTmpWindow = w;
-                return true;
-            }
-
-            if (!focusedApp.windowsAreFocusable()) {
-                // Current focused app windows aren't focusable...
-                if (DEBUG_FOCUS_LIGHT) Slog.v(TAG_WM, "findFocusedWindow: focusedApp windows not"
-                        + " focusable using new focus @ " + w);
-                mTmpWindow = w;
-                return true;
-            }
-
-            // Descend through all of the app tokens and find the first that either matches
-            // win.mAppToken (return win) or mFocusedApp (return null).
-            if (wtoken != null && w.mAttrs.type != TYPE_APPLICATION_STARTING) {
-                if (focusedApp.compareTo(wtoken) > 0) {
-                    // App stack below focused app stack. No focus for you!!!
-                    if (DEBUG_FOCUS_LIGHT) Slog.v(TAG_WM,
-                            "findFocusedWindow: Reached focused app=" + focusedApp);
-                    mTmpWindow = null;
-                    return true;
-                }
-            }
-
-            if (DEBUG_FOCUS_LIGHT) Slog.v(TAG_WM, "findFocusedWindow: Found new focus @ " + w);
-            mTmpWindow = w;
-            return true;
-        }, true /* traverseTopToBottom */);
+        forAllWindows(mFindFocusedWindow, true /* traverseTopToBottom */);
 
         if (mTmpWindow == null) {
             if (DEBUG_FOCUS_LIGHT) Slog.v(TAG_WM, "findFocusedWindow: No focusable windows.");
@@ -1251,11 +1584,8 @@
         // TODO(multidisplay): Needs some serious rethought when the target and IME are not on the
         // same display. Or even when the current IME/target are not on the same screen as the next
         // IME/target. For now only look for input windows on the main screen.
-        WindowState target = getWindow(w -> {
-            if (DEBUG_INPUT_METHOD && updateImeTarget) Slog.i(TAG_WM, "Checking window @"
-                    + w + " fl=0x" + Integer.toHexString(w.mAttrs.flags));
-            return w.canBeImeTarget();
-        });
+        mUpdateImeTarget = updateImeTarget;
+        WindowState target = getWindow(mComputeImeTargetPredicate);
 
 
         // Yet more tricksyness!  If this window is a "starting" window, we do actually want
@@ -1518,51 +1848,8 @@
     }
 
     void updateWindowsForAnimator(WindowAnimator animator) {
-        forAllWindows(w -> {
-            WindowStateAnimator winAnimator = w.mWinAnimator;
-            if (winAnimator.hasSurface()) {
-                final boolean wasAnimating = winAnimator.mWasAnimating;
-                final boolean nowAnimating = winAnimator.stepAnimationLocked(animator.mCurrentTime);
-                winAnimator.mWasAnimating = nowAnimating;
-                animator.orAnimating(nowAnimating);
-
-                if (DEBUG_WALLPAPER) Slog.v(TAG,
-                        w + ": wasAnimating=" + wasAnimating + ", nowAnimating=" + nowAnimating);
-
-                if (wasAnimating && !winAnimator.mAnimating
-                        && mWallpaperController.isWallpaperTarget(w)) {
-                    animator.mBulkUpdateParams |= SET_WALLPAPER_MAY_CHANGE;
-                    pendingLayoutChanges |= FINISH_LAYOUT_REDO_WALLPAPER;
-                    if (DEBUG_LAYOUT_REPEATS) {
-                        mService.mWindowPlacerLocked.debugLayoutRepeats(
-                                "updateWindowsAndWallpaperLocked 2", pendingLayoutChanges);
-                    }
-                }
-            }
-
-            final AppWindowToken atoken = w.mAppToken;
-            if (winAnimator.mDrawState == READY_TO_SHOW) {
-                if (atoken == null || atoken.allDrawn) {
-                    if (w.performShowLocked()) {
-                        pendingLayoutChanges |= FINISH_LAYOUT_REDO_ANIM;
-                        if (DEBUG_LAYOUT_REPEATS) {
-                            mService.mWindowPlacerLocked.debugLayoutRepeats(
-                                    "updateWindowsAndWallpaperLocked 5", pendingLayoutChanges);
-                        }
-                    }
-                }
-            }
-            final AppWindowAnimator appAnimator = winAnimator.mAppAnimator;
-            if (appAnimator != null && appAnimator.thumbnail != null) {
-                if (appAnimator.thumbnailTransactionSeq != animator.mAnimTransactionSequence) {
-                    appAnimator.thumbnailTransactionSeq = animator.mAnimTransactionSequence;
-                    appAnimator.thumbnailLayer = 0;
-                }
-                if (appAnimator.thumbnailLayer < winAnimator.mAnimLayer) {
-                    appAnimator.thumbnailLayer = winAnimator.mAnimLayer;
-                }
-            }
-        }, true /* traverseTopToBottom */);
+        mTmpWindowAnimator = animator;
+        forAllWindows(mUpdateWindowsForAnimator, true /* traverseTopToBottom */);
     }
 
     void updateWallpaperForAnimator(WindowAnimator animator) {
@@ -1570,53 +1857,9 @@
 
         // Used to indicate a detached wallpaper.
         mTmpWindow = null;
+        mTmpWindowAnimator = animator;
 
-        forAllWindows(w -> {
-            final WindowStateAnimator winAnimator = w.mWinAnimator;
-            if (winAnimator.mSurfaceController == null || !winAnimator.hasSurface()) {
-                return;
-            }
-
-            final int flags = w.mAttrs.flags;
-
-            // If this window is animating, make a note that we have an animating window and take
-            // care of a request to run a detached wallpaper animation.
-            if (winAnimator.mAnimating) {
-                if (winAnimator.mAnimation != null) {
-                    if ((flags & FLAG_SHOW_WALLPAPER) != 0
-                            && winAnimator.mAnimation.getDetachWallpaper()) {
-                        mTmpWindow = w;
-                    }
-                    final int color = winAnimator.mAnimation.getBackgroundColor();
-                    if (color != 0) {
-                        final TaskStack stack = w.getStack();
-                        if (stack != null) {
-                            stack.setAnimationBackground(winAnimator, color);
-                        }
-                    }
-                }
-                animator.setAnimating(true);
-            }
-
-            // If this window's app token is running a detached wallpaper animation, make a note so
-            // we can ensure the wallpaper is displayed behind it.
-            final AppWindowAnimator appAnimator = winAnimator.mAppAnimator;
-            if (appAnimator != null && appAnimator.animation != null
-                    && appAnimator.animating) {
-                if ((flags & FLAG_SHOW_WALLPAPER) != 0
-                        && appAnimator.animation.getDetachWallpaper()) {
-                    mTmpWindow = w;
-                }
-
-                final int color = appAnimator.animation.getBackgroundColor();
-                if (color != 0) {
-                    final TaskStack stack = w.getStack();
-                    if (stack != null) {
-                        stack.setAnimationBackground(winAnimator, color);
-                    }
-                }
-            }
-        }, true /* traverseTopToBottom */);
+        forAllWindows(mUpdateWallpaperForAnimator, true /* traverseTopToBottom */);
 
         if (animator.mWindowDetachedWallpaper != mTmpWindow) {
             if (DEBUG_WALLPAPER) Slog.v(TAG, "Detached wallpaper changed from "
@@ -1627,9 +1870,7 @@
     }
 
     void prepareWindowSurfaces() {
-        forAllWindows(w -> {
-            w.mWinAnimator.prepareSurfaceLocked(true);
-        }, false /* traverseTopToBottom */);
+        forAllWindows(mPrepareWindowSurfaces, false /* traverseTopToBottom */);
     }
 
     boolean inputMethodClientHasFocus(IInputMethodClient client) {
@@ -1765,136 +2006,18 @@
 
             if (isDefaultDisplay) {
                 mService.mPolicy.beginPostLayoutPolicyLw(dw, dh);
-                forAllWindows(w -> {
-                    mService.mPolicy.applyPostLayoutPolicyLw(w, w.mAttrs, w.getParentWindow(),
-                            mService.mInputMethodTarget);
-                }, true /* traverseTopToBottom */);
+                forAllWindows(mApplyPostLayoutPolicy, true /* traverseTopToBottom */);
                 pendingLayoutChanges |= mService.mPolicy.finishPostLayoutPolicyLw();
                 if (DEBUG_LAYOUT_REPEATS) surfacePlacer.debugLayoutRepeats(
                         "after finishPostLayoutPolicyLw", pendingLayoutChanges);
             }
         } while (pendingLayoutChanges != 0);
 
-        final RootWindowContainer root = mService.mRoot;
         mTmpApplySurfaceChangesTransactionState.reset();
         resetDimming();
 
-        // Only used if default window
-        final boolean someoneLosingFocus = !mService.mLosingFocus.isEmpty();
-
-        forAllWindows(w -> {
-            final boolean obscuredChanged = w.mObscured !=
-                    mTmpApplySurfaceChangesTransactionState.obscured;
-
-            // Update effect.
-            w.mObscured = mTmpApplySurfaceChangesTransactionState.obscured;
-            if (!mTmpApplySurfaceChangesTransactionState.obscured) {
-                final boolean isDisplayed = w.isDisplayedLw();
-
-                if (isDisplayed && w.isObscuringDisplay()) {
-                    // This window completely covers everything behind it, so we want to leave all
-                    // of them as undimmed (for performance reasons).
-                    root.mObscuringWindow = w;
-                    mTmpApplySurfaceChangesTransactionState.obscured = true;
-                }
-
-                mTmpApplySurfaceChangesTransactionState.displayHasContent |=
-                        root.handleNotObscuredLocked(w,
-                                mTmpApplySurfaceChangesTransactionState.obscured,
-                                mTmpApplySurfaceChangesTransactionState.syswin);
-
-                if (w.mHasSurface && isDisplayed) {
-                    final int type = w.mAttrs.type;
-                    if (type == TYPE_SYSTEM_DIALOG || type == TYPE_SYSTEM_ERROR
-                            || (w.mAttrs.privateFlags & PRIVATE_FLAG_KEYGUARD) != 0) {
-                        mTmpApplySurfaceChangesTransactionState.syswin = true;
-                    }
-                    if (mTmpApplySurfaceChangesTransactionState.preferredRefreshRate == 0
-                            && w.mAttrs.preferredRefreshRate != 0) {
-                        mTmpApplySurfaceChangesTransactionState.preferredRefreshRate
-                                = w.mAttrs.preferredRefreshRate;
-                    }
-                    if (mTmpApplySurfaceChangesTransactionState.preferredModeId == 0
-                            && w.mAttrs.preferredDisplayModeId != 0) {
-                        mTmpApplySurfaceChangesTransactionState.preferredModeId
-                                = w.mAttrs.preferredDisplayModeId;
-                    }
-                }
-            }
-
-            w.applyDimLayerIfNeeded();
-
-            if (isDefaultDisplay && obscuredChanged && w.isVisibleLw()
-                    && mWallpaperController.isWallpaperTarget(w)) {
-                // This is the wallpaper target and its obscured state changed... make sure the
-                // current wallpaper's visibility has been updated accordingly.
-                mWallpaperController.updateWallpaperVisibility();
-            }
-
-            w.handleWindowMovedIfNeeded();
-
-            final WindowStateAnimator winAnimator = w.mWinAnimator;
-
-            //Slog.i(TAG, "Window " + this + " clearing mContentChanged - done placing");
-            w.mContentChanged = false;
-
-            // Moved from updateWindowsAndWallpaperLocked().
-            if (w.mHasSurface) {
-                // Take care of the window being ready to display.
-                final boolean committed = winAnimator.commitFinishDrawingLocked();
-                if (isDefaultDisplay && committed) {
-                    if (w.mAttrs.type == TYPE_DREAM) {
-                        // HACK: When a dream is shown, it may at that point hide the lock screen.
-                        // So we need to redo the layout to let the phone window manager make this
-                        // happen.
-                        pendingLayoutChanges |= FINISH_LAYOUT_REDO_LAYOUT;
-                        if (DEBUG_LAYOUT_REPEATS) {
-                            surfacePlacer.debugLayoutRepeats(
-                                    "dream and commitFinishDrawingLocked true",
-                                    pendingLayoutChanges);
-                        }
-                    }
-                    if ((w.mAttrs.flags & FLAG_SHOW_WALLPAPER) != 0) {
-                        if (DEBUG_WALLPAPER_LIGHT) Slog.v(TAG,
-                                "First draw done in potential wallpaper target " + w);
-                        root.mWallpaperMayChange = true;
-                        pendingLayoutChanges |= FINISH_LAYOUT_REDO_WALLPAPER;
-                        if (DEBUG_LAYOUT_REPEATS) {
-                            surfacePlacer.debugLayoutRepeats(
-                                    "wallpaper and commitFinishDrawingLocked true",
-                                    pendingLayoutChanges);
-                        }
-                    }
-                }
-                if (!winAnimator.isAnimationStarting() && !winAnimator.isWaitingForOpening()) {
-                    // Updates the shown frame before we set up the surface. This is needed
-                    // because the resizing could change the top-left position (in addition to
-                    // size) of the window. setSurfaceBoundariesLocked uses mShownPosition to
-                    // position the surface.
-                    //
-                    // If an animation is being started, we can't call this method because the
-                    // animation hasn't processed its initial transformation yet, but in general
-                    // we do want to update the position if the window is animating.
-                    winAnimator.computeShownFrameLocked();
-                }
-                winAnimator.setSurfaceBoundariesLocked(recoveringMemory);
-            }
-
-            final AppWindowToken atoken = w.mAppToken;
-            if (atoken != null) {
-                final boolean updateAllDrawn = atoken.updateDrawnWindowStates(w);
-                if (updateAllDrawn && !mTmpUpdateAllDrawn.contains(atoken)) {
-                    mTmpUpdateAllDrawn.add(atoken);
-                }
-            }
-
-            if (isDefaultDisplay && someoneLosingFocus && w == mService.mCurrentFocus
-                    && w.isDisplayedLw()) {
-                mTmpApplySurfaceChangesTransactionState.focusDisplayed = true;
-            }
-
-            w.updateResizingWindowIfNeeded();
-        }, true /* traverseTopToBottom */);
+        mTmpRecoveringMemory = recoveringMemory;
+        forAllWindows(mApplySurfaceChangesTransaction, true /* traverseTopToBottom */);
 
         mService.mDisplayManagerInternal.setDisplayProperties(mDisplayId,
                 mTmpApplySurfaceChangesTransactionState.displayHasContent,
@@ -1923,8 +2046,6 @@
         final int dw = mDisplayInfo.logicalWidth;
         final int dh = mDisplayInfo.logicalHeight;
 
-        int i;
-
         if (DEBUG_LAYOUT) {
             Slog.v(TAG, "-------------------------------------");
             Slog.v(TAG, "performLayout: needed=" + isLayoutNeeded() + " dw=" + dw + " dh=" + dh);
@@ -1947,107 +2068,20 @@
         // Used to indicate that we have processed the dream window and all additional windows are
         // behind it.
         mTmpWindow = null;
+        mTmpInitial = initial;
 
         // First perform layout of any root windows (not attached to another window).
-        forAllWindows(w -> {
-            // Don't do layout of a window if it is not visible, or soon won't be visible, to avoid
-            // wasting time and funky changes while a window is animating away.
-            final boolean gone = (mTmpWindow != null && mService.mPolicy.canBeHiddenByKeyguardLw(w))
-                    || w.isGoneForLayoutLw();
-
-            if (DEBUG_LAYOUT && !w.mLayoutAttached) {
-                Slog.v(TAG, "1ST PASS " + w + ": gone=" + gone + " mHaveFrame=" + w.mHaveFrame
-                        + " mLayoutAttached=" + w.mLayoutAttached
-                        + " screen changed=" + w.isConfigChanged());
-                final AppWindowToken atoken = w.mAppToken;
-                if (gone) Slog.v(TAG, "  GONE: mViewVisibility=" + w.mViewVisibility
-                        + " mRelayoutCalled=" + w.mRelayoutCalled + " hidden=" + w.mToken.hidden
-                        + " hiddenRequested=" + (atoken != null && atoken.hiddenRequested)
-                        + " parentHidden=" + w.isParentWindowHidden());
-                else Slog.v(TAG, "  VIS: mViewVisibility=" + w.mViewVisibility
-                        + " mRelayoutCalled=" + w.mRelayoutCalled + " hidden=" + w.mToken.hidden
-                        + " hiddenRequested=" + (atoken != null && atoken.hiddenRequested)
-                        + " parentHidden=" + w.isParentWindowHidden());
-            }
-
-            // If this view is GONE, then skip it -- keep the current frame, and let the caller know
-            // so they can ignore it if they want.  (We do the normal layout for INVISIBLE windows,
-            // since that means "perform layout as normal, just don't display").
-            if (!gone || !w.mHaveFrame || w.mLayoutNeeded
-                    || ((w.isConfigChanged() || w.setReportResizeHints())
-                    && !w.isGoneForLayoutLw() &&
-                    ((w.mAttrs.privateFlags & PRIVATE_FLAG_KEYGUARD) != 0 ||
-                            (w.mHasSurface && w.mAppToken != null &&
-                                    w.mAppToken.layoutConfigChanges)))) {
-                if (!w.mLayoutAttached) {
-                    if (initial) {
-                        //Slog.i(TAG, "Window " + this + " clearing mContentChanged - initial");
-                        w.mContentChanged = false;
-                    }
-                    if (w.mAttrs.type == TYPE_DREAM) {
-                        // Don't layout windows behind a dream, so that if it does stuff like hide
-                        // the status bar we won't get a bad transition when it goes away.
-                        mTmpWindow = w;
-                    }
-                    w.mLayoutNeeded = false;
-                    w.prelayout();
-                    mService.mPolicy.layoutWindowLw(w, null);
-                    w.mLayoutSeq = mService.mLayoutSeq;
-
-                    // Window frames may have changed. Update dim layer with the new bounds.
-                    final Task task = w.getTask();
-                    if (task != null) {
-                        mDimLayerController.updateDimLayer(task);
-                    }
-
-                    if (DEBUG_LAYOUT) Slog.v(TAG, "  LAYOUT: mFrame=" + w.mFrame
-                            + " mContainingFrame=" + w.mContainingFrame
-                            + " mDisplayFrame=" + w.mDisplayFrame);
-                }
-            }
-        }, true /* traverseTopToBottom */);
+        forAllWindows(mPerformLayout, true /* traverseTopToBottom */);
 
         // Used to indicate that we have processed the dream window and all additional attached
         // windows are behind it.
-        final WindowState dreamWin = mTmpWindow;
+        mTmpWindow2 = mTmpWindow;
         mTmpWindow = null;
 
         // Now perform layout of attached windows, which usually depend on the position of the
         // window they are attached to. XXX does not deal with windows that are attached to windows
         // that are themselves attached.
-        forAllWindows(w -> {
-            if (w.mLayoutAttached) {
-                if (DEBUG_LAYOUT) Slog.v(TAG, "2ND PASS " + w + " mHaveFrame=" + w.mHaveFrame
-                        + " mViewVisibility=" + w.mViewVisibility
-                        + " mRelayoutCalled=" + w.mRelayoutCalled);
-                // If this view is GONE, then skip it -- keep the current frame, and let the
-                // caller know so they can ignore it if they want.  (We do the normal layout for
-                // INVISIBLE windows, since that means "perform layout as normal, just don't
-                // display").
-                if (mTmpWindow != null && mService.mPolicy.canBeHiddenByKeyguardLw(w)) {
-                    return;
-                }
-                if ((w.mViewVisibility != GONE && w.mRelayoutCalled) || !w.mHaveFrame
-                        || w.mLayoutNeeded) {
-                    if (initial) {
-                        //Slog.i(TAG, "Window " + this + " clearing mContentChanged - initial");
-                        w.mContentChanged = false;
-                    }
-                    w.mLayoutNeeded = false;
-                    w.prelayout();
-                    mService.mPolicy.layoutWindowLw(w, w.getParentWindow());
-                    w.mLayoutSeq = mService.mLayoutSeq;
-                    if (DEBUG_LAYOUT)
-                        Slog.v(TAG, "  LAYOUT: mFrame=" + w.mFrame
-                                + " mContainingFrame=" + w.mContainingFrame
-                                + " mDisplayFrame=" + w.mDisplayFrame);
-                }
-            } else if (w.mAttrs.type == TYPE_DREAM) {
-                // Don't layout windows behind a dream, so that if it does stuff like hide the
-                // status bar we won't get a bad transition when it goes away.
-                mTmpWindow = dreamWin;
-            }
-        }, true /* traverseTopToBottom */);
+        forAllWindows(mPerformLayoutAttached, true /* traverseTopToBottom */);
 
         // Window frames may have changed. Tell the input dispatcher about it.
         mService.mInputMonitor.layoutInputConsumers(dw, dh);
@@ -2592,11 +2626,23 @@
          * Compares two child window tokens returns -1 if the first is lesser than the second in
          * terms of z-order and 1 otherwise.
          */
-        final Comparator<WindowToken> mWindowComparator = (token1, token2) ->
+        private final Comparator<WindowToken> mWindowComparator = (token1, token2) ->
                 // Tokens with higher base layer are z-ordered on-top.
                 mService.mPolicy.windowTypeToLayerLw(token1.windowType)
                 < mService.mPolicy.windowTypeToLayerLw(token2.windowType) ? -1 : 1;
 
+        private final Predicate<WindowState> mGetOrientingWindow = w -> {
+            if (!w.isVisibleLw() || !w.mPolicyVisibilityAfterAnim) {
+                return false;
+            }
+            final int req = w.mAttrs.screenOrientation;
+            if(req == SCREEN_ORIENTATION_UNSPECIFIED || req == SCREEN_ORIENTATION_BEHIND
+                    || req == SCREEN_ORIENTATION_UNSET) {
+                return false;
+            }
+            return true;
+        };
+
         private final String mName;
         NonAppWindowContainers(String name) {
             mName = name;
@@ -2610,17 +2656,7 @@
         int getOrientation() {
             final WindowManagerPolicy policy = mService.mPolicy;
             // Find a window requesting orientation.
-            final WindowState win = getWindow((w) -> {
-                if (!w.isVisibleLw() || !w.mPolicyVisibilityAfterAnim) {
-                    return false;
-                }
-                final int req = w.mAttrs.screenOrientation;
-                if(req == SCREEN_ORIENTATION_UNSPECIFIED || req == SCREEN_ORIENTATION_BEHIND
-                        || req == SCREEN_ORIENTATION_UNSET) {
-                    return false;
-                }
-                return true;
-            });
+            final WindowState win = getWindow(mGetOrientingWindow);
 
             if (win != null) {
                 final int req = win.mAttrs.screenOrientation;
diff --git a/services/core/java/com/android/server/wm/InputMonitor.java b/services/core/java/com/android/server/wm/InputMonitor.java
index 01d1c30..495be09 100644
--- a/services/core/java/com/android/server/wm/InputMonitor.java
+++ b/services/core/java/com/android/server/wm/InputMonitor.java
@@ -45,12 +45,14 @@
 import android.view.WindowManager;
 
 import android.view.WindowManagerPolicy;
+
 import com.android.server.input.InputApplicationHandle;
 import com.android.server.input.InputManagerService;
 import com.android.server.input.InputWindowHandle;
 
 import java.io.PrintWriter;
 import java.util.Arrays;
+import java.util.function.Consumer;
 
 final class InputMonitor implements InputManagerService.WindowManagerCallbacks {
     private final WindowManagerService mService;
@@ -79,6 +81,9 @@
     private boolean mAddPipInputConsumerHandle;
     private boolean mAddWallpaperInputConsumerHandle;
     private boolean mDisableWallpaperTouchEvents;
+    private final Rect mTmpRect = new Rect();
+    private final UpdateInputForAllWindowsConsumer mUpdateInputForAllWindowsConsumer =
+            new UpdateInputForAllWindowsConsumer();
 
     // Set to true when the first input device configuration change notification
     // is received to indicate that the input devices are ready.
@@ -381,94 +386,11 @@
         }
 
         // Add all windows on the default display.
-        updateInputWindows(inDrag);
+        mUpdateInputForAllWindowsConsumer.updateInputWindows(inDrag);
 
         if (false) Slog.d(TAG_WM, "<<<<<<< EXITED updateInputWindowsLw");
     }
 
-    private void updateInputWindows(boolean inDrag) {
-
-        clearInputWindowHandlesLw();
-
-        // TODO: multi-display
-        final InputConsumerImpl navInputConsumer =
-                getInputConsumer(INPUT_CONSUMER_NAVIGATION, DEFAULT_DISPLAY);
-        final InputConsumerImpl pipInputConsumer =
-                getInputConsumer(INPUT_CONSUMER_PIP, DEFAULT_DISPLAY);
-        final InputConsumerImpl wallpaperInputConsumer =
-                getInputConsumer(INPUT_CONSUMER_WALLPAPER, DEFAULT_DISPLAY);
-        mAddInputConsumerHandle = navInputConsumer != null;
-        mAddPipInputConsumerHandle = pipInputConsumer != null;
-        mAddWallpaperInputConsumerHandle = wallpaperInputConsumer != null;
-        final Rect pipTouchableBounds = mAddPipInputConsumerHandle ? new Rect() : null;
-        mDisableWallpaperTouchEvents = false;
-
-        final WallpaperController wallpaperController = mService.mRoot.mWallpaperController;
-        mService.mRoot.forAllWindows(w -> {
-            final InputChannel inputChannel = w.mInputChannel;
-            final InputWindowHandle inputWindowHandle = w.mInputWindowHandle;
-            if (inputChannel == null || inputWindowHandle == null || w.mRemoved
-                    || w.isAdjustedForMinimizedDock()) {
-                // Skip this window because it cannot possibly receive input.
-                return;
-            }
-
-            if (mAddPipInputConsumerHandle
-                    && w.getStackId() == PINNED_STACK_ID
-                    && inputWindowHandle.layer <= pipInputConsumer.mWindowHandle.layer) {
-                // Update the bounds of the Pip input consumer to match the Pinned stack
-                w.getStack().getBounds(pipTouchableBounds);
-                pipInputConsumer.mWindowHandle.touchableRegion.set(pipTouchableBounds);
-                addInputWindowHandle(pipInputConsumer.mWindowHandle);
-                mAddPipInputConsumerHandle = false;
-            }
-
-            if (mAddInputConsumerHandle
-                    && inputWindowHandle.layer <= navInputConsumer.mWindowHandle.layer) {
-                addInputWindowHandle(navInputConsumer.mWindowHandle);
-                mAddInputConsumerHandle = false;
-            }
-
-            if (mAddWallpaperInputConsumerHandle) {
-                if (w.mAttrs.type == TYPE_WALLPAPER && w.isVisibleLw()) {
-                    // Add the wallpaper input consumer above the first visible wallpaper.
-                    addInputWindowHandle(wallpaperInputConsumer.mWindowHandle);
-                    mAddWallpaperInputConsumerHandle = false;
-                }
-            }
-
-            final int flags = w.mAttrs.flags;
-            final int privateFlags = w.mAttrs.privateFlags;
-            final int type = w.mAttrs.type;
-
-            final boolean hasFocus = w == mInputFocus;
-            final boolean isVisible = w.isVisibleLw();
-            if ((privateFlags & PRIVATE_FLAG_DISABLE_WALLPAPER_TOUCH_EVENTS) != 0) {
-                mDisableWallpaperTouchEvents = true;
-            }
-            final boolean hasWallpaper = wallpaperController.isWallpaperTarget(w)
-                    && (privateFlags & PRIVATE_FLAG_KEYGUARD) == 0
-                    && !mDisableWallpaperTouchEvents;
-
-            // If there's a drag in progress and 'child' is a potential drop target,
-            // make sure it's been told about the drag
-            if (inDrag && isVisible && w.getDisplayContent().isDefaultDisplay) {
-                mService.mDragState.sendDragStartedIfNeededLw(w);
-            }
-
-            addInputWindowHandle(
-                    inputWindowHandle, w, flags, type, isVisible, hasFocus, hasWallpaper);
-        }, true /* traverseTopToBottom */);
-
-        if (mAddWallpaperInputConsumerHandle) {
-            // No visible wallpaper found, add the wallpaper input consumer at the end.
-            addInputWindowHandle(wallpaperInputConsumer.mWindowHandle);
-        }
-
-        // Send windows to native code.
-        mService.mInputManager.setInputWindows(mInputWindowHandles);
-    }
-
     /* Notifies that the input device configuration has changed. */
     @Override
     public void notifyConfigurationChanged() {
@@ -657,4 +579,98 @@
             pw.println(prefix + "mInputFreezeReason=" + mInputFreezeReason);
         }
     }
+
+    private final class UpdateInputForAllWindowsConsumer implements Consumer<WindowState> {
+
+        InputConsumerImpl navInputConsumer;
+        InputConsumerImpl pipInputConsumer;
+        InputConsumerImpl wallpaperInputConsumer;
+        Rect pipTouchableBounds;
+        boolean inDrag;
+        WallpaperController wallpaperController;
+
+        private void updateInputWindows(boolean inDrag) {
+
+            clearInputWindowHandlesLw();
+
+            // TODO: multi-display
+            navInputConsumer = getInputConsumer(INPUT_CONSUMER_NAVIGATION, DEFAULT_DISPLAY);
+            pipInputConsumer = getInputConsumer(INPUT_CONSUMER_PIP, DEFAULT_DISPLAY);
+            wallpaperInputConsumer = getInputConsumer(INPUT_CONSUMER_WALLPAPER, DEFAULT_DISPLAY);
+            mAddInputConsumerHandle = navInputConsumer != null;
+            mAddPipInputConsumerHandle = pipInputConsumer != null;
+            mAddWallpaperInputConsumerHandle = wallpaperInputConsumer != null;
+            mTmpRect.setEmpty();
+            pipTouchableBounds = mAddPipInputConsumerHandle ? mTmpRect : null;
+            mDisableWallpaperTouchEvents = false;
+            this.inDrag = inDrag;
+            wallpaperController = mService.mRoot.mWallpaperController;
+
+            mService.mRoot.forAllWindows(this, true /* traverseTopToBottom */);
+            if (mAddWallpaperInputConsumerHandle) {
+                // No visible wallpaper found, add the wallpaper input consumer at the end.
+                addInputWindowHandle(wallpaperInputConsumer.mWindowHandle);
+            }
+
+            // Send windows to native code.
+            mService.mInputManager.setInputWindows(mInputWindowHandles);
+        }
+
+        @Override
+        public void accept(WindowState w) {
+            final InputChannel inputChannel = w.mInputChannel;
+            final InputWindowHandle inputWindowHandle = w.mInputWindowHandle;
+            if (inputChannel == null || inputWindowHandle == null || w.mRemoved
+                    || w.isAdjustedForMinimizedDock()) {
+                // Skip this window because it cannot possibly receive input.
+                return;
+            }
+
+            if (mAddPipInputConsumerHandle
+                    && w.getStackId() == PINNED_STACK_ID
+                    && inputWindowHandle.layer <= pipInputConsumer.mWindowHandle.layer) {
+                // Update the bounds of the Pip input consumer to match the Pinned stack
+                w.getStack().getBounds(pipTouchableBounds);
+                pipInputConsumer.mWindowHandle.touchableRegion.set(pipTouchableBounds);
+                addInputWindowHandle(pipInputConsumer.mWindowHandle);
+                mAddPipInputConsumerHandle = false;
+            }
+
+            if (mAddInputConsumerHandle
+                    && inputWindowHandle.layer <= navInputConsumer.mWindowHandle.layer) {
+                addInputWindowHandle(navInputConsumer.mWindowHandle);
+                mAddInputConsumerHandle = false;
+            }
+
+            if (mAddWallpaperInputConsumerHandle) {
+                if (w.mAttrs.type == TYPE_WALLPAPER && w.isVisibleLw()) {
+                    // Add the wallpaper input consumer above the first visible wallpaper.
+                    addInputWindowHandle(wallpaperInputConsumer.mWindowHandle);
+                    mAddWallpaperInputConsumerHandle = false;
+                }
+            }
+
+            final int flags = w.mAttrs.flags;
+            final int privateFlags = w.mAttrs.privateFlags;
+            final int type = w.mAttrs.type;
+
+            final boolean hasFocus = w == mInputFocus;
+            final boolean isVisible = w.isVisibleLw();
+            if ((privateFlags & PRIVATE_FLAG_DISABLE_WALLPAPER_TOUCH_EVENTS) != 0) {
+                mDisableWallpaperTouchEvents = true;
+            }
+            final boolean hasWallpaper = wallpaperController.isWallpaperTarget(w)
+                    && (privateFlags & PRIVATE_FLAG_KEYGUARD) == 0
+                    && !mDisableWallpaperTouchEvents;
+
+            // If there's a drag in progress and 'child' is a potential drop target,
+            // make sure it's been told about the drag
+            if (inDrag && isVisible && w.getDisplayContent().isDefaultDisplay) {
+                mService.mDragState.sendDragStartedIfNeededLw(w);
+            }
+
+            addInputWindowHandle(
+                    inputWindowHandle, w, flags, type, isVisible, hasFocus, hasWallpaper);
+        }
+    }
 }
diff --git a/services/core/java/com/android/server/wm/RootWindowContainer.java b/services/core/java/com/android/server/wm/RootWindowContainer.java
index ee6d4ea..349740b 100644
--- a/services/core/java/com/android/server/wm/RootWindowContainer.java
+++ b/services/core/java/com/android/server/wm/RootWindowContainer.java
@@ -40,6 +40,7 @@
 import java.io.FileDescriptor;
 import java.io.PrintWriter;
 import java.util.ArrayList;
+import java.util.function.Consumer;
 
 import static android.app.AppOpsManager.MODE_ALLOWED;
 import static android.app.AppOpsManager.MODE_DEFAULT;
@@ -129,6 +130,23 @@
     private final WindowLayersController mLayersController;
     final WallpaperController mWallpaperController;
 
+    private String mCloseSystemDialogsReason;
+    private final Consumer<WindowState> mCloseSystemDialogsConsumer = w -> {
+        if (w.mHasSurface) {
+            try {
+                w.mClient.closeSystemDialogs(mCloseSystemDialogsReason);
+            } catch (RemoteException e) {
+            }
+        }
+    };
+
+    private static final Consumer<WindowState> sRemoveReplacedWindowsConsumer = w -> {
+        final AppWindowToken aToken = w.mAppToken;
+        if (aToken != null) {
+            aToken.removeReplacedWindowIfNeeded(w);
+        }
+    };
+
     RootWindowContainer(WindowManagerService service) {
         mService = service;
         mLayersController = new WindowLayersController(mService);
@@ -376,26 +394,15 @@
     }
 
     void closeSystemDialogs(String reason) {
-        forAllWindows((w) -> {
-            if (w.mHasSurface) {
-                try {
-                    w.mClient.closeSystemDialogs(reason);
-                } catch (RemoteException e) {
-                }
-            }
-        }, false /* traverseTopToBottom */);
+        mCloseSystemDialogsReason = reason;
+        forAllWindows(mCloseSystemDialogsConsumer, false /* traverseTopToBottom */);
     }
 
     void removeReplacedWindows() {
         if (SHOW_TRANSACTIONS) Slog.i(TAG, ">>> OPEN TRANSACTION removeReplacedWindows");
         mService.openSurfaceTransaction();
         try {
-            forAllWindows((w) -> {
-                final AppWindowToken aToken = w.mAppToken;
-                if (aToken != null) {
-                    aToken.removeReplacedWindowIfNeeded(w);
-                }
-            }, true /* traverseTopToBottom */);
+            forAllWindows(sRemoveReplacedWindowsConsumer, true /* traverseTopToBottom */);
         } finally {
             mService.closeSurfaceTransaction();
             if (SHOW_TRANSACTIONS) Slog.i(TAG, "<<< CLOSE TRANSACTION removeReplacedWindows");
diff --git a/services/core/java/com/android/server/wm/WallpaperController.java b/services/core/java/com/android/server/wm/WallpaperController.java
index 7fd8028..feceb8e 100644
--- a/services/core/java/com/android/server/wm/WallpaperController.java
+++ b/services/core/java/com/android/server/wm/WallpaperController.java
@@ -16,6 +16,8 @@
 
 package com.android.server.wm;
 
+import com.android.internal.util.ToBooleanFunction;
+
 import static android.app.ActivityManager.StackId.FREEFORM_WORKSPACE_STACK_ID;
 import static android.view.ViewGroup.LayoutParams.MATCH_PARENT;
 import static android.view.WindowManager.LayoutParams.FLAG_SHOW_WALLPAPER;
@@ -51,7 +53,7 @@
  */
 class WallpaperController {
     private static final String TAG = TAG_WITH_CLASS_NAME ? "WallpaperController" : TAG_WM;
-    final private WindowManagerService mService;
+    private WindowManagerService mService;
 
     private final ArrayList<WallpaperWindowToken> mWallpaperTokens = new ArrayList<>();
 
@@ -97,6 +99,78 @@
 
     private final FindWallpaperTargetResult mFindResults = new FindWallpaperTargetResult();
 
+    private final ToBooleanFunction<WindowState> mFindWallpaperTargetFunction = w -> {
+        final WindowAnimator winAnimator = mService.mAnimator;
+        if ((w.mAttrs.type == TYPE_WALLPAPER)) {
+            if (mFindResults.topWallpaper == null || mFindResults.resetTopWallpaper) {
+                mFindResults.setTopWallpaper(w);
+                mFindResults.resetTopWallpaper = false;
+            }
+            return false;
+        }
+
+        mFindResults.resetTopWallpaper = true;
+        if (w != winAnimator.mWindowDetachedWallpaper && w.mAppToken != null) {
+            // If this window's app token is hidden and not animating, it is of no interest to us.
+            if (w.mAppToken.hidden && w.mAppToken.mAppAnimator.animation == null) {
+                if (DEBUG_WALLPAPER) Slog.v(TAG,
+                        "Skipping hidden and not animating token: " + w);
+                return false;
+            }
+        }
+        if (DEBUG_WALLPAPER) Slog.v(TAG, "Win " + w + ": isOnScreen=" + w.isOnScreen()
+                + " mDrawState=" + w.mWinAnimator.mDrawState);
+
+        if (w.mWillReplaceWindow && mWallpaperTarget == null
+                && !mFindResults.useTopWallpaperAsTarget) {
+            // When we are replacing a window and there was wallpaper before replacement, we want to
+            // keep the window until the new windows fully appear and can determine the visibility,
+            // to avoid flickering.
+            mFindResults.setUseTopWallpaperAsTarget(true);
+        }
+
+        final boolean keyguardGoingAwayWithWallpaper = (w.mAppToken != null
+                && AppTransition.isKeyguardGoingAwayTransit(
+                w.mAppToken.mAppAnimator.getTransit())
+                && (w.mAppToken.mAppAnimator.getTransitFlags()
+                & TRANSIT_FLAG_KEYGUARD_GOING_AWAY_WITH_WALLPAPER) != 0);
+
+        boolean needsShowWhenLockedWallpaper = false;
+        if ((w.mAttrs.flags & FLAG_SHOW_WHEN_LOCKED) != 0
+                && mService.mPolicy.isKeyguardLocked()
+                && mService.mPolicy.isKeyguardOccluded()) {
+            // The lowest show when locked window decides whether we need to put the wallpaper
+            // behind.
+            needsShowWhenLockedWallpaper = !isFullscreen(w.mAttrs)
+                    || (w.mAppToken != null && !w.mAppToken.fillsParent());
+        }
+
+        if (keyguardGoingAwayWithWallpaper || needsShowWhenLockedWallpaper) {
+            // Keep the wallpaper during Keyguard exit but also when it's needed for a
+            // non-fullscreen show when locked activity.
+            mFindResults.setUseTopWallpaperAsTarget(true);
+        }
+
+        final boolean hasWallpaper = (w.mAttrs.flags & FLAG_SHOW_WALLPAPER) != 0;
+        if (hasWallpaper && w.isOnScreen() && (mWallpaperTarget == w || w.isDrawFinishedLw())) {
+            if (DEBUG_WALLPAPER) Slog.v(TAG, "Found wallpaper target: " + w);
+            mFindResults.setWallpaperTarget(w);
+            if (w == mWallpaperTarget && w.mWinAnimator.isAnimationSet()) {
+                // The current wallpaper target is animating, so we'll look behind it for
+                // another possible target and figure out what is going on later.
+                if (DEBUG_WALLPAPER) Slog.v(TAG,
+                        "Win " + w + ": token animating, looking behind.");
+            }
+            // Found a target! End search.
+            return true;
+        } else if (w == winAnimator.mWindowDetachedWallpaper) {
+            if (DEBUG_WALLPAPER_LIGHT) Slog.v(TAG,
+                    "Found animating detached wallpaper target win: " + w);
+            mFindResults.setUseTopWallpaperAsTarget(true);
+        }
+        return false;
+    };
+
     public WallpaperController(WindowManagerService service) {
         mService = service;
     }
@@ -371,89 +445,18 @@
         return mWallpaperAnimLayerAdjustment;
     }
 
-    private void findWallpaperTarget(DisplayContent dc , FindWallpaperTargetResult result) {
-        final WindowAnimator winAnimator = mService.mAnimator;
-        result.reset();
+    private void findWallpaperTarget(DisplayContent dc) {
+        mFindResults.reset();
         if (mService.isStackVisibleLocked(FREEFORM_WORKSPACE_STACK_ID)) {
             // In freeform mode we set the wallpaper as its own target, so we don't need an
             // additional window to make it visible.
-            result.setUseTopWallpaperAsTarget(true);
+            mFindResults.setUseTopWallpaperAsTarget(true);
         }
 
-        dc.forAllWindows(w -> {
-            if ((w.mAttrs.type == TYPE_WALLPAPER)) {
-                if (result.topWallpaper == null || result.resetTopWallpaper) {
-                    result.setTopWallpaper(w);
-                    result.resetTopWallpaper = false;
-                }
-                return false;
-            }
+        dc.forAllWindows(mFindWallpaperTargetFunction, true /* traverseTopToBottom */);
 
-            result.resetTopWallpaper = true;
-            if (w != winAnimator.mWindowDetachedWallpaper && w.mAppToken != null) {
-                // If this window's app token is hidden and not animating,
-                // it is of no interest to us.
-                if (w.mAppToken.hidden && w.mAppToken.mAppAnimator.animation == null) {
-                    if (DEBUG_WALLPAPER) Slog.v(TAG,
-                            "Skipping hidden and not animating token: " + w);
-                    return false;
-                }
-            }
-            if (DEBUG_WALLPAPER) Slog.v(TAG, "Win " + w + ": isOnScreen=" + w.isOnScreen()
-                    + " mDrawState=" + w.mWinAnimator.mDrawState);
-
-            if (w.mWillReplaceWindow && mWallpaperTarget == null
-                    && !result.useTopWallpaperAsTarget) {
-                // When we are replacing a window and there was wallpaper before replacement, we
-                // want to keep the window until the new windows fully appear and can determine the
-                // visibility, to avoid flickering.
-                result.setUseTopWallpaperAsTarget(true);
-            }
-
-            final boolean keyguardGoingAwayWithWallpaper = (w.mAppToken != null
-                    && AppTransition.isKeyguardGoingAwayTransit(
-                            w.mAppToken.mAppAnimator.getTransit())
-                    && (w.mAppToken.mAppAnimator.getTransitFlags()
-                            & TRANSIT_FLAG_KEYGUARD_GOING_AWAY_WITH_WALLPAPER) != 0);
-
-            boolean needsShowWhenLockedWallpaper = false;
-            if ((w.mAttrs.flags & FLAG_SHOW_WHEN_LOCKED) != 0
-                    && mService.mPolicy.isKeyguardLocked()
-                    && mService.mPolicy.isKeyguardOccluded()) {
-                // The lowest show when locked window decides whether we need to put the wallpaper
-                // behind.
-                needsShowWhenLockedWallpaper = !isFullscreen(w.mAttrs)
-                        || (w.mAppToken != null && !w.mAppToken.fillsParent());
-            }
-
-            if (keyguardGoingAwayWithWallpaper || needsShowWhenLockedWallpaper) {
-                // Keep the wallpaper during Keyguard exit but also when it's needed for a
-                // non-fullscreen show when locked activity.
-                result.setUseTopWallpaperAsTarget(true);
-            }
-
-            final boolean hasWallpaper = (w.mAttrs.flags & FLAG_SHOW_WALLPAPER) != 0;
-            if (hasWallpaper && w.isOnScreen() && (mWallpaperTarget == w || w.isDrawFinishedLw())) {
-                if (DEBUG_WALLPAPER) Slog.v(TAG, "Found wallpaper target: " + w);
-                result.setWallpaperTarget(w);
-                if (w == mWallpaperTarget && w.mWinAnimator.isAnimationSet()) {
-                    // The current wallpaper target is animating, so we'll look behind it for
-                    // another possible target and figure out what is going on later.
-                    if (DEBUG_WALLPAPER) Slog.v(TAG,
-                            "Win " + w + ": token animating, looking behind.");
-                }
-                // Found a target! End search.
-                return true;
-            } else if (w == winAnimator.mWindowDetachedWallpaper) {
-                if (DEBUG_WALLPAPER_LIGHT) Slog.v(TAG,
-                        "Found animating detached wallpaper target win: " + w);
-                result.setUseTopWallpaperAsTarget(true);
-            }
-            return false;
-        }, true /* traverseTopToBottom */);
-
-        if (result.wallpaperTarget == null && result.useTopWallpaperAsTarget) {
-            result.setWallpaperTarget(result.topWallpaper);
+        if (mFindResults.wallpaperTarget == null && mFindResults.useTopWallpaperAsTarget) {
+            mFindResults.setWallpaperTarget(mFindResults.topWallpaper);
         }
     }
 
@@ -552,7 +555,7 @@
 
         // First find top-most window that has asked to be on top of the wallpaper;
         // all wallpapers go behind it.
-        findWallpaperTarget(dc, mFindResults);
+        findWallpaperTarget(dc);
         updateWallpaperWindowsTarget(dc, mFindResults);
 
         // The window is visible to the compositor...but is it visible to the user?
diff --git a/services/core/java/com/android/server/wm/WindowContainer.java b/services/core/java/com/android/server/wm/WindowContainer.java
index 03769e9..02376e9 100644
--- a/services/core/java/com/android/server/wm/WindowContainer.java
+++ b/services/core/java/com/android/server/wm/WindowContainer.java
@@ -25,7 +25,6 @@
 import java.util.Comparator;
 import java.util.LinkedList;
 import java.util.function.Consumer;
-import java.util.function.Function;
 import java.util.function.Predicate;
 
 import static android.content.pm.ActivityInfo.SCREEN_ORIENTATION_BEHIND;
@@ -501,7 +500,7 @@
      * @param   traverseTopToBottom If true traverses the hierarchy from top-to-bottom in terms of
      *                              z-order, else from bottom-to-top.
      * @return  True if the search ended before we reached the end of the hierarchy due to
-     *          {@link Function#apply} returning true.
+     *          {@link ToBooleanFunction#apply} returning true.
      */
     boolean forAllWindows(ToBooleanFunction<WindowState> callback, boolean traverseTopToBottom) {
         if (traverseTopToBottom) {
diff --git a/services/core/java/com/android/server/wm/WindowLayersController.java b/services/core/java/com/android/server/wm/WindowLayersController.java
index c2988ad..e0c3d60 100644
--- a/services/core/java/com/android/server/wm/WindowLayersController.java
+++ b/services/core/java/com/android/server/wm/WindowLayersController.java
@@ -20,6 +20,7 @@
 import android.view.Display;
 
 import java.util.ArrayDeque;
+import java.util.function.Consumer;
 
 import static android.app.ActivityManager.StackId.DOCKED_STACK_ID;
 import static android.app.ActivityManager.StackId.PINNED_STACK_ID;
@@ -69,45 +70,47 @@
     private boolean mAboveImeTarget;
     private ArrayDeque<WindowState> mAboveImeTargetAppWindows = new ArrayDeque();
 
+    private final Consumer<WindowState> mAssignWindowLayersConsumer = w -> {
+        boolean layerChanged = false;
+
+        int oldLayer = w.mLayer;
+        if (w.mBaseLayer == mCurBaseLayer) {
+            mCurLayer += WINDOW_LAYER_MULTIPLIER;
+        } else {
+            mCurBaseLayer = mCurLayer = w.mBaseLayer;
+        }
+        assignAnimLayer(w, mCurLayer);
+
+        // TODO: Preserved old behavior of code here but not sure comparing oldLayer to
+        // mAnimLayer and mLayer makes sense...though the worst case would be unintentional
+        // layer reassignment.
+        if (w.mLayer != oldLayer || w.mWinAnimator.mAnimLayer != oldLayer) {
+            layerChanged = true;
+            mAnyLayerChanged = true;
+        }
+
+        if (w.mAppToken != null) {
+            mHighestApplicationLayer = Math.max(mHighestApplicationLayer,
+                    w.mWinAnimator.mAnimLayer);
+        }
+        if (mImeTarget != null && w.mBaseLayer == mImeTarget.mBaseLayer) {
+            mHighestLayerInImeTargetBaseLayer = Math.max(mHighestLayerInImeTargetBaseLayer,
+                    w.mWinAnimator.mAnimLayer);
+        }
+
+        collectSpecialWindows(w);
+
+        if (layerChanged) {
+            w.scheduleAnimationIfDimming();
+        }
+    };
+
     final void assignWindowLayers(DisplayContent dc) {
         if (DEBUG_LAYERS) Slog.v(TAG_WM, "Assigning layers based",
                 new RuntimeException("here").fillInStackTrace());
 
         reset();
-        dc.forAllWindows((w) -> {
-            boolean layerChanged = false;
-
-            int oldLayer = w.mLayer;
-            if (w.mBaseLayer == mCurBaseLayer) {
-                mCurLayer += WINDOW_LAYER_MULTIPLIER;
-            } else {
-                mCurBaseLayer = mCurLayer = w.mBaseLayer;
-            }
-            assignAnimLayer(w, mCurLayer);
-
-            // TODO: Preserved old behavior of code here but not sure comparing oldLayer to
-            // mAnimLayer and mLayer makes sense...though the worst case would be unintentional
-            // layer reassignment.
-            if (w.mLayer != oldLayer || w.mWinAnimator.mAnimLayer != oldLayer) {
-                layerChanged = true;
-                mAnyLayerChanged = true;
-            }
-
-            if (w.mAppToken != null) {
-                mHighestApplicationLayer = Math.max(mHighestApplicationLayer,
-                        w.mWinAnimator.mAnimLayer);
-            }
-            if (mImeTarget != null && w.mBaseLayer == mImeTarget.mBaseLayer) {
-                mHighestLayerInImeTargetBaseLayer = Math.max(mHighestLayerInImeTargetBaseLayer,
-                        w.mWinAnimator.mAnimLayer);
-            }
-
-            collectSpecialWindows(w);
-
-            if (layerChanged) {
-                w.scheduleAnimationIfDimming();
-            }
-        }, false /* traverseTopToBottom */);
+        dc.forAllWindows(mAssignWindowLayersConsumer, false /* traverseTopToBottom */);
 
         adjustSpecialWindows();