Respect relative layer value if in tree when screenshotting.

Relative layers need to be screenshotted so the dim layer can be
captured in the correct place. However, there are other relative layers
that are not in the tree that shouldn't be part of the screenshot. This
change allows relative layer value to be handled correctly as long as the
layer it's relative to is in the tree to capture. This ensures relative
layers in the tree are captured the correct order, but also ignores any
relative layers that are not in the tree.

Change-Id: If0e0ea1417f01ad1bfe518744ae2b9827009c4c6
Fixes: 77823704
Test: Open app with dialog and go to recents. Dim is screenshotted.
Test: Open app with IME and go to recents. IME is not screenshotted.
Test: Transaction_test#DontCaptureRelativeOutsideTree, CaptureRelativeInTree
diff --git a/services/surfaceflinger/Layer.cpp b/services/surfaceflinger/Layer.cpp
index 8c0050e..2300048 100644
--- a/services/surfaceflinger/Layer.cpp
+++ b/services/surfaceflinger/Layer.cpp
@@ -23,6 +23,7 @@
 #include <stdint.h>
 #include <stdlib.h>
 #include <sys/types.h>
+#include <algorithm>
 
 #include <cutils/compiler.h>
 #include <cutils/native_handle.h>
@@ -1787,29 +1788,81 @@
     }
 }
 
-/**
- * Traverse only children in z order, ignoring relative layers.
- */
-void Layer::traverseChildrenInZOrder(LayerVector::StateSet stateSet,
-                                     const LayerVector::Visitor& visitor) {
+LayerVector Layer::makeChildrenTraversalList(LayerVector::StateSet stateSet,
+                                             const std::vector<Layer*>& layersInTree) {
+    LOG_ALWAYS_FATAL_IF(stateSet == LayerVector::StateSet::Invalid,
+                        "makeTraversalList received invalid stateSet");
     const bool useDrawing = stateSet == LayerVector::StateSet::Drawing;
     const LayerVector& children = useDrawing ? mDrawingChildren : mCurrentChildren;
+    const State& state = useDrawing ? mDrawingState : mCurrentState;
+
+    LayerVector traverse;
+    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.
+        // If a relative layer is not a descendent, then it should be ignored.
+        if (std::binary_search(layersInTree.begin(), layersInTree.end(), strongRelative.get())) {
+            traverse.add(strongRelative);
+        }
+    }
+
+    for (const sp<Layer>& child : children) {
+        const State& childState = useDrawing ? child->mDrawingState : child->mCurrentState;
+        // If a layer has a relativeOf layer, only ignore if the layer it's relative to is a
+        // descendent of the top most parent of the tree. If it's not a descendent, then just add
+        // the child here since it won't be added later as a relative.
+        if (std::binary_search(layersInTree.begin(), layersInTree.end(),
+                               childState.zOrderRelativeOf.promote().get())) {
+            continue;
+        }
+        traverse.add(child);
+    }
+
+    return traverse;
+}
+
+void Layer::traverseChildrenInZOrderInner(const std::vector<Layer*>& layersInTree,
+                                          LayerVector::StateSet stateSet,
+                                          const LayerVector::Visitor& visitor) {
+    const LayerVector list = makeChildrenTraversalList(stateSet, layersInTree);
 
     size_t i = 0;
-    for (; i < children.size(); i++) {
-        const auto& relative = children[i];
+    for (; i < list.size(); i++) {
+        const auto& relative = list[i];
         if (relative->getZ() >= 0) {
             break;
         }
-        relative->traverseChildrenInZOrder(stateSet, visitor);
+        relative->traverseChildrenInZOrderInner(layersInTree, stateSet, visitor);
     }
+
     visitor(this);
-    for (; i < children.size(); i++) {
-        const auto& relative = children[i];
-        relative->traverseChildrenInZOrder(stateSet, visitor);
+    for (; i < list.size(); i++) {
+        const auto& relative = list[i];
+        relative->traverseChildrenInZOrderInner(layersInTree, stateSet, visitor);
     }
 }
 
+std::vector<Layer*> Layer::getLayersInTree(LayerVector::StateSet stateSet) {
+    const bool useDrawing = stateSet == LayerVector::StateSet::Drawing;
+    const LayerVector& children = useDrawing ? mDrawingChildren : mCurrentChildren;
+
+    std::vector<Layer*> layersInTree = {this};
+    for (size_t i = 0; i < children.size(); i++) {
+        const auto& child = children[i];
+        std::vector<Layer*> childLayers = child->getLayersInTree(stateSet);
+        layersInTree.insert(layersInTree.end(), childLayers.cbegin(), childLayers.cend());
+    }
+
+    return layersInTree;
+}
+
+void Layer::traverseChildrenInZOrder(LayerVector::StateSet stateSet,
+                                     const LayerVector::Visitor& visitor) {
+    std::vector<Layer*> layersInTree = getLayersInTree(stateSet);
+    std::sort(layersInTree.begin(), layersInTree.end());
+    traverseChildrenInZOrderInner(layersInTree, stateSet, visitor);
+}
+
 Transform Layer::getTransform() const {
     Transform t;
     const auto& p = mDrawingParent.promote();
diff --git a/services/surfaceflinger/Layer.h b/services/surfaceflinger/Layer.h
index d382a1a..be3967b 100644
--- a/services/surfaceflinger/Layer.h
+++ b/services/surfaceflinger/Layer.h
@@ -51,6 +51,7 @@
 #include "RenderEngine/Texture.h"
 
 #include <math/vec4.h>
+#include <vector>
 
 using namespace android::surfaceflinger;
 
@@ -564,6 +565,10 @@
                                  const LayerVector::Visitor& visitor);
     void traverseInZOrder(LayerVector::StateSet stateSet, const LayerVector::Visitor& visitor);
 
+    /**
+     * Traverse only children in z order, ignoring relative layers that are not children of the
+     * parent.
+     */
     void traverseChildrenInZOrder(LayerVector::StateSet stateSet,
                                   const LayerVector::Visitor& visitor);
 
@@ -778,6 +783,22 @@
     wp<Layer> mDrawingParent;
 
     mutable LayerBE mBE;
+
+private:
+    /**
+     * Returns an unsorted vector of all layers that are part of this tree.
+     * That includes the current layer and all its descendants.
+     */
+    std::vector<Layer*> getLayersInTree(LayerVector::StateSet stateSet);
+    /**
+     * Traverses layers that are part of this tree in the correct z order.
+     * layersInTree must be sorted before calling this method.
+     */
+    void traverseChildrenInZOrderInner(const std::vector<Layer*>& layersInTree,
+                                       LayerVector::StateSet stateSet,
+                                       const LayerVector::Visitor& visitor);
+    LayerVector makeChildrenTraversalList(LayerVector::StateSet stateSet,
+                                          const std::vector<Layer*>& layersInTree);
 };
 
 // ---------------------------------------------------------------------------
diff --git a/services/surfaceflinger/tests/Transaction_test.cpp b/services/surfaceflinger/tests/Transaction_test.cpp
index a0f12f1..5108279 100644
--- a/services/surfaceflinger/tests/Transaction_test.cpp
+++ b/services/surfaceflinger/tests/Transaction_test.cpp
@@ -2376,6 +2376,58 @@
     mCapture->expectColor(Rect(0, 10, 9, 19), {0, 0, 0, 0});
 }
 
