From c9a0ed385373a0e31e1ccaf1e838c7b7d51894d9 Mon Sep 17 00:00:00 2001
From: Michael Schneider <mischneider1@gmail.com>
Date: Tue, 18 Dec 2018 15:17:09 +0100
Subject: [PATCH] Further improvements

---
 Source/ASDisplayNode+Beta.h  |  4 ++--
 Source/ASDisplayNode+Yoga.mm | 26 ++++++++++++++++++++++----
 2 files changed, 24 insertions(+), 6 deletions(-)

diff --git a/Source/ASDisplayNode+Beta.h b/Source/ASDisplayNode+Beta.h
index f2e450d92..8bd825216 100644
--- a/Source/ASDisplayNode+Beta.h
+++ b/Source/ASDisplayNode+Beta.h
@@ -188,8 +188,7 @@ AS_EXTERN void ASDisplayNodePerformBlockOnEveryYogaChild(ASDisplayNode * _Nullab
 
 @interface ASDisplayNode (Yoga)
 
-// TODO: Make this and yogaCalculatedLayout atomic (lock).
-@property (nonatomic, copy) NSArray *yogaChildren;
+@property (copy) NSArray *yogaChildren;
 
 - (void)addYogaChild:(ASDisplayNode *)child;
 - (void)removeYogaChild:(ASDisplayNode *)child;
@@ -198,6 +197,7 @@ AS_EXTERN void ASDisplayNodePerformBlockOnEveryYogaChild(ASDisplayNode * _Nullab
 - (void)semanticContentAttributeDidChange:(UISemanticContentAttribute)attribute;
 
 @property BOOL yogaLayoutInProgress;
+// TODO: Make this atomic (lock).
 @property (nullable, nonatomic) ASLayout *yogaCalculatedLayout;
 
 // Will walk up the Yoga tree and returns the root node
diff --git a/Source/ASDisplayNode+Yoga.mm b/Source/ASDisplayNode+Yoga.mm
index 9211c2aba..c770cbd3a 100644
--- a/Source/ASDisplayNode+Yoga.mm
+++ b/Source/ASDisplayNode+Yoga.mm
@@ -48,32 +48,43 @@ - (void)setYogaChildren:(NSArray *)yogaChildren
   for (ASDisplayNode *child in [_yogaChildren copy]) {
     // Make sure to un-associate the YGNodeRef tree before replacing _yogaChildren
     // If this becomes a performance bottleneck, it can be optimized by not doing the NSArray removals here.
-    [self removeYogaChild:child];
+    [self _locked_removeYogaChild:child];
   }
   _yogaChildren = nil;
   for (ASDisplayNode *child in yogaChildren) {
-    [self addYogaChild:child];
+    [self _locked_addYogaChild:child];
   }
 }
 
 - (NSArray *)yogaChildren
 {
+  ASLockScope(self.yogaRoot);
   return [_yogaChildren copy] ?: @[];
 }
 
 - (void)addYogaChild:(ASDisplayNode *)child
 {
   ASLockScope(self.yogaRoot);
+  [self _locked_addYogaChild:child];
+}
+
+- (void)_locked_addYogaChild:(ASDisplayNode *)child
+{
   [self insertYogaChild:child atIndex:_yogaChildren.count];
 }
 
 - (void)removeYogaChild:(ASDisplayNode *)child
 {
   ASLockScope(self.yogaRoot);
+  [self _locked_removeYogaChild:child];
+}
+
+- (void)_locked_removeYogaChild:(ASDisplayNode *)child
+{
   if (child == nil) {
     return;
   }
-  
+
   [_yogaChildren removeObjectIdenticalTo:child];
 
   // YGNodeRef removal is done in setParent:
@@ -83,6 +94,11 @@ - (void)removeYogaChild:(ASDisplayNode *)child
 - (void)insertYogaChild:(ASDisplayNode *)child atIndex:(NSUInteger)index
 {
   ASLockScope(self.yogaRoot);
+  [self _locked_insertYogaChild:child atIndex:index];
+}
+
+- (void)_locked_insertYogaChild:(ASDisplayNode *)child atIndex:(NSUInteger)index
+{
   if (child == nil) {
     return;
   }
@@ -91,7 +107,7 @@ - (void)insertYogaChild:(ASDisplayNode *)child atIndex:(NSUInteger)index
   }
 
   // Clean up state in case this child had another parent.
-  [self removeYogaChild:child];
+  [self _locked_removeYogaChild:child];
 
   [_yogaChildren insertObject:child atIndex:index];
 
@@ -99,6 +115,8 @@ - (void)insertYogaChild:(ASDisplayNode *)child atIndex:(NSUInteger)index
   child.yogaParent = self;
 }
 
+#pragma mark - Subclass Hooks
+
 - (void)semanticContentAttributeDidChange:(UISemanticContentAttribute)attribute
 {
   UIUserInterfaceLayoutDirection layoutDirection =