Fixed _nodesOnCurrentUndoBranch not always reloaded to contain curren…

…t node when handling COUndoTrackDidChangeNotification (#67)

Tests/Undo/TestUndoTrack.m

@@ -22,6 +22,7 @@ @interface TestUndoTrack : EditingContextTestCase <UKTest>
     NSUInteger _trackNotificationCount;
     NSUInteger _track2NotificationCount;
     NSUInteger _patternTrackNotificationCount;
+    void (^_notificationBlock)(NSNotification *);
 }
 
 @end
@@ -67,6 +68,8 @@ - (void)trackDidChange: (NSNotification *)notif
     {
         ETAssertUnreachable();
     }
+
+    _notificationBlock(notif);
 }
 
 /**
@@ -136,6 +139,9 @@ - (instancetype)init
                                              selector: @selector(trackDidChange:)
                                                  name: COUndoTrackDidChangeNotification
                                                object: _patternTrack];
+
+    _notificationBlock = ^(NSNotification *notif) { };
+
     return self;
 }
 
@@ -299,6 +305,27 @@ - (void)testSetCurrentNodeToDivergentNode
     UKObjectsEqual(A(placeholderNode, group1, group1a), _track.nodes);
 }
 
+
+- (void)testValidNodesOnNotificationForSetCurrentNodeToDivergentNode
+{
+    COCommandGroup *group1 = [[COCommandGroup alloc] init];
+    COCommandGroup *group1a = [[COCommandGroup alloc] init];
+    COCommandGroup *group1b = [[COCommandGroup alloc] init];
+
+    [_track recordCommand: group1];
+    [_track recordCommand: group1a];
+    [_track setCurrentNode: group1];
+    [_track recordCommand: group1b];
+
+    COUndoTrack *track = _track;
+
+    _notificationBlock = ^(NSNotification *notif)
+    {
+        UKDoesNotRaiseException((void)track.canRedo);
+    };
+    _track.currentNode = group1a;
+}
+
 - (void)testDivergentNodesWhereCommonAncestorIsPlaceholderNode
 {
     COCommandGroup *group1a = [[COCommandGroup alloc] init];

Tests/Undo/TestUndoTrackStore.m

@@ -7,6 +7,10 @@
 
 #import "TestCommon.h"
 
+@interface COUndoTrackStore (TestUndoTrackStore)
+- (BOOL)commitTransaction;
+@end
+
 @interface TestUndoTrackStore : NSObject <UKTest>
 {
     COUndoTrackStore *_store;
@@ -188,3 +192,14 @@ - (void)testDivergent
 }
 
 @end
+
+
+@implementation COUndoTrackStore (TestUndoTrackStore)
+
+- (BOOL)commitTransaction
+{
+    return [self commitTransactionWithCompletionHandler: ^() { }];
+}
+
+@end
+

Undo/COUndoTrack.m

@@ -203,8 +203,7 @@ - (BOOL)setCurrentNode: (id <COTrackNode>)node
 
     [self undo: undo1 redo: redo1 undo: @[] redo: @[]];
 
-    BOOL ok = [self.store commitTransaction];
-    if (ok)
+    return [self.store commitTransactionWithCompletionHandler: ^()
     {
         if ([self isKindOfClass: [COPatternUndoTrack class]])
         {
@@ -229,8 +228,7 @@ command between X and Y (D sequence number is higher than C).
         {
             [self didUpdate];
         }
-    }
-    return ok;
+    }];
 }
 
 - (NSArray *)nodesFromNode: (id <COTrackNode>)node toTargetNode: (id <COTrackNode>)targetNode
@@ -290,14 +288,12 @@ - (BOOL)setCurrentNodeToDivergentNode: (id <COTrackNode>)node
 
     [self undo: undo1 redo: redo1 undo: @[] redo: redo2];
 
-    BOOL ok = [self.store commitTransaction];
-    if (ok)
+    return [self.store commitTransactionWithCompletionHandler: ^()
     {
         /* When we set the current command to a divergent one, we switch to
            another command branch (the head command changes) */
         [self reloadNodesOnCurrentBranch];
-    }
-    return ok;
+    }];
 }
 
 - (void)undoNode: (id <COTrackNode>)aNode
@@ -430,15 +426,17 @@ - (void)recordCommand: (COCommandGroup *)aCommand
         [_commandsByUUID removeObjectForKey: coalescedCommandUUIDToDelete];
     }
 
-    ETAssert([self.store commitTransaction]);
-    [self reloadNodesOnCurrentBranch];
+    ETAssert([self.store commitTransactionWithCompletionHandler: ^()
+    {
+        [self reloadNodesOnCurrentBranch];
+    }]);
 }
 
 - (void)clear
 {
     [self.store beginTransaction];
     [self.store removeTrackWithName: _name];
-    ETAssert([self.store commitTransaction]);
+    ETAssert([self.store commitTransactionWithCompletionHandler: ^() { }]);
 
     _nodesOnCurrentUndoBranch = nil;
     [_commandsByUUID removeAllObjects];

Undo/COUndoTrackStore+Private.h

@@ -97,9 +97,12 @@ NSString *const COUndoTrackStoreTrackCompacted;
  *
  * Returns whether committing the transaction has succeeded.
  *
+ * On success, will execute the completion block in the main thread prior 
+ * to posting COUndoTrackStoreTrackDidChangeNotification.
+ *
  * This method must be run in the main thread.
  */
-- (BOOL)commitTransaction;
+- (BOOL)commitTransactionWithCompletionHandler: (void (^)())completion;
 
 
 /** @taskunit Managing Undo Tracks */

Undo/COUndoTrackStore.m

@@ -314,7 +314,7 @@ - (BOOL)beginTransaction
     return ok;
 }
 
-- (BOOL)commitTransaction
+- (BOOL)commitTransactionWithCompletionHandler: (void (^)())completion
 {
     ETAssert([NSThread isMainThread]);
     __block BOOL ok = NO;
@@ -327,6 +327,7 @@ - (BOOL)commitTransaction
 
     if (ok)
     {
+        completion();
         [self postCommitNotifications];
     }
     dispatch_semaphore_signal(_transactionLock);