Use correct StateSet for LayerVector compare.

Currently LayerVector compare function was using the current StateSet.
This is incorect since the LayerVector may be created with the intention
of sorting the layers by drawing state. Instead, create the LayerVector
with a specified StateSet so the compare function always uses the
correct state.

This fixes an issue where the layers were getting added and sorted by
current state z order but the caller expected the order to be by drawing
state z order.

Change-Id: I7afef556fa72f687bcfeb0a642465488cc72f40b
Fixes: 80516823
Test: No longer flicker when IME closes. Logs show correct z order.

4 files changed, 32 insertions, 15 deletions

diff --git a/services/surfaceflinger/Layer.cpp b/services/surfaceflinger/Layer.cpp
index 8c1199b8a1..10b3443ac4 100644
--- a/services/surfaceflinger/Layer.cpp
+++ b/services/surfaceflinger/Layer.cpp
@@ -94,6 +94,8 @@ Layer::Layer(SurfaceFlinger* flinger, const sp<Client>& client, const String8& n
         mLastFrameNumberReceived(0),
         mAutoRefresh(false),
         mFreezeGeometryUpdates(false),
+        mCurrentChildren(LayerVector::StateSet::Current),
+        mDrawingChildren(LayerVector::StateSet::Drawing),
         mBE{this, name.string()} {
 
     mCurrentCrop.makeInvalid();
@@ -133,7 +135,6 @@ Layer::Layer(SurfaceFlinger* flinger, const sp<Client>& client, const String8& n
     CompositorTiming compositorTiming;
     flinger->getCompositorTiming(&compositorTiming);
     mFrameEventHistory.initializeCompositorTiming(compositorTiming);
-
 }
 
 void Layer::onFirstRef() {}
@@ -1677,7 +1678,7 @@ __attribute__((no_sanitize("unsigned-integer-overflow"))) LayerVector Layer::mak
         return children;
     }
 
-    LayerVector traverse;
+    LayerVector traverse(stateSet);
     for (const wp<Layer>& weakRelative : state.zOrderRelatives) {
         sp<Layer> strongRelative = weakRelative.promote();
         if (strongRelative != nullptr) {
@@ -1775,7 +1776,7 @@ LayerVector Layer::makeChildrenTraversalList(LayerVector::StateSet stateSet,
     const LayerVector& children = useDrawing ? mDrawingChildren : mCurrentChildren;
     const State& state = useDrawing ? mDrawingState : mCurrentState;
 
-    LayerVector traverse;
+    LayerVector traverse(stateSet);
     for (const wp<Layer>& weakRelative : state.zOrderRelatives) {
         sp<Layer> strongRelative = weakRelative.promote();
         // Only add relative layers that are also descendents of the top most parent of the tree.
diff --git a/services/surfaceflinger/LayerVector.cpp b/services/surfaceflinger/LayerVector.cpp
index 47156c19d2..84945247a6 100644
--- a/services/surfaceflinger/LayerVector.cpp
+++ b/services/surfaceflinger/LayerVector.cpp
@@ -19,26 +19,37 @@
 
 namespace android {
 
-LayerVector::LayerVector() = default;
+LayerVector::LayerVector(const StateSet stateSet) : mStateSet(stateSet) {}
 
-LayerVector::LayerVector(const LayerVector& rhs) : SortedVector<sp<Layer>>(rhs) {
-}
+LayerVector::LayerVector(const LayerVector& rhs, const StateSet stateSet)
+      : SortedVector<sp<Layer>>(rhs), mStateSet(stateSet) {}
 
 LayerVector::~LayerVector() = default;
 
+// This operator override is needed to prevent mStateSet from getting copied over.
+LayerVector& LayerVector::operator=(const LayerVector& rhs) {
+    SortedVector::operator=(rhs);
+    return *this;
+}
+
 int LayerVector::do_compare(const void* lhs, const void* rhs) const
 {
     // sort layers per layer-stack, then by z-order and finally by sequence
     const auto& l = *reinterpret_cast<const sp<Layer>*>(lhs);
     const auto& r = *reinterpret_cast<const sp<Layer>*>(rhs);
 
-    uint32_t ls = l->getCurrentState().layerStack;
-    uint32_t rs = r->getCurrentState().layerStack;
+    const auto& lState =
+            (mStateSet == StateSet::Current) ? l->getCurrentState() : l->getDrawingState();
+    const auto& rState =
+            (mStateSet == StateSet::Current) ? r->getCurrentState() : r->getDrawingState();
+
+    uint32_t ls = lState.layerStack;
+    uint32_t rs = rState.layerStack;
     if (ls != rs)
         return (ls > rs) ? 1 : -1;
 
-    int32_t lz = l->getCurrentState().z;
-    int32_t rz = r->getCurrentState().z;
+    int32_t lz = lState.z;
+    int32_t rz = rState.z;
     if (lz != rz)
         return (lz > rz) ? 1 : -1;
 
diff --git a/services/surfaceflinger/LayerVector.h b/services/surfaceflinger/LayerVector.h
index a9adb4113a..88d7711bb8 100644
--- a/services/surfaceflinger/LayerVector.h
+++ b/services/surfaceflinger/LayerVector.h
@@ -32,22 +32,27 @@ class Layer;
  */
 class LayerVector : public SortedVector<sp<Layer>> {
 public:
-    LayerVector();
-    LayerVector(const LayerVector& rhs);
-    ~LayerVector() override;
-
     enum class StateSet {
         Invalid,
         Current,
         Drawing,
     };
 
+    explicit LayerVector(const StateSet stateSet);
+    LayerVector(const LayerVector& rhs, const StateSet stateSet);
+    ~LayerVector() override;
+
+    LayerVector& operator=(const LayerVector& rhs);
+
     // Sorts layer by layer-stack, Z order, and finally creation order (sequence).
     int do_compare(const void* lhs, const void* rhs) const override;
 
     using Visitor = std::function<void(Layer*)>;
     void traverseInReverseZOrder(StateSet stateSet, const Visitor& visitor) const;
     void traverseInZOrder(StateSet stateSet, const Visitor& visitor) const;
+
+private:
+    const StateSet mStateSet;
 };
 }
 
diff --git a/services/surfaceflinger/SurfaceFlinger.h b/services/surfaceflinger/SurfaceFlinger.h
index 750f2ec39d..e58287a6b9 100644
--- a/services/surfaceflinger/SurfaceFlinger.h
+++ b/services/surfaceflinger/SurfaceFlinger.h
@@ -371,7 +371,7 @@ private:
 
     class State {
     public:
-        explicit State(LayerVector::StateSet set) : stateSet(set) {}
+        explicit State(LayerVector::StateSet set) : stateSet(set), layersSortedByZ(set) {}
         State& operator=(const State& other) {
             // We explicitly don't copy stateSet so that, e.g., mDrawingState
             // always uses the Drawing StateSet.