Setup for back/dismiss buttons in VCs, useful for MVVM+C pattern to finish coordinators or don't care about navigation controller and it's dismiss button (whether it's a back button or dismiss or w/e).

ClosureBarButtonItem.swift

final class ClosureBarButtonItem: UIBarButtonItem {

    private var customAction: (() -> Void)?

    convenience init(barButtonSystemItem: UIBarButtonItem.SystemItem, action: (() -> Void)?) {
        self.init(barButtonSystemItem: barButtonSystemItem, target: nil, action: nil)
        self.target = self
        self.action = #selector(didTap)
        self.customAction = action
    }

    @objc private func didTap() {
        customAction?()
    }
}

DismissButton.swift

@objc enum DismissButton: Int {
    case backButton // default for vcs pushed on the stack
    case xButton // default for modally presented screens
    case none // if you want to hide the back/dismiss button
}

DismissButtonProviding.swift

protocol DismissButtonProviding {
    var dismissButton: DismissButton { get }
    var shouldDismiss: () -> Bool { get }

    func setupDismissButton()
    func dismiss()
}

extension UIViewController: DismissButtonProviding {
    @objc var dismissButton: DismissButton {
        if isBeingPresented || (navigationController?.isBeingPresented == true && navigationController?.viewControllers.count == 1) {
            return .xButton
        } else {
            return .backButton
        }
    }

    @objc func setupDismissButton() {
        switch dismissButton {
        case .none:
            navigationItem.hidesBackButton = true
        case .xButton:
            navigationItem.leftBarButtonItem = ActionBarButtonItem(barButtonSystemItem: .stop, action: Action { [weak self] in
                guard let self = self else { return }

                if self.shouldDismiss() {
                    self.dismiss()
                }
            })
        case .backButton:
            break
        }
    }

    @objc var shouldDismiss: () -> Bool {
        return { return true }
    }

    @objc func dismiss() {
        dismissOrPop()
    }
}

DismissNavigationController.swift

final class DismissNavigationController: UINavigationController, UINavigationBarDelegate, UINavigationControllerDelegate {
    override init(nibName nibNameOrNil: String? = nil, bundle nibBundleOrNil: Bundle? = nil) {
        super.init(nibName: nibNameOrNil, bundle: nibBundleOrNil)

        setup()
    }

    required init?(coder aDecoder: NSCoder) {
        super.init(coder: aDecoder)

        setup()
    }

    override init(rootViewController: UIViewController) {
        super.init(rootViewController: rootViewController)

        setup()
    }

    private func setup() {
        delegate = self
    }

    func navigationBar(_ navigationBar: UINavigationBar, shouldPop item: UINavigationItem) -> Bool {
        guard let controller = topViewController as DismissButtonProviding? else { return true }

        return controller.shouldDismiss()
    }

    func navigationController(_ navigationController: UINavigationController, willShow viewController: UIViewController, animated: Bool) {
        viewController.setupDismissButton()
    }
}

UIViewController+Dismiss.swift

extension UIViewController {
    func dismissOrPop() {
        if presentingViewController != nil {
            dismiss(animated: true, completion: nil)
        } else if navigationController?.presentingViewController != nil, navigationController?.viewControllers.count == 1 {
            navigationController?.dismiss(animated: true, completion: nil)
        } else {
            navigationController?.popViewController(animated: true)
        }
    }
}