+TEST_F(ScreenCaptureTest, DontCaptureRelativeOutsideTree) {
+    auto fgHandle = mFGSurfaceControl->getHandle();
+
+    sp<SurfaceControl> child =
+            mComposerClient->createSurface(String8("Child surface"), 10, 10, PIXEL_FORMAT_RGBA_8888,
+                                           0, mFGSurfaceControl.get());
+    sp<SurfaceControl> relative = mComposerClient->createSurface(String8("Relative surface"), 10,
+                                                                 10, PIXEL_FORMAT_RGBA_8888, 0);
+    fillSurfaceRGBA8(child, 200, 200, 200);
+    fillSurfaceRGBA8(relative, 100, 100, 100);
+
+    SurfaceComposerClient::Transaction()
+            .show(child)
+            // Set relative layer above fg layer so should be shown above when computing all layers.
+            .setRelativeLayer(relative, fgHandle, 1)
+            .show(relative)
+            .apply(true);
+
+    // Captures mFGSurfaceControl layer and its child. Relative layer shouldn't be captured.
+    ScreenCapture::captureLayers(&mCapture, fgHandle);
+    mCapture->expectFGColor(10, 10);
+    mCapture->expectChildColor(0, 0);
+}
+
+TEST_F(ScreenCaptureTest, CaptureRelativeInTree) {
+    auto fgHandle = mFGSurfaceControl->getHandle();
+
+    sp<SurfaceControl> child =
+            mComposerClient->createSurface(String8("Child surface"), 10, 10, PIXEL_FORMAT_RGBA_8888,
+                                           0, mFGSurfaceControl.get());
+    sp<SurfaceControl> relative =
+            mComposerClient->createSurface(String8("Relative surface"), 10, 10,
+                                           PIXEL_FORMAT_RGBA_8888, 0, mFGSurfaceControl.get());
+    fillSurfaceRGBA8(child, 200, 200, 200);
+    fillSurfaceRGBA8(relative, 100, 100, 100);
+
+    SurfaceComposerClient::Transaction()
+            .show(child)
+            // Set relative layer below fg layer but relative to child layer so it should be shown
+            // above child layer.
+            .setLayer(relative, -1)
+            .setRelativeLayer(relative, child->getHandle(), 1)
+            .show(relative)
+            .apply(true);
+
+    // Captures mFGSurfaceControl layer and its children. Relative layer is a child of fg so its
+    // relative value should be taken into account, placing it above child layer.
+    ScreenCapture::captureLayers(&mCapture, fgHandle);
+    mCapture->expectFGColor(10, 10);
+    // Relative layer is showing on top of child layer
+    mCapture->expectColor(Rect(0, 0, 9, 9), {100, 100, 100, 255});
+}
 
 // In the following tests we verify successful skipping of a parent layer,
 // so we use the same verification logic and only change how we mutate