Reduce JNI allocations when importing native maps (#58274)

Summary:

`ReadableNativeMap` materialization copied keys and created temporary JNI references for every imported type. Cache pointers to the stable native values and reuse global `ReadableType` references so importing maps and arrays does less allocation and lookup work.

Changelog: [Internal]

Differential Revision: D118277119

Author: javache

packages/react-native/ReactAndroid/src/main/jni/react/jni/JDynamicNative.cpp

@@ -17,7 +17,7 @@ jboolean JDynamicNative::isNullNative() {
   return static_cast<jboolean>(payload_.isNull());
 }
 
-jni::local_ref<ReadableType> JDynamicNative::getTypeNative() {
+jni::alias_ref<ReadableType> JDynamicNative::getTypeNative() {
   return ReadableType::getType(payload_.type());
 }
 

packages/react-native/ReactAndroid/src/main/jni/react/jni/JDynamicNative.h

@@ -42,7 +42,7 @@ class JDynamicNative : public jni::HybridClass<JDynamicNative, JDynamic> {
  private:
   friend HybridBase;
 
-  jni::local_ref<ReadableType> getTypeNative();
+  jni::alias_ref<ReadableType> getTypeNative();
   jni::local_ref<jstring> asString();
   jboolean asBoolean();
   jdouble asDouble();

packages/react-native/ReactAndroid/src/main/jni/react/jni/NativeCommon.cpp

@@ -27,32 +27,32 @@ alias_ref<ReadableType> getTypeField(const char* fieldName) {
 
 } // namespace
 
-local_ref<ReadableType> ReadableType::getType(folly::dynamic::Type type) {
+alias_ref<ReadableType> ReadableType::getType(folly::dynamic::Type type) {
   switch (type) {
     case folly::dynamic::Type::NULLT: {
-      static alias_ref<ReadableType> val = getTypeField("Null");
-      return make_local(val);
+      static auto val = getTypeField("Null");
+      return val;
     }
     case folly::dynamic::Type::BOOL: {
-      static alias_ref<ReadableType> val = getTypeField("Boolean");
-      return make_local(val);
+      static auto val = getTypeField("Boolean");
+      return val;
     }
     case folly::dynamic::Type::DOUBLE:
     case folly::dynamic::Type::INT64: {
-      static alias_ref<ReadableType> val = getTypeField("Number");
-      return make_local(val);
+      static auto val = getTypeField("Number");
+      return val;
     }
     case folly::dynamic::Type::STRING: {
-      static alias_ref<ReadableType> val = getTypeField("String");
-      return make_local(val);
+      static auto val = getTypeField("String");
+      return val;
     }
     case folly::dynamic::Type::OBJECT: {
-      static alias_ref<ReadableType> val = getTypeField("Map");
-      return make_local(val);
+      static auto val = getTypeField("Map");
+      return val;
     }
     case folly::dynamic::Type::ARRAY: {
-      static alias_ref<ReadableType> val = getTypeField("Array");
-      return make_local(val);
+      static auto val = getTypeField("Array");
+      return val;
     }
     default:
       throwNewJavaException(

packages/react-native/ReactAndroid/src/main/jni/react/jni/NativeCommon.h

@@ -19,7 +19,7 @@ namespace facebook::react {
 struct ReadableType : public jni::JavaClass<ReadableType> {
   static auto constexpr kJavaDescriptor = "Lcom/facebook/react/bridge/ReadableType;";
 
-  static jni::local_ref<ReadableType> getType(folly::dynamic::Type type);
+  static jni::alias_ref<ReadableType> getType(folly::dynamic::Type type);
 };
 
 namespace exceptions {
@@ -29,7 +29,7 @@ extern const char *gUnexpectedNativeTypeExceptionClass;
 template <typename T>
 void throwIfObjectAlreadyConsumed(const T &t, const char *msg)
 {
-  if (t->isConsumed) {
+  if (t->isConsumed) [[unlikely]] {
     jni::throwNewJavaException("com/facebook/react/bridge/ObjectAlreadyConsumedException", msg);
   }
 }

packages/react-native/ReactAndroid/src/main/jni/react/jni/ReadableNativeArray.cpp

@@ -7,8 +7,6 @@
 
 #include "ReadableNativeArray.h"
 
-#include "ReadableNativeMap.h"
-
 using namespace facebook::jni;
 
 namespace facebook::react {
@@ -23,7 +21,7 @@ void ReadableNativeArray::mapException(std::exception_ptr ex) {
 }
 
 local_ref<JArrayClass<jobject>> ReadableNativeArray::importArray() {
-  auto size = static_cast<jint>(array_.size());
+  auto size = static_cast<jsize>(array_.size());
   auto jarray = JArrayClass<jobject>::newArray(size);
   for (jint ii = 0; ii < size; ii++) {
     addDynamicToJArray(jarray, ii, array_.at(ii));
@@ -32,10 +30,10 @@ local_ref<JArrayClass<jobject>> ReadableNativeArray::importArray() {
 }
 
 local_ref<JArrayClass<jobject>> ReadableNativeArray::importTypeArray() {
-  auto size = static_cast<jint>(array_.size());
+  auto size = static_cast<jsize>(array_.size());
   auto jarray = JArrayClass<jobject>::newArray(size);
   for (jint ii = 0; ii < size; ii++) {
-    (*jarray)[ii] = ReadableType::getType(array_.at(ii).type());
+    jarray->setElement(ii, ReadableType::getType(array_.at(ii).type()).get());
   }
   return jarray;
 }

packages/react-native/ReactAndroid/src/main/jni/react/jni/ReadableNativeArray.h

@@ -9,9 +9,6 @@
 
 #include "NativeArray.h"
 
-#include "NativeCommon.h"
-#include "NativeMap.h"
-
 namespace facebook::react {
 
 struct ReadableArray : jni::JavaClass<ReadableArray> {

packages/react-native/ReactAndroid/src/main/jni/react/jni/ReadableNativeMap.cpp

@@ -20,84 +20,86 @@ void ReadableNativeMap::mapException(std::exception_ptr ex) {
   }
 }
 
+void ReadableNativeMap::throwIfKeysNotImported() const {
+  if (!values_.has_value()) [[unlikely]] {
+    throwNewJavaException(
+        "java/lang/IllegalStateException",
+        "importKeys must be called before importing values or types");
+  }
+}
+
 void addDynamicToJArray(
-    local_ref<JArrayClass<jobject>> jarray,
+    alias_ref<JArrayClass<jobject>> jarray,
     jint index,
     const folly::dynamic& dyn) {
+  local_ref<jobject> value;
   switch (dyn.type()) {
-    case folly::dynamic::Type::NULLT: {
-      jarray->setElement(index, nullptr);
+    case folly::dynamic::Type::BOOL:
+      value = JBoolean::valueOf(static_cast<jboolean>(dyn.getBool()));
       break;
-    }
-    case folly::dynamic::Type::BOOL: {
-      (*jarray)[index] =
-          JBoolean::valueOf(static_cast<unsigned char>(dyn.getBool()));
+    case folly::dynamic::Type::INT64:
+      value = JDouble::valueOf(static_cast<double>(dyn.getInt()));
       break;
-    }
-    case folly::dynamic::Type::INT64: {
-      (*jarray)[index] = JDouble::valueOf(dyn.getInt());
+    case folly::dynamic::Type::DOUBLE:
+      value = JDouble::valueOf(dyn.getDouble());
       break;
-    }
-    case folly::dynamic::Type::DOUBLE: {
-      (*jarray)[index] = JDouble::valueOf(dyn.getDouble());
+    case folly::dynamic::Type::STRING:
+      value = make_jstring(dyn.getString());
       break;
-    }
-    case folly::dynamic::Type::STRING: {
-      (*jarray)[index] = make_jstring(dyn.getString());
+    case folly::dynamic::Type::OBJECT:
+      value = ReadableNativeMap::newObjectCxxArgs(dyn);
       break;
-    }
-    case folly::dynamic::Type::OBJECT: {
-      (*jarray)[index] = ReadableNativeMap::newObjectCxxArgs(dyn);
-      break;
-    }
-    case folly::dynamic::Type::ARRAY: {
-      (*jarray)[index] = ReadableNativeArray::newObjectCxxArgs(dyn);
+    case folly::dynamic::Type::ARRAY:
+      value = ReadableNativeArray::newObjectCxxArgs(dyn);
       break;
-    }
+    case folly::dynamic::Type::NULLT:
     default:
-      jarray->setElement(index, nullptr);
       break;
   }
+  jarray->setElement(index, value.get());
 }
 
 local_ref<JArrayClass<jstring>> ReadableNativeMap::importKeys() {
   throwIfConsumed();
 
-  keys_ = folly::dynamic::array();
-  if (map_ == nullptr) {
-    return JArrayClass<jstring>::newArray(0);
-  }
-  auto jarray = JArrayClass<jstring>::newArray(map_.size());
+  auto size = map_ == nullptr ? 0 : static_cast<jsize>(map_.size());
+  std::vector<const folly::dynamic*> values(size);
+
+  auto jarray = JArrayClass<jstring>::newArray(size);
   jint i = 0;
-  for (auto& pair : map_.items()) {
-    auto value = pair.first.asString();
-    (*keys_).push_back(value);
-    (*jarray)[i++] = make_jstring(value);
+  if (map_ != nullptr) {
+    for (auto& pair : map_.items()) {
+      values[i] = &pair.second;
+      jarray->setElement(i++, make_jstring(pair.first.getString()).get());
+    }
   }
+  values_ = std::move(values);
 
   return jarray;
 }
 
 local_ref<JArrayClass<jobject>> ReadableNativeMap::importValues() {
   throwIfConsumed();
+  throwIfKeysNotImported();
 
-  auto size = static_cast<jint>(keys_.value().size());
+  const auto& values = values_.value();
+  auto size = static_cast<jsize>(values.size());
   auto jarray = JArrayClass<jobject>::newArray(size);
   for (jint ii = 0; ii < size; ii++) {
-    const std::string& key = (*keys_)[ii].getString();
-    addDynamicToJArray(jarray, ii, map_.at(key));
+    addDynamicToJArray(jarray, ii, *values[ii]);
   }
   return jarray;
 }
 
 local_ref<JArrayClass<jobject>> ReadableNativeMap::importTypes() {
   throwIfConsumed();
+  throwIfKeysNotImported();
 
-  auto size = static_cast<jint>(keys_.value().size());
+  const auto& values = values_.value();
+  auto size = static_cast<jsize>(values.size());
   auto jarray = JArrayClass<jobject>::newArray(size);
   for (jint ii = 0; ii < size; ii++) {
-    const std::string& key = (*keys_)[ii].getString();
-    (*jarray)[ii] = ReadableType::getType(map_.at(key).type());
+    jarray->setElement(ii, ReadableType::getType(values[ii]->type()).get());
   }
   return jarray;
 }

packages/react-native/ReactAndroid/src/main/jni/react/jni/ReadableNativeMap.h

@@ -11,10 +11,9 @@
 #include <folly/dynamic.h>
 #include <folly/json.h>
 #include <optional>
+#include <vector>
 
-#include "NativeCommon.h"
 #include "NativeMap.h"
-#include "ReadableNativeArray.h"
 
 namespace facebook::react {
 
@@ -24,15 +23,14 @@ struct ReadableMap : jni::JavaClass<ReadableMap> {
   static auto constexpr kJavaDescriptor = "Lcom/facebook/react/bridge/ReadableMap;";
 };
 
-void addDynamicToJArray(jni::local_ref<jni::JArrayClass<jobject>> jarray, jint index, const folly::dynamic &dyn);
+void addDynamicToJArray(jni::alias_ref<jni::JArrayClass<jobject>> jarray, jint index, const folly::dynamic &dyn);
 
 struct ReadableNativeMap : jni::HybridClass<ReadableNativeMap, NativeMap> {
   static auto constexpr kJavaDescriptor = "Lcom/facebook/react/bridge/ReadableNativeMap;";
 
   jni::local_ref<jni::JArrayClass<jstring>> importKeys();
   jni::local_ref<jni::JArrayClass<jobject>> importValues();
   jni::local_ref<jni::JArrayClass<jobject>> importTypes();
-  std::optional<folly::dynamic> keys_;
   static jni::local_ref<jhybridobject> createWithContents(folly::dynamic &&map);
 
   static void mapException(std::exception_ptr ex);
@@ -41,6 +39,12 @@ struct ReadableNativeMap : jni::HybridClass<ReadableNativeMap, NativeMap> {
   using HybridBase::HybridBase;
   friend HybridBase;
   friend struct WritableNativeMap;
+
+ private:
+  void throwIfKeysNotImported() const;
+
+  // Mutations are rejected once these borrowed pointers have been populated.
+  std::optional<std::vector<const folly::dynamic *>> values_;
 };
 
 } // namespace facebook::react

packages/react-native/ReactAndroid/src/main/jni/react/jni/WritableNativeMap.cpp

@@ -24,28 +24,37 @@ void WritableNativeMap::initHybrid(alias_ref<jhybridobject> jobj) {
   setCxxInstance(jobj);
 }
 
-void WritableNativeMap::putNull(std::string key) {
+void WritableNativeMap::throwIfNotWritable() {
   throwIfConsumed();
+  if (values_.has_value()) [[unlikely]] {
+    throwNewJavaException(
+        "java/lang/IllegalStateException",
+        "Cannot modify a map after its keys have been imported");
+  }
+}
+
+void WritableNativeMap::putNull(std::string key) {
+  throwIfNotWritable();
   map_.insert(std::move(key), nullptr);
 }
 
 void WritableNativeMap::putBoolean(std::string key, bool val) {
-  throwIfConsumed();
+  throwIfNotWritable();
   map_.insert(std::move(key), val);
 }
 
 void WritableNativeMap::putDouble(std::string key, double val) {
-  throwIfConsumed();
+  throwIfNotWritable();
   map_.insert(std::move(key), val);
 }
 
 void WritableNativeMap::putInt(std::string key, int val) {
-  throwIfConsumed();
+  throwIfNotWritable();
   map_.insert(std::move(key), val);
 }
 
 void WritableNativeMap::putLong(std::string key, jlong val) {
-  throwIfConsumed();
+  throwIfNotWritable();
   map_.insert(std::move(key), val);
 }
 
@@ -54,7 +63,7 @@ void WritableNativeMap::putString(std::string key, alias_ref<jstring> val) {
     putNull(std::move(key));
     return;
   }
-  throwIfConsumed();
+  throwIfNotWritable();
   map_.insert(std::move(key), val->toString());
 }
 
@@ -65,7 +74,7 @@ void WritableNativeMap::putNativeArray(
     putNull(std::move(key));
     return;
   }
-  throwIfConsumed();
+  throwIfNotWritable();
   map_.insert(key, otherArray->consume());
 }
 
@@ -76,12 +85,12 @@ void WritableNativeMap::putNativeMap(
     putNull(std::move(key));
     return;
   }
-  throwIfConsumed();
+  throwIfNotWritable();
   map_.insert(std::move(key), otherMap->consume());
 }
 
 void WritableNativeMap::mergeNativeMap(ReadableNativeMap* other) {
-  throwIfConsumed();
+  throwIfNotWritable();
   other->throwIfConsumed();
 
   for (const auto& sourceIt : other->map_.items()) {

packages/react-native/ReactAndroid/src/main/jni/react/jni/WritableNativeMap.h

@@ -41,6 +41,9 @@ struct WritableNativeMap : jni::HybridClass<WritableNativeMap, ReadableNativeMap
   static void registerNatives();
 
   friend HybridBase;
+
+ private:
+  void throwIfNotWritable();
 };
 
 } // namespace facebook::react