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.
Merged-In: I7afef556fa72f687bcfeb0a642465488cc72f40b
diff --git a/services/surfaceflinger/LayerVector.cpp b/services/surfaceflinger/LayerVector.cpp
index 47156c1..8494524 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;