Handling navigation with UINavigationController


I feel like the only one using TCA with UIKit, but perhaps there are others out there. And if there are they've probably also run across the issue of handling popping view controllers in a navigation controller. Presumably the pushed view controller is tied to some optional state somewhere and whenever the view controller is popped you want to nil that out.

I wrote this UINavigationController subclass that handles that automatically by attaching an action to a pushed view controller. I don't know if this is a good idea though. So I'm posting here in hopes of getting some code review or other input on the topic!

import UIKit

open class NavigationController<Action>: UINavigationController {

  private var popActions: [UIViewController: Action] = [:]

  public var popActionHandler: ((Action) -> Void)?

  public func pushViewController(_ viewController: UIViewController, animated: Bool, withPopAction popAction: Action) {
    popActions[viewController] = popAction
    super.pushViewController(viewController, animated: true)

  open override func popViewController(animated: Bool) -> UIViewController? {
    let popped = super.popViewController(animated: animated)
    return popped

  open override func popToRootViewController(animated: Bool) -> [UIViewController]? {
    let popped = super.popToRootViewController(animated: animated)
    return popped

  open override func popToViewController(_ viewController: UIViewController, animated: Bool) -> [UIViewController]? {
    let popped = super.popToViewController(viewController, animated: animated)
    return popped

  private func performActionAndClear(for viewController: UIViewController) {
    guard let action = popActions[viewController] else { return }
    popActions.removeValue(forKey: viewController)

  private func performActionAndClear(for viewControllers: [UIViewController]) {


It's pretty straight forward, keep track of pushed view controllers and their associated action. Whenever we pop one or multiple (I think this should handle the new tap-and-hold-the-back-button navigation in iOS 14 but haven't double checked yet) view controllers we call the popActionHandler with that associated action.

So in practice we would have a an instance of this navigation controller setup like this:

let myNavigationController = NavigationController<MyFeatureAction>(rootViewController: featureRootViewController)
myNavigationController.popActionHandler = { [weak self] action in self?.viewStore.send(action) }

And then whenever we push a view controller we also pass along the associated pop-action:

      .scope(state: { $0.recipeCollectionViewState }, action: PlanFeatureAction.recipeCollectionViewAction)
      .ifLet { [weak self] scoped in
        let recipeCollectionViewController = RecipeCollectionView(store: scoped)
        self?.pushViewController(recipeCollectionViewController, animated: true, withPopAction: .didPopRecipeCollectionView)
      } else: { [weak self] in
        guard let self = self else { return }
        self.navigationController?.popToViewController(self, animated: true)
      .store(in: &cancellables)