Remove JNI Binding usage of layoutContext (#39402)

Summary:
Pull Request resolved: #39402

X-link: facebook/yoga#1377

To avoid keeping a per-node mapping on native Yoga nodes to Java nodes, a per-layout context was added, to be able to pass information from the start of the layout, to measure functions, log functions, etc.

The way this was done was super invasive, and added quite a few private APIs used only by the JNI functions.

This change removes the context-using functions from the JNI bindings in favor of it managing its own context. Next diff removes all the cruft.

Reviewed By: javache

Differential Revision: D49179243

fbshipit-source-id: 423a64dfadba8e43c51425b31572d2389a6f60e6

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>