Remove JNI Binding usage of layoutContext

packages/react-native/ReactAndroid/src/main/jni/first-party/yogajni/jni/LayoutContext.cpp

@@ -0,0 +1,34 @@
+/*
+ * 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 <stack>
+
+#include "LayoutContext.h"
+
+namespace facebook::yoga::vanillajni {
+
+namespace {
+std::stack<PtrJNodeMapVanilla*>& getContexts() {
+  static thread_local std::stack<PtrJNodeMapVanilla*> contexts;
+  return contexts;
+}
+
+} // namespace
+
+LayoutContext::Provider::Provider(PtrJNodeMapVanilla* data) {
+  getContexts().push(data);
+}
+
+LayoutContext::Provider::~Provider() {
+  getContexts().pop();
+}
+
+/*static*/ PtrJNodeMapVanilla* LayoutContext::getNodeMap() {
+  return getContexts().empty() ? nullptr : getContexts().top();
+}
+
+} // namespace facebook::yoga::vanillajni

packages/react-native/ReactAndroid/src/main/jni/first-party/yogajni/jni/LayoutContext.h

@@ -0,0 +1,29 @@
+/*
+ * 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.
+ */
+
+#pragma once
+
+#include <yoga/Yoga.h>
+#include "YGJTypesVanilla.h"
+
+namespace facebook::yoga::vanillajni {
+
+// TODO: This should not be exported or used outside of the JNI bindings
+class YG_EXPORT LayoutContext {
+public:
+  // Sets a context on the current thread for the duration of the Provider's
+  // lifetime. This context should be set during the layout process to allow
+  // layout callbacks to access context-data specific to the layout pass.
+  struct Provider {
+    explicit Provider(PtrJNodeMapVanilla* data);
+    ~Provider();
+  };
+
+  static PtrJNodeMapVanilla* getNodeMap();
+};
+
+} // namespace facebook::yoga::vanillajni

packages/react-native/ReactAndroid/src/main/jni/first-party/yogajni/jni/YGJNIVanilla.cpp

@@ -14,22 +14,17 @@
 #include <iostream>
 #include <memory>
 #include "YogaJniException.h"
+#include "LayoutContext.h"
 
 #include <yoga/Yoga-internal.h>
 #include <yoga/bits/BitCast.h>
 
-// TODO: Reconcile missing layoutContext functionality from callbacks in the C
-// API and use that
-#include <yoga/node/Node.h>
-
 using namespace facebook;
 using namespace facebook::yoga;
 using namespace facebook::yoga::vanillajni;
 
