Support for (de)serializing measure funcs

Summary:
In addition to all the state that gets set on the node that is easy to serialize - like floats, enums, bools, etc - we also need to serialize measure functions. This is because these functions take a nontrivial amount of time up during layout and we should capture that. Also, they are important to the ability to truly replay layout as it was captured as the results of the measure functions determine many of the steps the layout algorithm takes.

Capturing this is rather tricky however, but I think I found a solution that is relatively simple and non-error prone. Essentially, since we are capturing the entire tree and virtually every input that goes into the flexbox algorithm, we *should* be able to replay layout exactly as it was captured. This means that the order in which measure functions are called *should* be the same. If this is the case, then all we need to do to capture the measure functions is store their input, output, and duration in a big array. During deserialization we just keep track of an index and use that to determine which measure function we should call. That is the premise behind what happens in this diff. In theory the algorithm could change and the capture would be wrong but it is easy enough to recapture again. Additionally we need to dirty the tree so that we get rid of caching which might omit some measure func calls

In order to capture you need to insert a method exposed by CaptureTree.h into the client measure func, which is kind of annoying but not that bad. In future diffs I will put a macro in place to make this even easier.

I also add our first capture! Which is of a large react native desktop app

Reviewed By: NickGerleman

Differential Revision: D53581121

fbshipit-source-id: 876a230208d67f0ecf76844a4f1b80048353aae2

benchmark/Benchmark.cpp

@@ -5,11 +5,15 @@
  * LICENSE file in the root directory of this source tree.
  */
 
+#include <chrono>
 #include <filesystem>
 #include <fstream>
+#include <iostream>
+#include <thread>
 
 #include <benchmark/Benchmark.h>
 #include <benchmark/TreeDeserialization.h>
