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() {
diff --git a/services/surfaceflinger/Layer.h b/services/surfaceflinger/Layer.h
index 3fa935f..8507331 100644
--- a/services/surfaceflinger/Layer.h
+++ b/services/surfaceflinger/Layer.h
@@ -1088,6 +1088,10 @@
     // Find the root of the cloned hierarchy, this means the first non cloned parent.
     // This will return null if first non cloned parent is not found.
     sp<Layer> getClonedRoot();
+
+    // Finds the top most layer in the hierarchy. This will find the root Layer where the parent is
+    // null.
+    sp<Layer> getRootLayer();
 };
 
 } // namespace android
diff --git a/services/surfaceflinger/tests/RelativeZ_test.cpp b/services/surfaceflinger/tests/RelativeZ_test.cpp
index 1180cac..3e0b3c6 100644
--- a/services/surfaceflinger/tests/RelativeZ_test.cpp
+++ b/services/surfaceflinger/tests/RelativeZ_test.cpp
@@ -207,6 +207,71 @@
         sc->checkPixel(1, 1, Color::BLUE.r, Color::BLUE.g, Color::BLUE.b);
     }
 }
+
+// Preserve the relative z order when a layer is reparented to a layer that's already offscreen
+TEST_F(RelativeZTest, LayerWithRelativeReparentedToOffscreen) {
+    std::unique_ptr<ScreenCapture> sc;
+
+    Color testLayerColor = {255, 100, 0, 255};
+
+    // Background layer (RED)
+    // Foregroud layer (GREEN)
+    //   child level 1a (testLayerColor) (relative to child level 2b)
+    //   child level 1b (WHITE)
+    //     child level 2a (BLUE)
+    //     child level 2b (BLACK)
+    sp<SurfaceControl> childLevel1a =
+            createColorLayer("child level 1a", testLayerColor, mForegroundLayer.get());
+    sp<SurfaceControl> childLevel1b =
+            createColorLayer("child level 1b", Color::WHITE, mForegroundLayer.get());
+    sp<SurfaceControl> childLevel2a =
+            createColorLayer("child level 2a", Color::BLUE, childLevel1b.get());
+    sp<SurfaceControl> childLevel2b =
+            createColorLayer("child level 2b", Color::BLACK, childLevel1b.get());
+
+    Transaction{}
+            .setRelativeLayer(childLevel1a, childLevel2b->getHandle(), 1)
+            .show(childLevel1a)
+            .show(childLevel1b)
+            .show(childLevel2a)
+            .show(childLevel2b)
+            .apply();
+
+    {
+        // The childLevel1a should be in front of childLevel2b.
+        ScreenCapture::captureScreen(&sc);
+        sc->expectColor(Rect(0, 0, mDisplayWidth, mDisplayHeight), testLayerColor);
+    }
+
+    // Background layer (RED)
+    // Foregroud layer (GREEN)
+    //   child level 1a (testLayerColor) (relative to child level 2b)
+    Transaction{}.reparent(childLevel1b, nullptr).apply();
+
+    // // Background layer (RED)
+    // // Foregroud layer (GREEN)
+    Transaction{}.reparent(childLevel1a, childLevel2a->getHandle()).apply();
+
+    {
+        // The childLevel1a and childLevel1b are no longer on screen
+        ScreenCapture::captureScreen(&sc);
+        sc->expectColor(Rect(0, 0, mDisplayWidth, mDisplayHeight), Color::GREEN);
+    }
+
+    // Background layer (RED)
+    // Foregroud layer (GREEN)
+    //   child level 1b (WHITE)
+    //     child level 2a (BLUE)
+    //       child level 1a (testLayerColor) (relative to child level 2b)
+    //     child level 2b (BLACK)
+    Transaction{}.reparent(childLevel1b, mForegroundLayer->getHandle()).apply();
+
+    {
+        // Nothing should change at this point since relative z info was preserved.
+        ScreenCapture::captureScreen(&sc);
+        sc->expectColor(Rect(0, 0, mDisplayWidth, mDisplayHeight), testLayerColor);
+    }
+}
 } // namespace android
 
 // TODO(b/129481165): remove the #pragma below and fix conversion issues