Allow to use JNI without global refs

Summary:
@public

Adds the ability to opt into avoiding global weak JNI refs via `YogaConfig`.

Note that only homogeneous trees are supported, i.e. **mixing weak-ref and non-weak-ref nodes will break!**

Not using JNI refs hopefully will help with avoiding JNI reference table overflows, and will help creating trees on multiple threads, as no lock has to be acquired at any time.

Reviewed By: SidharthGuglani

Differential Revision: D14151037

fbshipit-source-id: 56d94713d39aee080d54be4cb4cdf5e3eccb473a

java/com/facebook/yoga/YogaNodeJNI.java

@@ -40,6 +40,7 @@ public class YogaNodeJNI extends YogaNode implements Cloneable {
   private int mEdgeSetFlag = 0;
 
   private boolean mHasSetPosition = false;
+  private final boolean mAvoidGlobalJNIRefs;
 
   @DoNotStrip
   private float mWidth = YogaConstants.UNDEFINED;
@@ -81,15 +82,17 @@ public class YogaNodeJNI extends YogaNode implements Cloneable {
 
   private native long jni_YGNodeNew();
   public YogaNodeJNI() {
+    mAvoidGlobalJNIRefs = false;
     mNativePointer = jni_YGNodeNew();
     if (mNativePointer == 0) {
       throw new IllegalStateException("Failed to allocate native memory");
     }
   }
 
-  private native long jni_YGNodeNewWithConfig(long configPointer);
+  private native long jni_YGNodeNewWithConfig(long configPointer, boolean avoidGlobalJNIRefs);
   public YogaNodeJNI(YogaConfig config) {
-    mNativePointer = jni_YGNodeNewWithConfig(config.mNativePointer);
+    mAvoidGlobalJNIRefs = config.avoidGlobalJNIRefs;
+    mNativePointer = jni_YGNodeNewWithConfig(config.mNativePointer, mAvoidGlobalJNIRefs);
     if (mNativePointer == 0) {
       throw new IllegalStateException("Failed to allocate native memory");
     }
@@ -188,13 +191,13 @@ public class YogaNodeJNI extends YogaNode implements Cloneable {
 
   private static native void jni_YGNodeSetOwner(long nativePointer, long newOwnerNativePointer);
 
-  private native long jni_YGNodeClone(long nativePointer, Object newNode);
+  private native long jni_YGNodeClone(long nativePointer, Object newNode, boolean avoidGlobalJNIRefs);
 
   @Override
   public YogaNodeJNI clone() {
     try {
       YogaNodeJNI clonedYogaNode = (YogaNodeJNI) super.clone();
-      long clonedNativePointer = jni_YGNodeClone(mNativePointer, clonedYogaNode);
+      long clonedNativePointer = jni_YGNodeClone(mNativePointer, clonedYogaNode, mAvoidGlobalJNIRefs);
 
       if (mChildren != null) {
         for (YogaNodeJNI child : mChildren) {
@@ -222,7 +225,7 @@ public class YogaNodeJNI extends YogaNode implements Cloneable {
   public YogaNodeJNI cloneWithNewChildren() {
     try {
       YogaNodeJNI clonedYogaNode = (YogaNodeJNI) super.clone();
-      long clonedNativePointer = jni_YGNodeClone(mNativePointer, clonedYogaNode);
+      long clonedNativePointer = jni_YGNodeClone(mNativePointer, clonedYogaNode, mAvoidGlobalJNIRefs);
       clonedYogaNode.mOwner = null;
       clonedYogaNode.mNativePointer = clonedNativePointer;
       clonedYogaNode.clearChildren();
@@ -276,9 +279,29 @@ public class YogaNodeJNI extends YogaNode implements Cloneable {
     return mChildren == null ? -1 : mChildren.indexOf(child);
   }
 
-  private static native void jni_YGNodeCalculateLayout(long nativePointer, float width, float height);
+  private static native void jni_YGNodeCalculateLayout(long nativePointer, float width, float height, long[] nativePointers, YogaNodeJNI[] nodes);
   public void calculateLayout(float width, float height) {
-    jni_YGNodeCalculateLayout(mNativePointer, width, height);
+    long[] nativePointers = null;
+    YogaNodeJNI[] nodes = null;
+
+    if (mAvoidGlobalJNIRefs) {
+      ArrayList<YogaNodeJNI> n = new ArrayList<>();
+      n.add(this);
+      for (int i = 0; i < n.size(); ++i) {
+        List<YogaNodeJNI> children  = n.get(i).mChildren;
+        if (children != null) {
+          n.addAll(children);
+        }
+      }
+
+      nodes = n.toArray(new YogaNodeJNI[n.size()]);
+      nativePointers = new long[nodes.length];
+      for (int i = 0; i < nodes.length; ++i) {
+        nativePointers[i] = nodes[i].mNativePointer;
+      }
+    }
+
+    jni_YGNodeCalculateLayout(mNativePointer, width, height, nativePointers, nodes);
   }
 
   public boolean hasNewLayout() {

java/jni/YGJNI.cpp

@@ -9,6 +9,7 @@
 #include <yoga/Yoga.h>
 #include <yoga/log.h>
 #include <iostream>
+#include <map>
 
 using namespace facebook::jni;
 using namespace std;
@@ -22,6 +23,34 @@ struct JYogaConfig : public JavaClass<JYogaConfig> {
   static constexpr auto kJavaDescriptor = "Lcom/facebook/yoga/YogaConfig;";
 };
 
+class PtrJNodeMap {
+  using JNodeArray = JArrayClass<JYogaNode::javaobject>;
+  std::map<YGNodeRef, size_t> ptrsToIdxs_;
+  alias_ref<JNodeArray> javaNodes_;
+
+public:
+  PtrJNodeMap() : ptrsToIdxs_{}, javaNodes_{} {}
+  PtrJNodeMap(
+      alias_ref<JArrayLong> nativePointers,
+      alias_ref<JNodeArray> javaNodes)
+      : javaNodes_{javaNodes} {
+    auto pin = nativePointers->pinCritical();
+    auto ptrs = pin.get();
+    for (size_t i = 0, n = pin.size(); i < n; ++i) {
+      ptrsToIdxs_[(YGNodeRef) ptrs[i]] = i;
+    }
+  }
+
+  local_ref<JYogaNode> ref(YGNodeRef node) {
+    auto idx = ptrsToIdxs_.find(node);
+    if (idx == ptrsToIdxs_.end()) {
+      return local_ref<JYogaNode>{};
+    } else {
+      return javaNodes_->getElement(idx->second);
+    }
+  }
+};
+
 struct YGConfigContext {
   global_ref<jobject>* logger;
   global_ref<jobject>* config;
@@ -34,8 +63,15 @@ struct YGConfigContext {
   }
 };
 
-static inline weak_ref<JYogaNode>* YGNodeJobject(YGNodeRef node) {
+static inline local_ref<JYogaNode> YGNodeJobject(
-  return reinterpret_cast<weak_ref<JYogaNode>*>(node->getContext());
+    YGNodeRef node,
+    void* layoutContext) {
+  if (layoutContext == nullptr) {
+    return reinterpret_cast<weak_ref<JYogaNode>*>(node->getContext())
+        ->lockLocal();
+  } else {
+    return reinterpret_cast<PtrJNodeMap*>(layoutContext)->ref(node);
+  }
 }
 
 static void YGTransferLayoutDirection(
@@ -47,11 +83,13 @@ static void YGTransferLayoutDirection(
       layoutDirectionField, static_cast<jint>(YGNodeLayoutGetDirection(node)));
 }
 
-static void YGTransferLayoutOutputsRecursive(YGNodeRef root) {
+static void YGTransferLayoutOutputsRecursive(
+    YGNodeRef root,
+    void* layoutContext) {
   if (!root->getHasNewLayout()) {
     return;
   }
-  auto obj = YGNodeJobject(root)->lockLocal();
+  auto obj = YGNodeJobject(root, layoutContext);
   if (!obj) {
     Log::log(
         root,
@@ -98,7 +136,7 @@ static void YGTransferLayoutOutputsRecursive(YGNodeRef root) {
   static auto doesLegacyStretchBehaviour = obj->getClass()->getField<jboolean>(
       "mDoesLegacyStretchFlagAffectsLayout");
 
-  /* Those flags needs be in sync with YogaNodeJNI.java */
+  /* Those flags needs be in sync with YogaNode.java */
   const int MARGIN = 1;
   const int PADDING = 2;
   const int BORDER = 4;
@@ -149,12 +187,12 @@ static void YGTransferLayoutOutputsRecursive(YGNodeRef root) {
   root->setHasNewLayout(false);
 
   for (uint32_t i = 0; i < YGNodeGetChildCount(root); i++) {
-    YGTransferLayoutOutputsRecursive(YGNodeGetChild(root, i));
+    YGTransferLayoutOutputsRecursive(YGNodeGetChild(root, i), layoutContext);
   }
 }
 
-static void YGPrint(YGNodeRef node) {
+static void YGPrint(YGNodeRef node, void* layoutContext) {
-  if (auto obj = YGNodeJobject(node)->lockLocal()) {
+  if (auto obj = YGNodeJobject(node, layoutContext)) {
     cout << obj->toString() << endl;
   } else {
     Log::log(
@@ -165,8 +203,12 @@ static void YGPrint(YGNodeRef node) {
   }
 }
 
-static float YGJNIBaselineFunc(YGNodeRef node, float width, float height) {
+static float YGJNIBaselineFunc(
-  if (auto obj = YGNodeJobject(node)->lockLocal()) {
+    YGNodeRef node,
+    float width,
+    float height,
+    void* layoutContext) {
+  if (auto obj = YGNodeJobject(node, layoutContext)) {
     static auto baselineFunc =
         findClassStatic("com/facebook/yoga/YogaNodeJNI")
             ->getMethod<jfloat(jfloat, jfloat)>("baseline");
@@ -187,7 +229,8 @@ static inline YGConfigRef _jlong2YGConfigRef(jlong addr) {
 static YGNodeRef YGJNIOnNodeClonedFunc(
     YGNodeRef oldNode,
     YGNodeRef owner,
-    int childIndex) {
+    int childIndex,
+    void* layoutContext) {
   auto config = oldNode->getConfig();
   if (!config) {
     return nullptr;
@@ -203,8 +246,8 @@ static YGNodeRef YGJNIOnNodeClonedFunc(
 
   auto newNode = onNodeClonedFunc(
       javaConfig->get(),
-      YGNodeJobject(oldNode)->lockLocal(),
+      YGNodeJobject(oldNode, layoutContext),
-      YGNodeJobject(owner)->lockLocal(),
+      YGNodeJobject(owner, layoutContext),
       childIndex);
 
   static auto replaceChild =
@@ -212,7 +255,7 @@ static YGNodeRef YGJNIOnNodeClonedFunc(
           ->getMethod<jlong(local_ref<JYogaNode>, jint)>("replaceChild");
 
   jlong newNodeNativePointer =
-      replaceChild(YGNodeJobject(owner)->lockLocal(), newNode, childIndex);
+      replaceChild(YGNodeJobject(owner, layoutContext), newNode, childIndex);
 
   return _jlong2YGNodeRef(newNodeNativePointer);
 }
@@ -222,8 +265,9 @@ static YGSize YGJNIMeasureFunc(
     float width,
     YGMeasureMode widthMode,
     float height,
-    YGMeasureMode heightMode) {
+    YGMeasureMode heightMode,
-  if (auto obj = YGNodeJobject(node)->lockLocal()) {
+    void* layoutContext) {
+  if (auto obj = YGNodeJobject(node, layoutContext)) {
     static auto measureFunc =
         findClassStatic("com/facebook/yoga/YogaNodeJNI")
             ->getMethod<jlong(jfloat, jint, jfloat, jint)>("measure");
@@ -264,6 +308,7 @@ static int YGJNILogFunc(
     const YGConfigRef config,
     const YGNodeRef node,
     YGLogLevel level,
+    void* layoutContext,
     const char* format,
     va_list args) {
   int result = vsnprintf(NULL, 0, format, args);
@@ -279,7 +324,7 @@ static int YGJNILogFunc(
       JYogaLogLevel::javaClassStatic()
           ->getStaticMethod<JYogaLogLevel::javaobject(jint)>("fromInt");
 
-  if (auto obj = YGNodeJobject(node)->lockLocal()) {
+  if (auto obj = YGNodeJobject(node, layoutContext)) {
     auto jlogger =
         reinterpret_cast<global_ref<jobject>*>(YGConfigGetContext(config));
     logFunc(
@@ -296,16 +341,18 @@ static int YGJNILogFunc(
 jlong jni_YGNodeNew(alias_ref<jobject> thiz) {
   const YGNodeRef node = YGNodeNew();
   node->setContext(new weak_ref<jobject>(make_weak(thiz)));
-  // YGNodeSetContext(node, new weak_ref<jobject>(make_weak(thiz)));
   node->setPrintFunc(YGPrint);
-  // YGNodeSetPrintFunc(node, YGPrint);
   return reinterpret_cast<jlong>(node);
 }
 
-jlong jni_YGNodeNewWithConfig(alias_ref<jobject> thiz, jlong configPointer) {
+jlong jni_YGNodeNewWithConfig(
+    alias_ref<jobject> thiz,
+    jlong configPointer,
+    jboolean avoidGlobalJNIRefs) {
   const YGNodeRef node = YGNodeNewWithConfig(_jlong2YGConfigRef(configPointer));
-  node->setContext(new weak_ref<jobject>(make_weak(thiz)));
+  if (!avoidGlobalJNIRefs) {
-  node->setPrintFunc(YGPrint);
+    node->setContext(new weak_ref<jobject>(make_weak(thiz)));
+  }
   return reinterpret_cast<jlong>(node);
 }
 
@@ -319,10 +366,13 @@ void jni_YGNodeSetOwner(jlong nativePointer, jlong newOwnerNativePointer) {
 jlong jni_YGNodeClone(
     alias_ref<jobject> thiz,
     jlong nativePointer,
-    alias_ref<jobject> clonedJavaObject) {
+    alias_ref<jobject> clonedJavaObject,
+    jboolean avoidGlobalJNIRefs) {
   const YGNodeRef clonedYogaNode = YGNodeClone(_jlong2YGNodeRef(nativePointer));
-  clonedYogaNode->setContext(
+  if (!avoidGlobalJNIRefs) {
-      new weak_ref<jobject>(make_weak(clonedJavaObject)));
+    clonedYogaNode->setContext(
+        new weak_ref<jobject>(make_weak(clonedJavaObject)));
+  }
   return reinterpret_cast<jlong>(clonedYogaNode);
 }
 
@@ -331,7 +381,10 @@ void jni_YGNodeFree(jlong nativePointer) {
     return;
   }
   const YGNodeRef node = _jlong2YGNodeRef(nativePointer);
-  delete YGNodeJobject(node);
+  auto context = node->getContext();
+  if (context != nullptr) {
+    delete reinterpret_cast<weak_ref<JYogaNode>*>(node->getContext());
+  }
   YGNodeFree(node);
 }
 
@@ -345,7 +398,6 @@ void jni_YGNodeReset(jlong nativePointer) {
   void* context = node->getContext();
   YGNodeReset(node);
   node->setContext(context);
-  node->setPrintFunc(YGPrint);
 }
 
 void jni_YGNodePrint(jlong nativePointer) {
@@ -384,14 +436,25 @@ void jni_YGNodeCalculateLayout(
     alias_ref<jclass>,
     jlong nativePointer,
     jfloat width,
-    jfloat height) {
+    jfloat height,
+    alias_ref<JArrayLong> nativePointers,
+    alias_ref<JArrayClass<JYogaNode::javaobject>> javaNodes) {
+
+  void* layoutContext = nullptr;
+  auto map = PtrJNodeMap{};
+  if (nativePointers) {
+    map = PtrJNodeMap{nativePointers, javaNodes};
+    layoutContext = &map;
+  }
+
   const YGNodeRef root = _jlong2YGNodeRef(nativePointer);
-  YGNodeCalculateLayout(
+  YGNodeCalculateLayoutWithContext(
       root,
       static_cast<float>(width),
       static_cast<float>(height),
-      YGNodeStyleGetDirection(_jlong2YGNodeRef(nativePointer)));
+      YGNodeStyleGetDirection(_jlong2YGNodeRef(nativePointer)),
-  YGTransferLayoutOutputsRecursive(root);
+      layoutContext);
+  YGTransferLayoutOutputsRecursive(root, layoutContext);
 }
 
 void jni_YGNodeMarkDirty(jlong nativePointer) {
@@ -622,9 +685,9 @@ void jni_YGConfigSetHasCloneNodeFunc(
       YGConfigSetContext(config, context);
     }
     context->config = new global_ref<jobject>(make_global(thiz));
-    YGConfigSetCloneNodeFunc(config, YGJNIOnNodeClonedFunc);
+    config->setCloneNodeCallback(YGJNIOnNodeClonedFunc);
   } else {
-    YGConfigSetCloneNodeFunc(config, nullptr);
+    config->setCloneNodeCallback(nullptr);
   }
 }
 
@@ -645,9 +708,9 @@ void jni_YGConfigSetLogger(
       YGConfigSetContext(config, context);
     }
     context->logger = new global_ref<jobject>(make_global(logger));
-    YGConfigSetLogger(config, YGJNILogFunc);
+    config->setLogger(YGJNILogFunc);
   } else {
-    YGConfigSetLogger(config, NULL);
+    config->setLogger(nullptr);
   }
 }
 

yoga/YGConfig.h

@@ -10,7 +10,7 @@
 #include "Yoga.h"
 
 struct YGConfig {
-  using LogWithContextFn = void (*)(
+  using LogWithContextFn = int (*)(
       YGConfigRef config,
       YGNodeRef node,
       YGLogLevel level,