diff --git a/java/jni/LayoutContext.cpp b/java/jni/LayoutContext.cpp
new file mode 100644
index 0000000000..db5812bfab
--- /dev/null
+++ b/java/jni/LayoutContext.cpp
@@ -0,0 +1,14 @@
+/*
+ * 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 "LayoutContext.h"
+
+namespace facebook::yoga::vanillajni {
+
+thread_local std::stack<PtrJNodeMapVanilla*> LayoutContext::contexts_;
+
+} // namespace facebook::yoga::vanillajni
diff --git a/java/jni/LayoutContext.h b/java/jni/LayoutContext.h
new file mode 100644
index 0000000000..9c6041f330
--- /dev/null
+++ b/java/jni/LayoutContext.h
@@ -0,0 +1,32 @@
+/*
+ * 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 <stack>
+#include "YGJTypesVanilla.h"
+
+namespace facebook::yoga::vanillajni {
+
+class LayoutContext {
+public:
+  struct Holder {
+    explicit Holder(PtrJNodeMapVanilla* data) {
+      LayoutContext::contexts_.push(data);
+    }
+    ~Holder() { LayoutContext::contexts_.pop(); }
+  };
+
+  static PtrJNodeMapVanilla* getCurrent() {
+    return contexts_.empty() ? nullptr : contexts_.top();
+  }
+
+private:
+  static thread_local std::stack<PtrJNodeMapVanilla*> contexts_;
+};
+
+} // namespace facebook::yoga::vanillajni
diff --git a/java/jni/YGJNIVanilla.cpp b/java/jni/YGJNIVanilla.cpp
index 86f740e1f5..aa855cd7c3 100644
--- a/java/jni/YGJNIVanilla.cpp
+++ b/java/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::getCurrent()->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::Holder holder(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(
diff --git a/java/jni/YGJTypesVanilla.h b/java/jni/YGJTypesVanilla.h
index e5c44c0cf6..705e1eda6f 100644
--- a/java/jni/YGJTypesVanilla.h
+++ b/java/jni/YGJTypesVanilla.h
@@ -5,6 +5,8 @@
  * LICENSE file in the root directory of this source tree.
  */
 
+#pragma once
+
 #include <map>
 #include <vector>