Node -> Measure func map instead of vec (#1581)

Summary:
Pull Request resolved: https://github.com/facebook/yoga/pull/1581

This is better than just trusting the order of the measure func call. Now each measure function I/O is associated with a node in the JSON.

Reviewed By: NickGerleman

Differential Revision: D53776790

fbshipit-source-id: 793cf2d9cbf6f663d24848af0af30aa297614eea

capture/CaptureTree.cpp

@@ -5,8 +5,8 @@
  * LICENSE file in the root directory of this source tree.
  */
 
+#include <cassert>
 #include <fstream>
-#include <vector>
 
 #include <capture/CaptureTree.h>
 #include <capture/NodeToString.h>
@@ -23,10 +23,9 @@ static void captureTree(
   file << serializedTree;
 }
 
-static std::vector<SerializedMeasureFunc>& currentSerializedMeasureFuncVec() {
-  static thread_local std::vector<SerializedMeasureFunc>
-      currentSerializedMeasureFuncVec;
-  return currentSerializedMeasureFuncVec;
+static SerializedMeasureFuncMap& currentSerializedMeasureFuncMap() {
+  static thread_local SerializedMeasureFuncMap map{};
+  return map;
 }
 
 /*
@@ -57,38 +56,45 @@ void YGNodeCalculateLayoutWithCapture(
     YGDirection ownerDirection,
     const std::filesystem::path& path) {
   dirtyTree(node);
+  YGNodeCalculateLayout(node, availableWidth, availableHeight, ownerDirection);
+
   json j;
   serializeLayoutInputs(j, availableWidth, availableHeight, ownerDirection);
   serializeTree(
       j,
+      currentSerializedMeasureFuncMap(),
       node,
       PrintOptions::Style | PrintOptions::Children | PrintOptions::Config |
           PrintOptions::Node);
-
-  YGNodeCalculateLayout(node, availableWidth, availableHeight, ownerDirection);
-
-  serializeMeasureFuncResults(j, currentSerializedMeasureFuncVec());
   // TODO: It is possible to have a measure function call layout again if, e.g.,
   // views are nested in text. Need to be able to resolve this special case.
-  currentSerializedMeasureFuncVec().clear();
+  currentSerializedMeasureFuncMap().clear();
   captureTree(j.dump(2), path);
 }
 
 void captureMeasureFunc(
+    YGNodeConstRef node,
     float width,
     YGMeasureMode widthMode,
     float height,
     YGMeasureMode heightMode,
     YGSize output,
     std::chrono::steady_clock::duration durationNs) {
-  currentSerializedMeasureFuncVec().push_back(SerializedMeasureFunc{
-      width,
-      widthMode,
-      height,
-      heightMode,
-      output.width,
-      output.height,
-      durationNs.count()});
+  auto measureFuncIt = currentSerializedMeasureFuncMap().find(node);
+  if (measureFuncIt == currentSerializedMeasureFuncMap().end()) {
+    std::vector<SerializedMeasureFunc> vec{};
+    currentSerializedMeasureFuncMap().insert(std::make_pair(node, vec));
+  }
+  measureFuncIt = currentSerializedMeasureFuncMap().find(node);
+  assert(measureFuncIt != currentSerializedMeasureFuncMap().end());
+  measureFuncIt->second.push_back(
+      {width,
+       widthMode,
+       height,
+       heightMode,
+       output.width,
+       output.height,
+       durationNs.count()});
 }
 
 } // namespace facebook::yoga

capture/CaptureTree.h

@@ -8,6 +8,8 @@
 #pragma once
 
 #include <filesystem>
+#include <unordered_map>
+#include <vector>
 
 #include <yoga/Yoga.h>
 
@@ -23,6 +25,9 @@ struct SerializedMeasureFunc {
   std::chrono::steady_clock::duration::rep durationNs;
 };
 
+using SerializedMeasureFuncMap =
+    std::unordered_map<YGNodeConstRef, std::vector<SerializedMeasureFunc>>;
+
 void YGNodeCalculateLayoutWithCapture(
     YGNodeRef node,
     float availableWidth,
@@ -31,6 +36,7 @@ void YGNodeCalculateLayoutWithCapture(
     const std::filesystem::path& path);
 
 void captureMeasureFunc(
+    YGNodeConstRef node,
     float width,
     YGMeasureMode widthMode,
     float height,

capture/NodeToString.cpp

@@ -6,6 +6,7 @@
  */
 
 #include <memory>
+#include <vector>
 
 #include <capture/NodeToString.h>
 
@@ -120,7 +121,26 @@ static YGValue borderFloatToYGValue(YGNodeRef node, YGEdge edge) {
   return YGValue{val, unit};
 }
 
-static void serializeTreeImpl(json& j, YGNodeRef node, PrintOptions options) {
+static void serializeMeasureFuncResults(
+    json& j,
+    std::vector<SerializedMeasureFunc>& measureFuncs) {
+  for (auto measureFunc : measureFuncs) {
+    j["measure-funcs"].push_back(
+        {{"width", measureFunc.inputWidth},
+         {"width-mode", YGMeasureModeToString(measureFunc.widthMode)},
+         {"height", measureFunc.inputHeight},
+         {"height-mode", YGMeasureModeToString(measureFunc.heightMode)},
+         {"output-width", measureFunc.outputWidth},
+         {"output-height", measureFunc.outputHeight},
+         {"duration-ns", measureFunc.durationNs}});
+  }
+}
+
+static void serializeTreeImpl(
+    json& j,
+    SerializedMeasureFuncMap& nodesToMeasureFuncs,
+    YGNodeRef node,
+    PrintOptions options) {
   if ((options & PrintOptions::Layout) == PrintOptions::Layout) {
     j["layout"]["width"] = YGNodeStyleGetWidth(node).value;
     j["layout"]["height"] = YGNodeStyleGetHeight(node).value;
@@ -293,7 +313,12 @@ static void serializeTreeImpl(json& j, YGNodeRef node, PrintOptions options) {
         YGNodeGetAlwaysFormsContainingBlock(node),
         YGNodeGetAlwaysFormsContainingBlock(defaultNode.get()));
     if (YGNodeHasMeasureFunc(node)) {
-      j["node"]["has-custom-measure"] = true;
+      auto measureFuncIt = nodesToMeasureFuncs.find(node);
+      if (measureFuncIt == nodesToMeasureFuncs.end()) {
+        j["node"]["measure-funcs"];
+      } else {
+        serializeMeasureFuncResults(j["node"], measureFuncIt->second);
+      }
     }
   }
 
@@ -302,13 +327,21 @@ static void serializeTreeImpl(json& j, YGNodeRef node, PrintOptions options) {
       childCount > 0) {
     for (size_t i = 0; i < childCount; i++) {
       j["children"].push_back({});
-      serializeTreeImpl(j["children"][i], YGNodeGetChild(node, i), options);
+      serializeTreeImpl(
+          j["children"][i],
+          nodesToMeasureFuncs,
+          YGNodeGetChild(node, i),
+          options);
     }
   }
 }
 
-void serializeTree(json& j, YGNodeRef node, PrintOptions options) {
-  serializeTreeImpl(j["tree"], node, options);
+void serializeTree(
+    json& j,
+    SerializedMeasureFuncMap& nodesToMeasureFuncs,
+    YGNodeRef node,
+    PrintOptions options) {
+  serializeTreeImpl(j["tree"], nodesToMeasureFuncs, node, options);
 }
 
 void serializeLayoutInputs(
@@ -323,19 +356,4 @@ void serializeLayoutInputs(
   };
 }
 
-void serializeMeasureFuncResults(
-    json& j,
-    std::vector<SerializedMeasureFunc>& measureFuncs) {
-  for (auto measureFunc : measureFuncs) {
-    j["measure-funcs"].push_back(
-        {{"width", measureFunc.inputWidth},
-         {"width-mode", YGMeasureModeToString(measureFunc.widthMode)},
-         {"height", measureFunc.inputHeight},
-         {"height-mode", YGMeasureModeToString(measureFunc.heightMode)},
-         {"output-width", measureFunc.outputWidth},
-         {"output-height", measureFunc.outputHeight},
-         {"duration-ns", measureFunc.durationNs}});
-  }
-}
-
 } // namespace facebook::yoga

capture/NodeToString.h

@@ -25,7 +25,11 @@ enum class PrintOptions : uint8_t {
 };
 YG_DEFINE_ENUM_FLAG_OPERATORS(PrintOptions);
 
-void serializeTree(nlohmann::json& j, YGNodeRef root, PrintOptions options);
+void serializeTree(
+    nlohmann::json& j,
+    SerializedMeasureFuncMap& nodesToMeasureFuncs,
+    YGNodeRef root,
+    PrintOptions options);
 
 void serializeLayoutInputs(
     nlohmann::json& j,
@@ -33,8 +37,4 @@ void serializeLayoutInputs(
     float availableHeight,
     YGDirection ownerDirection);
 
-void serializeMeasureFuncResults(
-    nlohmann::json& j,
-    std::vector<SerializedMeasureFunc>& measureFuncs);
-
 } // namespace facebook::yoga