-static inline ScopedLocalRef<jobject> YGNodeJobject(
-    YGNodeConstRef node,
-    void* layoutContext) {
-  return reinterpret_cast<PtrJNodeMapVanilla*>(layoutContext)->ref(node);
+static inline ScopedLocalRef<jobject> YGNodeJobject(YGNodeConstRef node) {
+  return LayoutContext::getNodeMap()->ref(node);
 }
 
 static inline YGNodeRef _jlong2YGNodeRef(jlong addr) {
@@ -129,7 +124,6 @@ static int YGJNILogFunc(
     const YGConfigConstRef config,
     const YGNodeConstRef /*node*/,
     YGLogLevel level,
-    void* /*layoutContext*/,
     const char* format,
     va_list args) {
   va_list argsCopy;
@@ -187,7 +181,7 @@ static void jni_YGConfigSetLoggerJNI(
     }
 
     *context = newGlobalRef(env, logger);
-    static_cast<yoga::Config*>(config)->setLogger(YGJNILogFunc);
+    YGConfigSetLogger(config, YGJNILogFunc);
   } else {
     if (context != nullptr) {
       delete context;
@@ -278,12 +272,11 @@ static void jni_YGNodeRemoveChildJNI(
 static void YGTransferLayoutOutputsRecursive(
     JNIEnv* env,
     jobject thiz,
-    YGNodeRef root,
-    void* layoutContext) {
+    YGNodeRef root) {
   if (!YGNodeGetHasNewLayout(root)) {
     return;
   }
-  auto obj = YGNodeJobject(root, layoutContext);
+  auto obj = YGNodeJobject(root);
   if (!obj) {
     return;
   }
@@ -351,8 +344,7 @@ static void YGTransferLayoutOutputsRecursive(
   YGNodeSetHasNewLayout(root, false);
 
   for (size_t i = 0; i < YGNodeGetChildCount(root); i++) {
-    YGTransferLayoutOutputsRecursive(
-        env, thiz, YGNodeGetChild(root, i), layoutContext);
+    YGTransferLayoutOutputsRecursive(env, thiz, YGNodeGetChild(root, i));
   }
 }
 
@@ -366,21 +358,22 @@ static void jni_YGNodeCalculateLayoutJNI(
     jobjectArray javaNodes) {
 
   try {
-    void* layoutContext = nullptr;
+    PtrJNodeMapVanilla* layoutContext = nullptr;
     auto map = PtrJNodeMapVanilla{};
     if (nativePointers) {
       map = PtrJNodeMapVanilla{nativePointers, javaNodes};
       layoutContext = &map;
     }
 
+    LayoutContext::Provider contextProvider(layoutContext);
+
     const YGNodeRef root = _jlong2YGNodeRef(nativePointer);
-    YGNodeCalculateLayoutWithContext(
+    YGNodeCalculateLayout(
         root,
         static_cast<float>(width),
         static_cast<float>(height),
-        YGNodeStyleGetDirection(_jlong2YGNodeRef(nativePointer)),
-        layoutContext);
-    YGTransferLayoutOutputsRecursive(env, obj, root, layoutContext);
+        YGNodeStyleGetDirection(_jlong2YGNodeRef(nativePointer)));
+    YGTransferLayoutOutputsRecursive(env, obj, root);
   } catch (const YogaJniException& jniException) {
     ScopedLocalRef<jthrowable> throwable = jniException.getThrowable();
     if (throwable.get()) {
@@ -647,9 +640,8 @@ static YGSize YGJNIMeasureFunc(
     float width,
     YGMeasureMode widthMode,
     float height,
-    YGMeasureMode heightMode,
-    void* layoutContext) {
-  if (auto obj = YGNodeJobject(node, layoutContext)) {
+    YGMeasureMode heightMode) {
+  if (auto obj = YGNodeJobject(node)) {
     YGTransferLayoutDirection(node, obj.get());
     JNIEnv* env = getCurrentEnv();
     auto objectClass = facebook::yoga::vanillajni::make_local_ref(
@@ -683,16 +675,13 @@ static void jni_YGNodeSetHasMeasureFuncJNI(
     jobject /*obj*/,
     jlong nativePointer,
     jboolean hasMeasureFunc) {
-  static_cast<yoga::Node*>(_jlong2YGNodeRef(nativePointer))
-      ->setMeasureFunc(hasMeasureFunc ? YGJNIMeasureFunc : nullptr);
+  YGNodeSetMeasureFunc(
+      _jlong2YGNodeRef(nativePointer),
+      hasMeasureFunc ? YGJNIMeasureFunc : nullptr);
 }
 
-static float YGJNIBaselineFunc(
-    YGNodeConstRef node,
-    float width,
-    float height,
-    void* layoutContext) {
-  if (auto obj = YGNodeJobject(node, layoutContext)) {
+static float YGJNIBaselineFunc(YGNodeConstRef node, float width, float height) {
+  if (auto obj = YGNodeJobject(node)) {
     JNIEnv* env = getCurrentEnv();
     auto objectClass = facebook::yoga::vanillajni::make_local_ref(
         env, env->GetObjectClass(obj.get()));
@@ -710,8 +699,9 @@ static void jni_YGNodeSetHasBaselineFuncJNI(
     jobject /*obj*/,
     jlong nativePointer,
     jboolean hasBaselineFunc) {
-  static_cast<yoga::Node*>(_jlong2YGNodeRef(nativePointer))
-      ->setBaselineFunc(hasBaselineFunc ? YGJNIBaselineFunc : nullptr);
+  YGNodeSetBaselineFunc(
+      _jlong2YGNodeRef(nativePointer),
+      hasBaselineFunc ? YGJNIBaselineFunc : nullptr);
 }
 
 static void jni_YGNodePrintJNI(

packages/react-native/ReactAndroid/src/main/jni/first-party/yogajni/jni/YGJTypesVanilla.h

@@ -5,6 +5,8 @@
  * LICENSE file in the root directory of this source tree.
  */
 
+#pragma once
+
 #include <map>
 #include <vector>
 

packages/react-native/ReactCommon/yoga/yoga/algorithm/Baseline.cpp

@@ -62,4 +62,23 @@ float calculateBaseline(const yoga::Node* node, void* layoutContext) {
   return baseline + baselineChild->getLayout().position[YGEdgeTop];
 }
 
+bool isBaselineLayout(const yoga::Node* node) {
+  if (isColumn(node->getStyle().flexDirection())) {
+    return false;
+  }
+  if (node->getStyle().alignItems() == YGAlignBaseline) {
+    return true;
+  }
+  const auto childCount = node->getChildCount();
+  for (size_t i = 0; i < childCount; i++) {
+    auto child = node->getChild(i);
+    if (child->getStyle().positionType() != YGPositionTypeAbsolute &&
+        child->getStyle().alignSelf() == YGAlignBaseline) {
+      return true;
+    }
+  }
+
+  return false;
+}
+
 } // namespace facebook::yoga

packages/react-native/ReactCommon/yoga/yoga/algorithm/Baseline.h

@@ -15,4 +15,7 @@ namespace facebook::yoga {
 // Calculate baseline represented as an offset from the top edge of the node.
 float calculateBaseline(const yoga::Node* node, void* layoutContext);
 
+// Whether any of the children of this node participate in baseline alignment
+bool isBaselineLayout(const yoga::Node* node);
+
 } // namespace facebook::yoga

packages/react-native/ReactCommon/yoga/yoga/algorithm/CalculateLayout.cpp

@@ -50,25 +50,6 @@ bool calculateLayoutInternal(
     const uint32_t depth,
     const uint32_t generationCount);
 
-static bool isBaselineLayout(const yoga::Node* node) {
-  if (isColumn(node->getStyle().flexDirection())) {
-    return false;
-  }
-  if (node->getStyle().alignItems() == YGAlignBaseline) {
-    return true;
-  }
-  const auto childCount = node->getChildCount();
-  for (size_t i = 0; i < childCount; i++) {
-    auto child = node->getChild(i);
-    if (child->getStyle().positionType() != YGPositionTypeAbsolute &&
-        child->getStyle().alignSelf() == YGAlignBaseline) {
-      return true;
-    }
-  }
-
-  return false;
-}
-
 static inline float dimensionWithMargin(
     const yoga::Node* const node,
     const YGFlexDirection axis,