diff --git a/Source/ASDisplayNode+Layout.mm b/Source/ASDisplayNode+Layout.mm
index 0eac1e3ab3..35b6708283 100644
--- a/Source/ASDisplayNode+Layout.mm
+++ b/Source/ASDisplayNode+Layout.mm
@@ -1046,10 +1046,7 @@ ASLayoutElementStyleExtensibilityForwarding
 #pragma mark -
 #pragma mark - ASDisplayNode (YogaLayout)
 
-@implementation ASDisplayNode (YogaInternal)
-
-#pragma mark -
-#pragma mark - ASDisplayNode (Yoga)
+@implementation ASDisplayNode (YogaLayout)
 
 - (BOOL)locked_shouldLayoutFromYogaRoot {
 #if YOGA
@@ -1071,32 +1068,4 @@ ASLayoutElementStyleExtensibilityForwarding
 #endif
 }
 
-- (ASLockSet)lockToRootIfNeededForLayout {
-  ASLockSet lockSet = ASLockSequence(^BOOL(ASAddLockBlock addLock) {
-    if (!addLock(self)) {
-      return NO;
-    }
-#if YOGA
-    if (![self locked_shouldLayoutFromYogaRoot]) {
-      return YES;
-    }
-    if (self.nodeController && !addLock(self.nodeController)) {
-      return NO;
-    }
-    ASDisplayNode *parent = _supernode;
-    while (parent) {
-      if (!addLock(parent)) {
-        return NO;
-      }
-      if (parent.nodeController && !addLock(parent.nodeController)) {
-        return NO;
-      }
-      parent = parent->_supernode;
-    }
-#endif
-    return true;
-  });
-  return lockSet;
-}
-
 @end
diff --git a/Source/ASDisplayNode+Yoga.h b/Source/ASDisplayNode+Yoga.h
index fc8d884642..93cbf7a7c1 100644
--- a/Source/ASDisplayNode+Yoga.h
+++ b/Source/ASDisplayNode+Yoga.h
@@ -33,6 +33,10 @@ AS_EXTERN void ASDisplayNodePerformBlockOnEveryYogaChild(ASDisplayNode * _Nullab
 // Will walk up the Yoga tree and returns the root node
 - (ASDisplayNode *)yogaRoot;
 
+
+@end
+
+@interface ASDisplayNode (YogaLocking)
 /**
  * @discussion Attempts(spinning) to lock all node up to root node when yoga is enabled.
  * This will lock self when yoga is not enabled;
diff --git a/Source/ASDisplayNode+Yoga.mm b/Source/ASDisplayNode+Yoga.mm
index 0f975c8854..97e7cb56b1 100644
--- a/Source/ASDisplayNode+Yoga.mm
+++ b/Source/ASDisplayNode+Yoga.mm
@@ -20,6 +20,7 @@
 #import <AsyncDisplayKit/ASDisplayNodeInternal.h>
 #import <AsyncDisplayKit/ASLayout.h>
 #import <AsyncDisplayKit/ASLayoutElementStylePrivate.h>
+#import <AsyncDisplayKit/ASNodeController+Beta.h>
 
 #import <AsyncDisplayKit/ASDisplayNode+LayoutSpec.h>
 
@@ -27,7 +28,7 @@
 
 #pragma mark - ASDisplayNode+Yoga
 
-@interface ASDisplayNode (YogaInternal)
+@interface ASDisplayNode (YogaPrivate)
 @property (nonatomic, weak) ASDisplayNode *yogaParent;
 - (ASSizeRange)_locked_constrainedSizeForLayoutPass;
 @end
@@ -409,6 +410,40 @@
 
 @end
 
+#pragma mark - ASDisplayNode (YogaLocking)
+
+@implementation ASDisplayNode (YogaLocking)
+
+- (ASLockSet)lockToRootIfNeededForLayout {
+  ASLockSet lockSet = ASLockSequence(^BOOL(ASAddLockBlock addLock) {
+    if (!addLock(self)) {
+      return NO;
+    }
+#if YOGA
+    if (![self locked_shouldLayoutFromYogaRoot]) {
+      return YES;
+    }
+    if (self.nodeController && !addLock(self.nodeController)) {
+      return NO;
+    }
+    ASDisplayNode *parent = _supernode;
+    while (parent) {
+      if (!addLock(parent)) {
+        return NO;
+      }
+      if (parent.nodeController && !addLock(parent.nodeController)) {
+        return NO;
+      }
+      parent = parent->_supernode;
+    }
+#endif
+    return true;
+  });
+  return lockSet;
+}
+
+@end
+
 @implementation ASDisplayNode (YogaDebugging)
 
 - (NSString *)yogaTreeDescription {