+#include <capture/CaptureTree.h>
 #include <nlohmann/json.hpp>
 
 namespace facebook::yoga {
@@ -17,10 +21,35 @@ namespace facebook::yoga {
 using namespace nlohmann;
 using namespace std::chrono;
 
-constexpr uint32_t kNumRepititions = 1000;
+constexpr uint32_t kNumRepititions = 100;
 using SteadyClockDurations =
     std::array<steady_clock::duration, kNumRepititions>;
 
+YGSize mockMeasureFunc(
+    YGNodeConstRef node,
+    float availableWidth,
+    YGMeasureMode widthMode,
+    float availableHeight,
+    YGMeasureMode heightMode) {
+  (void)node;
+  (void)availableHeight;
+  (void)availableWidth;
+  (void)widthMode;
+  (void)heightMode;
+  MeasureFuncVecWithIndex* fns =
+      static_cast<MeasureFuncVecWithIndex*>(YGNodeGetContext(node));
+  if (fns->index >= fns->vec.size()) {
+    return {10.0, 10.0};
+  }
+
+  auto values = fns->vec.at(fns->index);
+  fns->index++;
+
+  std::this_thread::sleep_for(std::chrono::nanoseconds(values.durationNs));
+
+  return {values.outputWidth, values.outputHeight};
+}
+
 std::shared_ptr<const YGConfig> buildConfigFromJson(const json& j) {
   json jsonConfig = j["config"];
   std::shared_ptr<YGConfig> config(YGConfigNew(), YGConfigFree);
@@ -186,24 +215,37 @@ void setStylesFromJson(const json& j, YGNodeRef node) {
 
 std::shared_ptr<YGNode> buildNodeFromJson(
     const json& j,
-    std::shared_ptr<const YGConfig> config) {
+    std::shared_ptr<const YGConfig> config,
+    std::shared_ptr<MeasureFuncVecWithIndex> fns) {
   std::shared_ptr<YGNode> node(YGNodeNewWithConfig(config.get()), YGNodeFree);
-  json nodeState = j["node"];
 
+  if (!j.contains("node") || j["node"].is_null()) {
+    return node;
+  }
+
+  json nodeState = j["node"];
   for (json::iterator it = nodeState.begin(); it != nodeState.end(); it++) {
     if (it.key() == "always-forms-containing-block") {
       YGNodeSetAlwaysFormsContainingBlock(node.get(), it.value());
+    } else if (it.key() == "has-custom-measure" && it.value()) {
+      YGNodeSetContext(node.get(), fns.get());
+      YGNodeSetMeasureFunc(node.get(), mockMeasureFunc);
     }
   }
 
   return node;
 }
 
-YogaNodeAndConfig
-buildTreeFromJson(const json& j, YogaNodeAndConfig* parent, size_t index) {
-  std::shared_ptr<const YGConfig> config = buildConfigFromJson(j);
-  std::shared_ptr<YGNode> node = buildNodeFromJson(j, config);
-  YogaNodeAndConfig wrapper{node, config, std::vector<YogaNodeAndConfig>{}};
+std::shared_ptr<YogaNodeAndConfig> buildTreeFromJson(
+    const json& j,
+    std::shared_ptr<MeasureFuncVecWithIndex> fns,
+    std::shared_ptr<YogaNodeAndConfig> parent,
+    size_t index) {
+  auto config = buildConfigFromJson(j);
+  auto node = buildNodeFromJson(j, config, fns);
+  auto wrapper = std::make_shared<YogaNodeAndConfig>(
+      node, config, std::vector<std::shared_ptr<YogaNodeAndConfig>>{});
+
   if (parent != nullptr) {
     YGNodeInsertChild(parent->node_.get(), node.get(), index);
     parent->children_.push_back(wrapper);
@@ -211,22 +253,25 @@ buildTreeFromJson(const json& j, YogaNodeAndConfig* parent, size_t index) {
 
   setStylesFromJson(j, node.get());
 
-  json children = j["children"];
-  size_t childIndex = 0;
-  for (json child : children) {
-    buildTreeFromJson(child, &wrapper, childIndex);
-    childIndex++;
+  if (j.contains("children")) {
+    json children = j["children"];
+    size_t childIndex = 0;
+    for (json child : children) {
+      buildTreeFromJson(child, fns, wrapper, childIndex);
+      childIndex++;
+    }
   }
 
   return wrapper;
 }
 
-BenchmarkResult generateBenchmark(const std::filesystem::path& capturePath) {
-  std::ifstream captureFile(capturePath);
-  json capture = json::parse(captureFile);
+BenchmarkResult generateBenchmark(json& capture) {
+  auto fns = std::make_shared<MeasureFuncVecWithIndex>();
+  populateMeasureFuncVec(capture["measure-funcs"], fns);
 
   auto treeCreationBegin = steady_clock::now();
-  YogaNodeAndConfig root = buildTreeFromJson(capture, nullptr, 0 /*index*/);
+  std::shared_ptr<YogaNodeAndConfig> root =
+      buildTreeFromJson(capture["tree"], fns, nullptr, 0 /*index*/);
   auto treeCreationEnd = steady_clock::now();
 
   json layoutInputs = capture["layout-inputs"];
@@ -236,7 +281,7 @@ BenchmarkResult generateBenchmark(const std::filesystem::path& capturePath) {
 
   auto layoutBegin = steady_clock::now();
   YGNodeCalculateLayout(
-      root.node_.get(), availableWidth, availableHeight, direction);
+      root->node_.get(), availableWidth, availableHeight, direction);
   auto layoutEnd = steady_clock::now();
 
   return BenchmarkResult{
@@ -278,8 +323,11 @@ void benchmark(std::filesystem::path& capturesDir) {
     SteadyClockDurations layoutDurations;
     SteadyClockDurations totalDurations;
 
+    std::ifstream captureFile(capture.path());
+    json j = json::parse(captureFile);
+
     for (uint32_t i = 0; i < kNumRepititions; i++) {
-      BenchmarkResult result = generateBenchmark(capture.path());
+      BenchmarkResult result = generateBenchmark(j);
       treeCreationDurations[i] = result.treeCreationDuration;
       layoutDurations[i] = result.layoutDuration;
       totalDurations[i] = result.treeCreationDuration + result.layoutDuration;

benchmark/Benchmark.h

@@ -17,9 +17,15 @@
 namespace facebook::yoga {
 
 struct YogaNodeAndConfig {
+  YogaNodeAndConfig(
+      std::shared_ptr<YGNode> node,
+      std::shared_ptr<const YGConfig> config,
+      std::vector<std::shared_ptr<YogaNodeAndConfig>> children)
+      : node_(node), config_(config), children_(children) {}
+
   std::shared_ptr<YGNode> node_;
   std::shared_ptr<const YGConfig> config_;
-  std::vector<YogaNodeAndConfig> children_;
+  std::vector<std::shared_ptr<YogaNodeAndConfig>> children_;
 };
 
 struct BenchmarkResult {

benchmark/TreeDeserialization.cpp

@@ -0,0 +1,245 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#include <memory>
+#include <vector>
+
+#include <benchmark/TreeDeserialization.h>
+#include <capture/CaptureTree.h>
+#include <nlohmann/json.hpp>
+#include <yoga/Yoga.h>
+
+namespace facebook::yoga {
+
+using namespace nlohmann;
+
+static inline bool isAuto(json& j) {
+  return j.is_string() && j == "auto";
+}
+
+static inline std::string invalidArgumentMessage(
+    std::string arg,
+    std::string enumName) {
+  return arg + " does not represent any " + enumName + " values";
+}
+
+static inline float floatFromJson(json& j) {
+  float result = YGUndefined;
+  if (!j.is_null()) {
+    result = j;
+  }
+
+  return result;
+}
+
+YGFlexDirection flexDirectionFromString(std::string str) {
+  if (str == "row") {
+    return YGFlexDirectionRow;
+  } else if (str == "row-reverse") {
+    return YGFlexDirectionRowReverse;
+  } else if (str == "column") {
+    return YGFlexDirectionColumn;
+  } else if (str == "column-reverse") {
+    return YGFlexDirectionColumnReverse;
+  } else {
+    throw std::invalid_argument(invalidArgumentMessage(str, "YGFlexDirection"));
+  }
+}
+
+YGJustify justifyContentFromString(std::string str) {
+  if (str == "flex-start") {
+    return YGJustifyFlexStart;
+  } else if (str == "center") {
+    return YGJustifyCenter;
+  } else if (str == "flex-end") {
+    return YGJustifyFlexEnd;
+  } else if (str == "space-between") {
+    return YGJustifySpaceBetween;
+  } else if (str == "space-around") {
+    return YGJustifySpaceAround;
+  } else if (str == "space-evenly") {
+    return YGJustifySpaceEvenly;
+  } else {
+    throw std::invalid_argument(invalidArgumentMessage(str, "YGJustify"));
+  }
+}
+
+YGAlign alignFromString(std::string str) {
+  if (str == "auto") {
+    return YGAlignAuto;
+  } else if (str == "flex-start") {
+    return YGAlignFlexStart;
+  } else if (str == "center") {
+    return YGAlignCenter;
+  } else if (str == "flex-end") {
+    return YGAlignFlexEnd;
+  } else if (str == "stretch") {
+    return YGAlignStretch;
+  } else if (str == "baseline") {
+    return YGAlignBaseline;
+  } else if (str == "space-between") {
+    return YGAlignSpaceBetween;
+  } else if (str == "space-around") {
+    return YGAlignSpaceAround;
+  } else if (str == "space-evenly") {
+    return YGAlignSpaceEvenly;
+  } else {
+    throw std::invalid_argument(invalidArgumentMessage(str, "YGAlign"));
+  }
+}
+
+YGWrap wrapFromString(std::string str) {
+  if (str == "no-wrap") {
+    return YGWrapNoWrap;
+  } else if (str == "wrap") {
+    return YGWrapWrap;
+  } else if (str == "wrap-reverse") {
+    return YGWrapWrapReverse;
+  } else {
+    throw std::invalid_argument(invalidArgumentMessage(str, "YGAlign"));
+  }
+}
+
+YGOverflow overflowFromString(std::string str) {
+  if (str == "visible") {
+    return YGOverflowVisible;
+  } else if (str == "hidden") {
+    return YGOverflowHidden;
+  } else if (str == "scroll") {
+    return YGOverflowScroll;
+  } else {
+    throw std::invalid_argument(invalidArgumentMessage(str, "YGAlign"));
+  }
+}
+
+YGDisplay displayFromString(std::string str) {
+  if (str == "flex") {
+    return YGDisplayFlex;
+  } else if (str == "none") {
+    return YGDisplayNone;
+  } else {
+    throw std::invalid_argument(invalidArgumentMessage(str, "YGAlign"));
+  }
+}
+
+YGPositionType positionTypeFromString(std::string str) {
+  if (str == "static") {
+    return YGPositionTypeStatic;
+  } else if (str == "relative") {
+    return YGPositionTypeRelative;
+  } else if (str == "absolute") {
+    return YGPositionTypeAbsolute;
+  } else {
+    throw std::invalid_argument(invalidArgumentMessage(str, "YGAlign"));
+  }
+}
+
+YGUnit unitFromJson(json& j) {
+  if (isAuto(j)) {
+    return YGUnitAuto;
+  }
+
+  std::string unit = j["unit"];
+  if (unit == "px") {
+    return YGUnitPoint;
+  } else if (unit == "pct") {
+    return YGUnitPercent;
+  } else {
+    throw std::invalid_argument(invalidArgumentMessage(unit, "YGUnit"));
+  }
+}
+
+YGEdge edgeFromString(std::string str) {
+  if (str == "left") {
+    return YGEdgeLeft;
+  } else if (str == "top") {
+    return YGEdgeTop;
+  } else if (str == "right") {
+    return YGEdgeRight;
+  } else if (str == "bottom") {
+    return YGEdgeBottom;
+  } else if (str == "start") {
+    return YGEdgeStart;
+  } else if (str == "end") {
+    return YGEdgeEnd;
+  } else if (str == "horizontal") {
+    return YGEdgeHorizontal;
+  } else if (str == "vertical") {
+    return YGEdgeVertical;
+  } else if (str == "all") {
+    return YGEdgeAll;
+  } else {
+    throw std::invalid_argument(invalidArgumentMessage(str, "YGEdge"));
+  }
+}
+
+YGErrata errataFromString(std::string str) {
+  if (str == "none") {
+    return YGErrataNone;
+  } else if (str == "all") {
+    return YGErrataAll;
+  } else if (str == "classic") {
+    return YGErrataClassic;
+  } else {
+    throw std::invalid_argument(invalidArgumentMessage(str, "YGErrata"));
+  }
+}
+
+YGExperimentalFeature experimentalFeatureFromString(std::string str) {
+  if (str == "web-flex-basis") {
+    return YGExperimentalFeatureWebFlexBasis;
+  } else {
+    throw std::invalid_argument(
+        invalidArgumentMessage(str, "YGExperimentalFeature"));
+  }
+}
+
+std::string edgeStringFromPropertyName(
+    json::iterator it,
+    std::string propertyName) {
+  return it.key().substr(propertyName.length() + 1);
+}
+
+YGDirection directionFromString(std::string str) {
+  if (str == "ltr") {
+    return YGDirectionLTR;
+  } else if (str == "rtl") {
+    return YGDirectionRTL;
+  } else if (str == "inherit") {
+    return YGDirectionInherit;
+  } else {
+    throw std::invalid_argument(invalidArgumentMessage(str, "YGDirection"));
+  }
+}
+
+YGMeasureMode measureModeFromString(std::string str) {
+  if (str == "at-most") {
+    return YGMeasureModeAtMost;
+  } else if (str == "exactly") {
+    return YGMeasureModeExactly;
+  } else if (str == "undefined") {
+    return YGMeasureModeUndefined;
+  } else {
+    throw std::invalid_argument(invalidArgumentMessage(str, "YGMeasureMode"));
+  }
+}
+
+void populateMeasureFuncVec(
+    json& j,
+    std::shared_ptr<MeasureFuncVecWithIndex> fns) {
+  for (auto measureFuncJson : j) {
+    fns->vec.push_back(SerializedMeasureFunc{
+        floatFromJson(measureFuncJson["width"]),
+        measureModeFromString(measureFuncJson["width-mode"]),
+        floatFromJson(measureFuncJson["height"]),
+        measureModeFromString(measureFuncJson["height-mode"]),
+        floatFromJson(measureFuncJson["output-width"]),
+        floatFromJson(measureFuncJson["output-height"]),
+        measureFuncJson["duration-ns"]});
+  }
+}
+} // namespace facebook::yoga

benchmark/TreeDeserialization.h

@@ -7,6 +7,10 @@
 
 #pragma once
 
+#include <memory>
+#include <vector>
+
+#include <capture/CaptureTree.h>
 #include <nlohmann/json.hpp>
 #include <yoga/Yoga.h>
 
@@ -14,193 +18,42 @@ namespace facebook::yoga {
 
 using namespace nlohmann;
 
-inline std::string invalidArgumentMessage(
-    std::string arg,
-    std::string enumName) {
-  return arg + " does not represent any " + enumName + " values";
-}
+struct MeasureFuncVecWithIndex {
+  std::vector<SerializedMeasureFunc> vec;
+  size_t index;
+};
 
-inline YGFlexDirection flexDirectionFromString(std::string str) {
-  if (str == "row") {
-    return YGFlexDirectionRow;
-  } else if (str == "row-reverse") {
-    return YGFlexDirectionRowReverse;
-  } else if (str == "column") {
-    return YGFlexDirectionColumn;
-  } else if (str == "column-reverse") {
-    return YGFlexDirectionColumnReverse;
-  } else {
-    throw std::invalid_argument(invalidArgumentMessage(str, "YGFlexDirection"));
-  }
-}
+YGFlexDirection flexDirectionFromString(std::string str);
 
-inline YGJustify justifyContentFromString(std::string str) {
-  if (str == "flex-start") {
-    return YGJustifyFlexStart;
-  } else if (str == "center") {
-    return YGJustifyCenter;
-  } else if (str == "flex-end") {
-    return YGJustifyFlexEnd;
-  } else if (str == "space-between") {
-    return YGJustifySpaceBetween;
-  } else if (str == "space-around") {
-    return YGJustifySpaceAround;
-  } else if (str == "space-evenly") {
-    return YGJustifySpaceEvenly;
-  } else {
-    throw std::invalid_argument(invalidArgumentMessage(str, "YGJustify"));
-  }
-}
+YGJustify justifyContentFromString(std::string str);
 
-inline YGAlign alignFromString(std::string str) {
-  if (str == "auto") {
-    return YGAlignAuto;
-  } else if (str == "flex-start") {
-    return YGAlignFlexStart;
-  } else if (str == "center") {
-    return YGAlignCenter;
-  } else if (str == "flex-end") {
-    return YGAlignFlexEnd;
-  } else if (str == "stretch") {
-    return YGAlignStretch;
-  } else if (str == "baseline") {
-    return YGAlignBaseline;
-  } else if (str == "space-between") {
-    return YGAlignSpaceBetween;
-  } else if (str == "space-around") {
-    return YGAlignSpaceAround;
-  } else if (str == "space-evenly") {
-    return YGAlignSpaceEvenly;
-  } else {
-    throw std::invalid_argument(invalidArgumentMessage(str, "YGAlign"));
-  }
-}
+YGAlign alignFromString(std::string str);
 
-inline YGWrap wrapFromString(std::string str) {
-  if (str == "no-wrap") {
-    return YGWrapNoWrap;
-  } else if (str == "wrap") {
-    return YGWrapWrap;
-  } else if (str == "wrap-reverse") {
-    return YGWrapWrapReverse;
-  } else {
-    throw std::invalid_argument(invalidArgumentMessage(str, "YGAlign"));
-  }
-}
+YGWrap wrapFromString(std::string str);
 
-inline YGOverflow overflowFromString(std::string str) {
-  if (str == "visible") {
-    return YGOverflowVisible;
-  } else if (str == "hidden") {
-    return YGOverflowHidden;
-  } else if (str == "scroll") {
-    return YGOverflowScroll;
-  } else {
-    throw std::invalid_argument(invalidArgumentMessage(str, "YGAlign"));
-  }
-}
+YGOverflow overflowFromString(std::string str);
 
-inline YGDisplay displayFromString(std::string str) {
-  if (str == "flex") {
-    return YGDisplayFlex;
-  } else if (str == "none") {
-    return YGDisplayNone;
-  } else {
-    throw std::invalid_argument(invalidArgumentMessage(str, "YGAlign"));
-  }
-}
+YGDisplay displayFromString(std::string str);
 
-inline YGPositionType positionTypeFromString(std::string str) {
-  if (str == "static") {
-    return YGPositionTypeStatic;
-  } else if (str == "relative") {
-    return YGPositionTypeRelative;
-  } else if (str == "absolute") {
-    return YGPositionTypeAbsolute;
-  } else {
-    throw std::invalid_argument(invalidArgumentMessage(str, "YGAlign"));
-  }
-}
+YGPositionType positionTypeFromString(std::string str);
 
-inline bool isAuto(json& j) {
-  return j.is_string() && j == "auto";
-}
+YGUnit unitFromJson(json& j);
 
-inline YGUnit unitFromJson(json& j) {
-  if (isAuto(j)) {
-    return YGUnitAuto;
-  }
+YGEdge edgeFromString(std::string str);
 
-  std::string unit = j["unit"];
-  if (unit == "px") {
-    return YGUnitPoint;
-  } else if (unit == "pct") {
-    return YGUnitPercent;
-  } else {
-    throw std::invalid_argument(invalidArgumentMessage(unit, "YGUnit"));
-  }
-}
+YGErrata errataFromString(std::string str);
 
-inline YGEdge edgeFromString(std::string str) {
-  if (str == "left") {
-    return YGEdgeLeft;
-  } else if (str == "top") {
-    return YGEdgeTop;
-  } else if (str == "right") {
-    return YGEdgeRight;
-  } else if (str == "bottom") {
-    return YGEdgeBottom;
-  } else if (str == "start") {
-    return YGEdgeStart;
-  } else if (str == "end") {
-    return YGEdgeEnd;
-  } else if (str == "horizontal") {
-    return YGEdgeHorizontal;
-  } else if (str == "vertical") {
-    return YGEdgeVertical;
-  } else if (str == "all") {
-    return YGEdgeAll;
-  } else {
-    throw std::invalid_argument(invalidArgumentMessage(str, "YGEdge"));
-  }
-}
+YGExperimentalFeature experimentalFeatureFromString(std::string str);
 
-inline YGErrata errataFromString(std::string str) {
-  if (str == "none") {
-    return YGErrataNone;
-  } else if (str == "all") {
-    return YGErrataAll;
-  } else if (str == "classic") {
-    return YGErrataClassic;
-  } else {
-    throw std::invalid_argument(invalidArgumentMessage(str, "YGErrata"));
-  }
-}
-
-inline YGExperimentalFeature experimentalFeatureFromString(std::string str) {
-  if (str == "web-flex-basis") {
-    return YGExperimentalFeatureWebFlexBasis;
-  } else {
-    throw std::invalid_argument(
-        invalidArgumentMessage(str, "YGExperimentalFeature"));
-  }
-}
-
-inline std::string edgeStringFromPropertyName(
+std::string edgeStringFromPropertyName(
     json::iterator it,
-    std::string propertyName) {
-  return it.key().substr(propertyName.length() + 1);
-}
+    std::string propertyName);
 
-inline YGDirection directionFromString(std::string str) {
-  if (str == "ltr") {
-    return YGDirectionLTR;
-  } else if (str == "rtl") {
-    return YGDirectionRTL;
-  } else if (str == "inherit") {
-    return YGDirectionInherit;
-  } else {
-    throw std::invalid_argument(invalidArgumentMessage(str, "YGDirection"));
-  }
-}
+YGDirection directionFromString(std::string str);
+
+YGMeasureMode measureModeFromString(std::string str);
+
+void populateMeasureFuncVec(
+    json& j,
+    std::shared_ptr<MeasureFuncVecWithIndex> fns);
 } // namespace facebook::yoga

capture/CaptureTree.cpp

@@ -6,6 +6,7 @@
  */
 
 #include <fstream>
+#include <vector>
 
 #include <capture/CaptureTree.h>
 #include <capture/NodeToString.h>
@@ -22,12 +23,40 @@ static void captureTree(
   file << serializedTree;
 }
 
+static std::vector<SerializedMeasureFunc>& currentSerializedMeasureFuncVec() {
+  static thread_local std::vector<SerializedMeasureFunc>
+      currentSerializedMeasureFuncVec;
+  return currentSerializedMeasureFuncVec;
+}
+
+/*
+ * Capturing a tree often means that we capturing multiple serial layouts over
+ * the course of the capture. Because of this, we need to make sure that we do
+ * a full layout pass with no caching. If we do not do this there is a chance
+ * we do not capture measure functions that were called and cached in previous
+ * layouts. Missing these captures would lead to inaccurate benchmarking where
+ * we do not have cached state.
+ *
+ * TODO: Dirty entire tree not just measure function nodes
+ */
+static void dirtyTree(YGNodeRef node) {
+  if (YGNodeHasMeasureFunc(node)) {
+    YGNodeMarkDirty(node);
+  }
+
+  const size_t childCount = YGNodeGetChildCount(node);
+  for (size_t i = 0; i < childCount; i++) {
+    dirtyTree(YGNodeGetChild(node, i));
+  }
+}
+
 void YGNodeCalculateLayoutWithCapture(
     YGNodeRef node,
     float availableWidth,
     float availableHeight,
     YGDirection ownerDirection,
     const std::filesystem::path& path) {
+  dirtyTree(node);
   json j;
   serializeLayoutInputs(j, availableWidth, availableHeight, ownerDirection);
   serializeTree(
@@ -35,8 +64,31 @@ void YGNodeCalculateLayoutWithCapture(
       node,
       PrintOptions::Style | PrintOptions::Children | PrintOptions::Config |
           PrintOptions::Node);
-  captureTree(j.dump(2), path);
+
   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();
+  captureTree(j.dump(2), path);
+}
+
+void captureMeasureFunc(
+    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()});
 }
 
 } // namespace facebook::yoga

capture/CaptureTree.h

@@ -13,6 +13,16 @@
 
 namespace facebook::yoga {
 
+struct SerializedMeasureFunc {
+  float inputWidth{0.0f};
+  YGMeasureMode widthMode{YGMeasureModeUndefined};
+  float inputHeight{0.0};
+  YGMeasureMode heightMode{YGMeasureModeUndefined};
+  float outputWidth{0.0f};
+  float outputHeight{0.0f};
+  std::chrono::steady_clock::duration::rep durationNs;
+};
+
 void YGNodeCalculateLayoutWithCapture(
     YGNodeRef node,
     float availableWidth,
@@ -20,4 +30,12 @@ void YGNodeCalculateLayoutWithCapture(
     YGDirection ownerDirection,
     const std::filesystem::path& path);
 
+void captureMeasureFunc(
+    float width,
+    YGMeasureMode widthMode,
+    float height,
+    YGMeasureMode heightMode,
+    YGSize output,
+    std::chrono::steady_clock::duration durationNs);
+
 } // namespace facebook::yoga

capture/NodeToString.cpp

@@ -114,13 +114,13 @@ static void appendEdges(
       (*Field)(defaultNode, YGEdgeHorizontal));
 }
 
-YGValue borderFloatToYGValue(YGNodeRef node, YGEdge edge) {
+static YGValue borderFloatToYGValue(YGNodeRef node, YGEdge edge) {
   float val = YGNodeStyleGetBorder(node, edge);
   YGUnit unit = YGFloatIsUndefined(val) ? YGUnitUndefined : YGUnitPoint;
   return YGValue{val, unit};
 }
 
-void serializeTree(json& j, YGNodeRef node, PrintOptions options) {
+static void serializeTreeImpl(json& j, YGNodeRef node, PrintOptions options) {
   if ((options & PrintOptions::Layout) == PrintOptions::Layout) {
     j["layout"]["width"] = YGNodeStyleGetWidth(node).value;
     j["layout"]["height"] = YGNodeStyleGetHeight(node).value;
@@ -302,11 +302,15 @@ void serializeTree(json& j, YGNodeRef node, PrintOptions options) {
       childCount > 0) {
     for (size_t i = 0; i < childCount; i++) {
       j["children"].push_back({});
-      serializeTree(j["children"][i], YGNodeGetChild(node, i), options);
+      serializeTreeImpl(j["children"][i], YGNodeGetChild(node, i), options);
     }
   }
 }
 
+void serializeTree(json& j, YGNodeRef node, PrintOptions options) {
+  serializeTreeImpl(j["tree"], node, options);
+}
+
 void serializeLayoutInputs(
     json& j,
     float availableWidth,
@@ -319,4 +323,19 @@ 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

@@ -7,8 +7,10 @@
 
 #pragma once
 
+#include <chrono>
 #include <string>
 
+#include <capture/CaptureTree.h>
 #include <nlohmann/json.hpp>
 #include <yoga/Yoga.h>
 
@@ -31,4 +33,8 @@ void serializeLayoutInputs(
     float availableHeight,
     YGDirection ownerDirection);
 
+void serializeMeasureFuncResults(
+    nlohmann::json& j,
+    std::vector<SerializedMeasureFunc>& measureFuncs);
+
 } // namespace facebook::yoga