Move calculate layout method to vanilla JNI

Summary: Using vanilla JNI for calculate layout

Reviewed By: amir-shalem

Differential Revision: D17714219

fbshipit-source-id: bb05de4a0112eefc2b731997a4c1ecef5c0c7361

java/com/facebook/yoga/YogaNative.java

@@ -132,6 +132,7 @@ public class YogaNative {
   static native boolean jni_YGNodeIsReferenceBaselineJNI(long nativePointer);
   static native void jni_YGNodeClearChildrenJNI(long nativePointer);
   static native void jni_YGNodeRemoveChildJNI(long nativePointer, long childPointer);
+  static native void jni_YGNodeCalculateLayoutJNI(long nativePointer, float width, float height, long[] nativePointers, YogaNodeJNIBase[] nodes);
   static native void jni_YGNodeMarkDirtyJNI(long nativePointer);
   static native void jni_YGNodeMarkDirtyAndPropogateToDescendantsJNI(long nativePointer);
   static native boolean jni_YGNodeIsDirtyJNI(long nativePointer);

java/com/facebook/yoga/YogaNodeJNIBase.java

@@ -197,7 +197,10 @@ public abstract class YogaNodeJNIBase extends YogaNode implements Cloneable {
       nativePointers[i] = nodes[i].mNativePointer;
     }
 
-    YogaNative.jni_YGNodeCalculateLayout(mNativePointer, width, height, nativePointers, nodes);
+    if (useVanillaJNI)
+      YogaNative.jni_YGNodeCalculateLayoutJNI(mNativePointer, width, height, nativePointers, nodes);
+    else
+      YogaNative.jni_YGNodeCalculateLayout(mNativePointer, width, height, nativePointers, nodes);
   }
 
   public void dirty() {

java/jni/YGJNI.cpp

@@ -22,23 +22,6 @@ using namespace facebook::jni;
 using namespace std;
 using facebook::yoga::detail::Log;
 
-const short int LAYOUT_EDGE_SET_FLAG_INDEX = 0;
-const short int LAYOUT_WIDTH_INDEX = 1;
-const short int LAYOUT_HEIGHT_INDEX = 2;
-const short int LAYOUT_LEFT_INDEX = 3;
-const short int LAYOUT_TOP_INDEX = 4;
-const short int LAYOUT_DIRECTION_INDEX = 5;
-const short int LAYOUT_MARGIN_START_INDEX = 6;
-const short int LAYOUT_PADDING_START_INDEX = 10;
-const short int LAYOUT_BORDER_START_INDEX = 14;
-
-namespace {
-
-const int DOES_LEGACY_STRETCH_BEHAVIOUR = 8;
-const int HAS_NEW_LAYOUT = 16;
-
-} // namespace
-
 static inline local_ref<JYogaNode> YGNodeJobject(
     YGNodeRef node,
     void* layoutContext) {

java/jni/YGJNI.h

@@ -47,8 +47,21 @@ enum YGStyleInput {
   IsReferenceBaseline,
 };
 
+const short int LAYOUT_EDGE_SET_FLAG_INDEX = 0;
+const short int LAYOUT_WIDTH_INDEX = 1;
+const short int LAYOUT_HEIGHT_INDEX = 2;
+const short int LAYOUT_LEFT_INDEX = 3;
+const short int LAYOUT_TOP_INDEX = 4;
+const short int LAYOUT_DIRECTION_INDEX = 5;
+const short int LAYOUT_MARGIN_START_INDEX = 6;
+const short int LAYOUT_PADDING_START_INDEX = 10;
+const short int LAYOUT_BORDER_START_INDEX = 14;
+
 namespace {
 
+const int DOES_LEGACY_STRETCH_BEHAVIOUR = 8;
+const int HAS_NEW_LAYOUT = 16;
+
 union YGNodeContext {
   uintptr_t edgesSet = 0;
   void* asVoidPtr;

java/jni/YGJNIVanilla.cpp

@@ -9,6 +9,19 @@
 #include <yoga/YGNode.h>
 #include <cstring>
 #include "YGJNI.h"
+#include "common.h"
+#include "YGJTypesVanilla.h"
+#include <yoga/log.h>
+
+using namespace facebook::yoga::vanillajni;
+using facebook::yoga::detail::Log;
+
+static inline ScopedLocalRef<jobject> YGNodeJobject(
+    YGNodeRef node,
+    void* layoutContext) {
+  return reinterpret_cast<PtrJNodeMap*>(layoutContext)
+      ->ref(getCurrentEnv(), node);
+}
 
 static inline YGNodeRef _jlong2YGNodeRef(jlong addr) {
   return reinterpret_cast<YGNodeRef>(static_cast<intptr_t>(addr));
@@ -144,6 +157,122 @@ static void jni_YGNodeRemoveChildJNI(
       _jlong2YGNodeRef(nativePointer), _jlong2YGNodeRef(childPointer));
 }
 
+static void YGTransferLayoutOutputsRecursive(
+    JNIEnv* env,
+    jobject thiz,
+    YGNodeRef root,
+    void* layoutContext) {
+  if (!root->getHasNewLayout()) {
+    return;
+  }
+  auto obj = YGNodeJobject(root, layoutContext);
+  if (!obj) {
+    Log::log(
+        root,
+        YGLogLevelError,
+        nullptr,
+        "Java YGNode was GCed during layout calculation\n");
+    return;
+  }
+
+  auto edgesSet = YGNodeEdges{root};
+
+  bool marginFieldSet = edgesSet.has(YGNodeEdges::MARGIN);
+  bool paddingFieldSet = edgesSet.has(YGNodeEdges::PADDING);
+  bool borderFieldSet = edgesSet.has(YGNodeEdges::BORDER);
+
+  int fieldFlags = edgesSet.get();
+  fieldFlags |= HAS_NEW_LAYOUT;
+  if (YGNodeLayoutGetDidLegacyStretchFlagAffectLayout(root)) {
+    fieldFlags |= DOES_LEGACY_STRETCH_BEHAVIOUR;
+  }
+
+  const int arrSize = 6 + (marginFieldSet ? 4 : 0) + (paddingFieldSet ? 4 : 0) +
+      (borderFieldSet ? 4 : 0);
+  float arr[18];
+  arr[LAYOUT_EDGE_SET_FLAG_INDEX] = fieldFlags;
+  arr[LAYOUT_WIDTH_INDEX] = YGNodeLayoutGetWidth(root);
+  arr[LAYOUT_HEIGHT_INDEX] = YGNodeLayoutGetHeight(root);
+  arr[LAYOUT_LEFT_INDEX] = YGNodeLayoutGetLeft(root);
+  arr[LAYOUT_TOP_INDEX] = YGNodeLayoutGetTop(root);
+  arr[LAYOUT_DIRECTION_INDEX] =
+      static_cast<jint>(YGNodeLayoutGetDirection(root));
+  if (marginFieldSet) {
+    arr[LAYOUT_MARGIN_START_INDEX] = YGNodeLayoutGetMargin(root, YGEdgeLeft);
+    arr[LAYOUT_MARGIN_START_INDEX + 1] = YGNodeLayoutGetMargin(root, YGEdgeTop);
+    arr[LAYOUT_MARGIN_START_INDEX + 2] =
+        YGNodeLayoutGetMargin(root, YGEdgeRight);
+    arr[LAYOUT_MARGIN_START_INDEX + 3] =
+        YGNodeLayoutGetMargin(root, YGEdgeBottom);
+  }
+  if (paddingFieldSet) {
+    int paddingStartIndex =
+        LAYOUT_PADDING_START_INDEX - (marginFieldSet ? 0 : 4);
+    arr[paddingStartIndex] = YGNodeLayoutGetPadding(root, YGEdgeLeft);
+    arr[paddingStartIndex + 1] = YGNodeLayoutGetPadding(root, YGEdgeTop);
+    arr[paddingStartIndex + 2] = YGNodeLayoutGetPadding(root, YGEdgeRight);
+    arr[paddingStartIndex + 3] = YGNodeLayoutGetPadding(root, YGEdgeBottom);
+  }
+
+  if (borderFieldSet) {
+    int borderStartIndex = LAYOUT_BORDER_START_INDEX -
+        (marginFieldSet ? 0 : 4) - (paddingFieldSet ? 0 : 4);
+    arr[borderStartIndex] = YGNodeLayoutGetBorder(root, YGEdgeLeft);
+    arr[borderStartIndex + 1] = YGNodeLayoutGetBorder(root, YGEdgeTop);
+    arr[borderStartIndex + 2] = YGNodeLayoutGetBorder(root, YGEdgeRight);
+    arr[borderStartIndex + 3] = YGNodeLayoutGetBorder(root, YGEdgeBottom);
+  }
+
+  // Don't change this field name without changing the name of the field in
+  // Database.java
+  auto objectClass = facebook::yoga::vanillajni::make_local_ref(
+      env, env->GetObjectClass(obj.get()));
+  static const jfieldID arrField = facebook::yoga::vanillajni::getFieldId(
+      env, objectClass.get(), "arr", "[F");
+
+  ScopedLocalRef<jfloatArray> arrFinal =
+      make_local_ref(env, env->NewFloatArray(arrSize));
+  env->SetFloatArrayRegion(arrFinal.get(), 0, arrSize, arr);
+  env->SetObjectField(obj.get(), arrField, arrFinal.get());
+
+  root->setHasNewLayout(false);
+
+  for (uint32_t i = 0; i < YGNodeGetChildCount(root); i++) {
+    YGTransferLayoutOutputsRecursive(
+        env, thiz, YGNodeGetChild(root, i), layoutContext);
+  }
+}
+
+static void jni_YGNodeCalculateLayoutJNI(
+    JNIEnv* env,
+    jobject obj,
+    jlong nativePointer,
+    jfloat width,
+    jfloat height,
+    jlongArray nativePointers,
+    jobjectArray javaNodes) {
+
+  void* layoutContext = nullptr;
+  auto map = PtrJNodeMap{};
+  if (nativePointers) {
+    size_t nativePointersSize = env->GetArrayLength(nativePointers);
+    jlong result[nativePointersSize];
+    env->GetLongArrayRegion(nativePointers, 0, nativePointersSize, result);
+
+    map = PtrJNodeMap{result, nativePointersSize, javaNodes};
+    layoutContext = &map;
+  }
+
+  const YGNodeRef root = _jlong2YGNodeRef(nativePointer);
+  YGNodeCalculateLayoutWithContext(
+      root,
+      static_cast<float>(width),
+      static_cast<float>(height),
+      YGNodeStyleGetDirection(_jlong2YGNodeRef(nativePointer)),
+      layoutContext);
+  YGTransferLayoutOutputsRecursive(env, obj, root, layoutContext);
+}
+
 static void jni_YGNodeMarkDirtyJNI(
     JNIEnv* env,
     jobject obj,
@@ -407,28 +536,6 @@ static void jni_YGNodeSetStyleInputsJNI(
   YGNodeSetStyleInputs(_jlong2YGNodeRef(nativePointer), result, size);
 }
 
-void assertNoPendingJniException(JNIEnv* env) {
-  // This method cannot call any other method of the library, since other
-  // methods of the library use it to check for exceptions too
-  if (env->ExceptionCheck()) {
-    env->ExceptionDescribe();
-  }
-}
-
-void registerNativeMethods(
-    JNIEnv* env,
-    const char* className,
-    JNINativeMethod methods[],
-    size_t numMethods) {
-  jclass clazz = env->FindClass(className);
-
-  assertNoPendingJniException(env);
-
-  env->RegisterNatives(clazz, methods, numMethods);
-
-  assertNoPendingJniException(env);
-}
-
 static JNINativeMethod methods[] = {
     {"jni_YGConfigNewJNI", "()J", (void*) jni_YGConfigNewJNI},
     //    {"jni_YGConfigFreeJNI", "(J)V", (void*) jni_YGConfigFreeJNI},
@@ -462,6 +569,9 @@ static JNINativeMethod methods[] = {
      (void*) jni_YGNodeIsReferenceBaselineJNI},
     {"jni_YGNodeClearChildrenJNI", "(J)V", (void*) jni_YGNodeClearChildrenJNI},
     {"jni_YGNodeRemoveChildJNI", "(JJ)V", (void*) jni_YGNodeRemoveChildJNI},
+    {"jni_YGNodeCalculateLayoutJNI",
+     "(JFF[J[Lcom/facebook/yoga/YogaNodeJNIBase;)V",
+     (void*) jni_YGNodeCalculateLayoutJNI},
     {"jni_YGNodeMarkDirtyJNI", "(J)V", (void*) jni_YGNodeMarkDirtyJNI},
     {"jni_YGNodeMarkDirtyAndPropogateToDescendantsJNI",
      "(J)V",
@@ -660,7 +770,7 @@ static JNINativeMethod methods[] = {
 };
 
 void YGJNIVanilla::registerNatives(JNIEnv* env) {
-  registerNativeMethods(
+  facebook::yoga::vanillajni::registerNatives(
       env,
       "com/facebook/yoga/YogaNative",
       methods,

java/jni/YGJTypesVanilla.h

@@ -0,0 +1,41 @@
+/*
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * This source code is licensed under the MIT license found in the LICENSE
+ * file in the root directory of this source tree.
+ */
+#include "jni.h"
+#include <yoga/YGValue.h>
+#include <yoga/Yoga.h>
+#include <map>
+#include "common.h"
+
+using namespace facebook::yoga::vanillajni;
+using namespace std;
+
+class PtrJNodeMap {
+  std::map<YGNodeRef, size_t> ptrsToIdxs_;
+  jobjectArray javaNodes_;
+
+public:
+  PtrJNodeMap() : ptrsToIdxs_{}, javaNodes_{} {}
+  PtrJNodeMap(
+      jlong* nativePointers,
+      size_t nativePointersSize,
+      jobjectArray javaNodes)
+      : javaNodes_{javaNodes} {
+    for (size_t i = 0; i < nativePointersSize; ++i) {
+      ptrsToIdxs_[(YGNodeRef) nativePointers[i]] = i;
+    }
+  }
+
+  ScopedLocalRef<jobject> ref(JNIEnv* env, YGNodeRef node) {
+    auto idx = ptrsToIdxs_.find(node);
+    if (idx == ptrsToIdxs_.end()) {
+      return ScopedLocalRef<jobject>(env);
+    } else {
+      return make_local_ref(
+          env, env->GetObjectArrayElement(javaNodes_, idx->second));
+    }
+  }
+};