From 2cc75f7d4781a7ec7b5f4566964b988435859e5e Mon Sep 17 00:00:00 2001
From: Jonathan Dann <jonathan@fb.com>
Date: Sun, 25 Mar 2018 13:56:13 -0700
Subject: [PATCH] Add YGNodeSetChildren(), YGNodeTraversePreOrder()

Reviewed By: Woody17

Differential Revision: D7360203

fbshipit-source-id: 32df8e1213ead03bc0a026ec4bf453bc799bb9ce
---
 ReactCommon/yoga/yoga/Yoga.cpp | 61 ++++++++++++++++++++++++++++++++--
 ReactCommon/yoga/yoga/Yoga.h   | 13 ++++++++
 2 files changed, 72 insertions(+), 2 deletions(-)

diff --git a/ReactCommon/yoga/yoga/Yoga.cpp b/ReactCommon/yoga/yoga/Yoga.cpp
index 570f7ab96..1600b1b05 100644
--- a/ReactCommon/yoga/yoga/Yoga.cpp
+++ b/ReactCommon/yoga/yoga/Yoga.cpp
@@ -295,8 +295,8 @@ static YGNodeRef YGNodeDeepClone(YGNodeRef oldNode) {
 }
 
 void YGNodeFree(const YGNodeRef node) {
-  if (node->getParent()) {
-    node->getParent()->removeChild(node);
+  if (YGNodeRef parent = node->getParent()) {
+    parent->removeChild(node);
     node->setParent(nullptr);
   }
 
@@ -477,6 +477,48 @@ void YGNodeRemoveAllChildren(const YGNodeRef parent) {
   parent->markDirtyAndPropogate();
 }
 
+static void YGNodeSetChildrenInternal(YGNodeRef const parent, const std::vector<YGNodeRef> &children)
+{
+  if (!parent) {
+    return;
+  }
+  if (children.size() == 0) {
+    if (YGNodeGetChildCount(parent) > 0) {
+      for (YGNodeRef const child : parent->getChildren()) {
+        child->setLayout(YGLayout());
+        child->setParent(nullptr);
+      }
+      parent->setChildren(YGVector());
+      parent->markDirtyAndPropogate();
+    }
+  } else {
+    if (YGNodeGetChildCount(parent) > 0) {
+      for (YGNodeRef const oldChild : parent->getChildren()) {
+        // Our new children may have nodes in common with the old children. We don't reset these common nodes.
+        if (std::find(children.begin(), children.end(), oldChild) == children.end()) {
+          oldChild->setLayout(YGLayout());
+          oldChild->setParent(nullptr);
+        }
+      }
+    }
+    parent->setChildren(children);
+    for (YGNodeRef child : children) {
+      child->setParent(parent);
+    }
+    parent->markDirtyAndPropogate();
+  }
+}
+
+void YGNodeSetChildren(YGNodeRef const parent, const YGNodeRef c[], const uint32_t count) {
+  const YGVector children = {c, c + count};
+  YGNodeSetChildrenInternal(parent, children);
+}
+
+void YGNodeSetChildren(YGNodeRef const parent, const std::vector<YGNodeRef> &children)
+{
+  YGNodeSetChildrenInternal(parent, children);
+}
+
 YGNodeRef YGNodeGetChild(const YGNodeRef node, const uint32_t index) {
   if (index < node->getChildren().size()) {
     return node->getChild(index);
@@ -3925,3 +3967,18 @@ void *YGConfigGetContext(const YGConfigRef config) {
 void YGConfigSetNodeClonedFunc(const YGConfigRef config, const YGNodeClonedFunc callback) {
   config->cloneNodeCallback = callback;
 }
+
+static void YGTraverseChildrenPreOrder(const YGVector& children, const std::function<void(YGNodeRef node)>& f) {
+  for (YGNodeRef node : children) {
+    f(node);
+    YGTraverseChildrenPreOrder(node->getChildren(), f);
+  }
+}
+
+void YGTraversePreOrder(YGNodeRef const node, std::function<void(YGNodeRef node)>&& f) {
+  if (!node) {
+    return;
+  }
+  f(node);
+  YGTraverseChildrenPreOrder(node->getChildren(), f);
+}
diff --git a/ReactCommon/yoga/yoga/Yoga.h b/ReactCommon/yoga/yoga/Yoga.h
index 2ec6a687e..1a985fa56 100644
--- a/ReactCommon/yoga/yoga/Yoga.h
+++ b/ReactCommon/yoga/yoga/Yoga.h
@@ -84,6 +84,7 @@ WIN_EXPORT void YGNodeRemoveAllChildren(const YGNodeRef node);
 WIN_EXPORT YGNodeRef YGNodeGetChild(const YGNodeRef node, const uint32_t index);
 WIN_EXPORT YGNodeRef YGNodeGetParent(const YGNodeRef node);
 WIN_EXPORT uint32_t YGNodeGetChildCount(const YGNodeRef node);
+WIN_EXPORT void YGNodeSetChildren(YGNodeRef const parent, const YGNodeRef children[], const uint32_t count);
 
 WIN_EXPORT void YGNodeCalculateLayout(const YGNodeRef node,
                                       const float availableWidth,
@@ -298,3 +299,15 @@ WIN_EXPORT float YGRoundValueToPixelGrid(
     const bool forceFloor);
 
 YG_EXTERN_C_END
+
+#ifdef __cplusplus
+
+#include <functional>
+#include <vector>
+
+// Calls f on each node in the tree including the given node argument.
+extern void YGTraversePreOrder(YGNodeRef const node, std::function<void(YGNodeRef node)>&& f);
+
+extern void YGNodeSetChildren(YGNodeRef const parent, const std::vector<YGNodeRef> &children);
+
+#endif