Leak aware resetting root ViewController.

MRCAppDelegate.m

@implementation MRCAppDelegate
  
- (void)p_updateWindowWithRootViewController:(nullable UIViewController *)controller
                             completionBlock:(nullable dispatch_block_t)completionBlock {
    if (!self.window) {
        self.window = [[UIWindow alloc] initWithFrame:[[UIScreen mainScreen] bounds]];
        [self.window makeKeyAndVisible];
    }
    val rootViewController = self.window.rootViewController;
    val resetRootViewControllerBlock = ^{
        [self p_setRootViewController:controller completionBlock:completionBlock];
    };
    if (rootViewController) {
        [rootViewController mrc_dismissAllAnimated:NO completion:resetRootViewControllerBlock];
    } else {
        resetRootViewControllerBlock();
    }
}

- (void)p_setRootViewController:(UIViewController *)actualController
                completionBlock:(nullable dispatch_block_t)completionBlock {
    __weak UIViewController *previousController = self.window.rootViewController;
    [self.window setRootViewController:actualController];
    if (previousController) {
        NSString *previousControllerClassName = NSStringFromClass([previousController class]);
        // Sometimes controller deallocates after 2-3 run loops, looks like UIKit detaches it
        // from internal data structes in async blocks...
        [previousController mrcapp_enusureDeallocWithRunLoopLimit:64 completionBlock:^(BOOL deallocated) {
            NSParameterAssert(deallocated);
            if (!deallocated) {
                NSLog(@"[UI] %@ leaked, presents: %@ {beingPresented:%@, beingDismissed:%@}",
                        previousController,
                        previousController.presentedViewController,
                        @(previousController.presentedViewController.beingPresented),
                        @(previousController.presentedViewController.beingDismissed));
            }
            if (completionBlock) {
                completionBlock();
            }
        }];
    } else if (completionBlock) {
        completionBlock();
    }
}

@end

//--  
  
@implementation UIViewController (MRCApp)
  
- (void)mrc_dismissAllAnimated:(BOOL)animated completion:(void(^)(void))completion {
    val transitionCoordinator = self.presentedViewController.transitionCoordinator;
    if (transitionCoordinator) {
        [transitionCoordinator animateAlongsideTransition:nil completion:^(id _) {
            val presentedViewController = self.presentedViewController;
            if (presentedViewController && !presentedViewController.beingDismissed) {
                [self dismissViewControllerAnimated:animated completion:completion];
            } else {
                completion();
            }
        }];
    } else if (self.presentedViewController) {
        [self dismissViewControllerAnimated:animated completion:completion];
    } else {
        completion();
    }
}

@end

//--
  
@implementation NSObject (MRCApp)

- (void)mrcapp_enusureDeallocWithRunLoopLimit:(NSInteger)runLoopLimit
                              completionBlock:(MRCObjectDeallocCompletionBlock)completionBlock {
    POS_CHECK(completionBlock);
    [self p_enusureTerminationWithRunLoopLimit:runLoopLimit runLoopCount:1 completionBlock:completionBlock];
}

- (void)p_enusureTerminationWithRunLoopLimit:(NSInteger)runLoopLimit
                                runLoopCount:(NSInteger)runLoopCount
                             completionBlock:(MRCObjectDeallocCompletionBlock)completionBlock {
    @weakify(self);
    dispatch_async(dispatch_get_main_queue(), ^{
        @strongify(self);
        if (!self) {
            completionBlock(YES);
        } else if (runLoopCount >= runLoopLimit) {
            completionBlock(NO);
        } else {
            [self p_enusureTerminationWithRunLoopLimit:runLoopLimit
                                          runLoopCount:runLoopCount + 1
                                       completionBlock:completionBlock];
        }
    });
}

@end