From 7025f80b2689699f85e1d6d865691797e22499ea Mon Sep 17 00:00:00 2001
From: zhouke <zhouke35@huawei.com>
Date: Wed, 26 Oct 2022 20:18:36 +0800
Subject: [PATCH] =?UTF-8?q?dumpLayout=20=E4=BF=AE=E5=A4=8Dbug=EF=BC=8C?=
 =?UTF-8?q?=E5=90=8C=E4=B8=80=E5=B1=82=E7=BA=A7=E7=9A=84=E8=8A=82=E7=82=B9?=
 =?UTF-8?q?=E4=BF=A1=E6=81=AF=E8=8B=A5=E6=8E=92=E5=9C=A8=E5=89=8D=E9=9D=A2?=
 =?UTF-8?q?=E7=9A=84=E8=8A=82=E7=82=B9=E8=A2=AB=E8=BF=87=E6=BB=A4=EF=BC=8C?=
 =?UTF-8?q?=E5=90=8E=E9=9D=A2=E7=9A=84=E8=8A=82=E7=82=B9=E5=B0=86=E4=B8=8D?=
 =?UTF-8?q?=E4=BC=9A=E8=A2=AB=E6=89=93=E5=8D=B0.Signed-off-by:=20<zhouke35?=
 =?UTF-8?q?@huawei.com>.?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Signed-off-by: zhouke <zhouke35@huawei.com>
---
 uitest/core/ui_model.cpp | 38 +++++++++++++++++++++++++++++---------
 1 file changed, 29 insertions(+), 9 deletions(-)

diff --git a/uitest/core/ui_model.cpp b/uitest/core/ui_model.cpp
index 2588f28..3bb20a5 100644
--- a/uitest/core/ui_model.cpp
+++ b/uitest/core/ui_model.cpp
@@ -311,7 +311,8 @@ namespace OHOS::uitest {
         widgetsConstructed_ = true;
     }
 
-    static void DfsMarshalWidget(const WidgetTree& tree, const Widget& root, nlohmann::json& dom)
+    static void DfsMarshalWidget(const WidgetTree& tree, const Widget& root, nlohmann::json& dom,
+        const std::map<string, size_t>& widgetChildCountMap)
     {
         auto attributesData = json();
         // "< UiAttr::HIERARCHY" : do not expose inner used attributes
@@ -326,14 +327,21 @@ namespace OHOS::uitest {
         attributesData[ATTR_NAMES[UiAttr::BOUNDS].data()] = stream.str();
 
         auto childrenData = json::array();
-        uint32_t childIndex = 0;
-        auto child = tree.GetChildWidget(root, childIndex);
-        while (child != nullptr) {
-            auto childData = json();
-            DfsMarshalWidget(tree, *child, childData);
-            childrenData.emplace_back(childData);
+        uint32_t childIndex = 0, childCount = 0, visitCount = 0;
+        auto hierarchy = root.GetHierarchy();
+        if (widgetChildCountMap.find(hierarchy) != widgetChildCountMap.end()) {
+            childCount = widgetChildCountMap.find(hierarchy)->second;
+        }
+        while (visitCount < childCount) {
+            auto child = tree.GetChildWidget(root, childIndex);
             childIndex++;
-            child = tree.GetChildWidget(root, childIndex);
+            if (child == nullptr) {
+                continue;
+            }
+            auto childData = json();
+            DfsMarshalWidget(tree, *child, childData, widgetChildCountMap);
+            childrenData.emplace_back(childData);
+            visitCount++;
         }
 
         dom["attributes"] = attributesData;
@@ -344,8 +352,20 @@ namespace OHOS::uitest {
     {
         DCHECK(widgetsConstructed_);
         auto root = GetRootWidget();
+        std::map<string, size_t> widgetChildCountMap;
+        for (auto &hierarchy : widgetHierarchyIdDfsOrder_) {
+            if (hierarchy == ROOT_HIERARCHY) {
+                continue;
+            }
+            auto parentHierarchy = WidgetHierarchyBuilder::GetParentWidgetHierarchy(hierarchy);
+            if (widgetChildCountMap.find(parentHierarchy) == widgetChildCountMap.end()) {
+                widgetChildCountMap[parentHierarchy] = 1;
+            } else {
+                widgetChildCountMap[parentHierarchy] = widgetChildCountMap[parentHierarchy] + 1;
+            }
+        }
         if (root != nullptr) {
-            DfsMarshalWidget(*this, *root, dom);
+            DfsMarshalWidget(*this, *root, dom, widgetChildCountMap);
         }
     }