Maintain Z order for all layers in offscreen hierarchy

There was already code in place to ensure that z order was maintained in
the offscreen tree. However, this didn't work if a layer was moved to become
a child of an offscreen layer. The z order was only maintained beneath
the new offscreen layer, not the entire subtree.

This change uses the root of the offscreen subtree to ensure the z order
is maintained even when new offscreen layers are added to the subtree.

Fixes: 157188227
Test: IME gets correct relative z
Test: RelativeZTest.LayerWithRelativeReparentedToOffscreen
Change-Id: I62553ce245dacd2a8684d8bb02de67f60ddc6774
diff --git a/services/surfaceflinger/Layer.cpp b/services/surfaceflinger/Layer.cpp
index dcc213f..f5acbae 100644
--- a/services/surfaceflinger/Layer.cpp
+++ b/services/surfaceflinger/Layer.cpp
@@ -229,13 +229,23 @@
     mFlinger->markLayerPendingRemovalLocked(this);
 }
 
+sp<Layer> Layer::getRootLayer() {
+    sp<Layer> parent = getParent();
+    if (parent == nullptr) {
+        return this;
+    }
+    return parent->getRootLayer();
+}
+
 void Layer::onRemovedFromCurrentState() {
-    auto layersInTree = getLayersInTree(LayerVector::StateSet::Current);
+    // Use the root layer since we want to maintain the hierarchy for the entire subtree.
+    auto layersInTree = getRootLayer()->getLayersInTree(LayerVector::StateSet::Current);
     std::sort(layersInTree.begin(), layersInTree.end());
-    for (const auto& layer : layersInTree) {
+
+    traverse(LayerVector::StateSet::Current, [&](Layer* layer) {
         layer->removeFromCurrentState();
         layer->removeRelativeZ(layersInTree);
-    }
+    });
 }
 
 void Layer::addToCurrentState() {