AOP-Style UIViewController lifecycle management

ViewControllerLifecycleBehavior.swift

import UIKit

protocol BehaviorableViewController {
    var behaviors: [ViewControllerLifecycleBehavior] { get }
}

protocol ViewControllerLifecycleBehavior {
    func afterLoading(_ viewController: UIViewController)

    func beforeAppearing(_ viewController: UIViewController)

    func afterAppearing(_ viewController: UIViewController)

    func beforeDisappearing(_ viewController: UIViewController)

    func afterDisappearing(_ viewController: UIViewController)

    func beforeLayingOutSubviews(_ viewController: UIViewController)

    func afterLayingOutSubviews(_ viewController: UIViewController)

    func preferredStatusBarStyle(_ viewController: UIViewController) -> UIStatusBarStyle?
}

extension ViewControllerLifecycleBehavior {
    func afterLoading(_ viewController: UIViewController) {}

    func beforeAppearing(_ viewController: UIViewController) {}

    func afterAppearing(_ viewController: UIViewController) {}

    func beforeDisappearing(_ viewController: UIViewController) {}

    func afterDisappearing(_ viewController: UIViewController) {}

    func beforeLayingOutSubviews(_ viewController: UIViewController) {}

    func afterLayingOutSubviews(_ viewController: UIViewController) {}

    func preferredStatusBarStyle(_ viewController: UIViewController) -> UIStatusBarStyle? { return nil }
}

private extension UIViewController {
    @objc
    func vcb_viewDidLoad() {
        vcb_viewDidLoad()

        applyBehaviors { behavior, viewController in
            behavior.afterLoading(viewController)
        }
    }

    @objc var vcb_preferredStatusBarStyle: UIStatusBarStyle {
        var statusBarStyle: UIStatusBarStyle?
        applyBehaviors { behavior, viewController in
            if let newStatusBarStyle = behavior.preferredStatusBarStyle(viewController) {
                assert(statusBarStyle == nil,
                       "More than one behavior which implements preferredStatusBarStyle(_:) are applied to the controller \(self)")
                statusBarStyle = newStatusBarStyle
            }
        }

        // what we got from one of the behaviors OR just return super's
        return statusBarStyle ?? self.vcb_preferredStatusBarStyle
    }

    @objc
    func vcb_viewWillAppear(_ animated: Bool) {
        vcb_viewWillAppear(animated)

        applyBehaviors { behavior, viewController in
            behavior.beforeAppearing(viewController)
        }
    }

    @objc
    func vcb_viewDidAppear(_ animated: Bool) {
        vcb_viewDidAppear(animated)

        applyBehaviors { behavior, viewController in
            behavior.afterAppearing(viewController)
        }
    }

    @objc
    func vcb_viewWillDisappear(_ animated: Bool) {
        vcb_viewWillDisappear(animated)

        applyBehaviors { behavior, viewController in
            behavior.beforeDisappearing(viewController)
        }
    }

    @objc
    func vcb_viewDidDisappear(_ animated: Bool) {
        vcb_viewDidDisappear(animated)

        applyBehaviors { behavior, viewController in
            behavior.afterDisappearing(viewController)
        }
    }

    @objc
    func vcb_viewWillLayoutSubviews() {
        vcb_viewWillLayoutSubviews()

        applyBehaviors { behavior, viewController in
            behavior.beforeLayingOutSubviews(viewController)
        }
    }

    @objc
    func vcb_viewDidLayoutSubviews() {
        vcb_viewDidLayoutSubviews()

        applyBehaviors { behavior, viewController in
            behavior.afterLayingOutSubviews(viewController)
        }
    }

    // MARK: - Private

    private func applyBehaviors(body: (_ behavior: ViewControllerLifecycleBehavior, _ viewController: UIViewController) -> Void) {
        guard let behaviorableViewController = self as? BehaviorableViewController else { return }

        let behaviors = behaviorableViewController.behaviors
        for behavior in behaviors {
            // since we are inside UIViewController's category it's safe to force cast
            // swiftlint:disable force_cast
            body(behavior, behaviorableViewController as! UIViewController)
            // swiftlint:enable force_cast
        }
    }
}

// MARK: Swizzling

private protocol Swizzlable: AnyObject {
    static func swizzle()
}

private class SwizzlingHelper {
    // use lazy propery as `dispatch_once()`
    private static let performSwizzlingOnce: Any? = {
        UIViewController.swizzle()
        return nil
    }()

    static func performSwizzling() {
        _ = SwizzlingHelper.performSwizzlingOnce
    }
}

// Since `initialize` class method is deprecated we have to hook on UIApplication's lifecycle
extension UIApplication {
    // swiftlint:disable override_in_extension

    open override var next: UIResponder? {
        // Called before applicationDidFinishLaunching
        SwizzlingHelper.performSwizzling()
        return super.next
    }

    // swiftlint:enable override_in_extension
}

extension UIViewController: Swizzlable {
    static func swizzle() {
        let cls = UIViewController.self
        guard cls === self else { return } // make sure this isn't a subclass

        Swizzle(forClass: cls,
                originalSelector: #selector(cls.viewDidLoad),
                alternateSelector: #selector(cls.vcb_viewDidLoad))

        Swizzle(forClass: cls,
                originalSelector: #selector(getter: cls.preferredStatusBarStyle),
                alternateSelector: #selector(getter: cls.vcb_preferredStatusBarStyle))

        Swizzle(forClass: cls,
                originalSelector: #selector(cls.viewWillAppear(_:)),
                alternateSelector: #selector(cls.vcb_viewWillAppear(_:)))

        Swizzle(forClass: cls,
                originalSelector: #selector(cls.viewDidAppear(_:)),
                alternateSelector: #selector(cls.vcb_viewDidAppear(_:)))

        Swizzle(forClass: cls,
                originalSelector: #selector(cls.viewWillDisappear(_:)),
                alternateSelector: #selector(cls.vcb_viewWillDisappear(_:)))

        Swizzle(forClass: cls,
                originalSelector: #selector(cls.viewDidDisappear(_:)),
                alternateSelector: #selector(cls.vcb_viewDidDisappear(_:)))

        Swizzle(forClass: cls,
                originalSelector: #selector(cls.viewWillLayoutSubviews),
                alternateSelector: #selector(cls.vcb_viewWillLayoutSubviews))

        Swizzle(forClass: cls,
                originalSelector: #selector(cls.viewDidLayoutSubviews),
                alternateSelector: #selector(cls.vcb_viewDidLayoutSubviews))
    }
}

private func Swizzle(forClass cls: AnyClass, originalSelector: Selector, alternateSelector: Selector) {
    guard let originalMethod = class_getInstanceMethod(cls, originalSelector) else {
        fatalError("Original method \(originalSelector) not found for class \(cls)")
    }

    guard let alternateMethod = class_getInstanceMethod(cls, alternateSelector) else {
        fatalError("Alternate method \(alternateSelector) not found for class \(cls)")
    }

    guard let originalImpl = class_getMethodImplementation(cls, originalSelector) else {
        fatalError("Original implementation for \(originalSelector) not found for class \(cls)")
    }

    guard let alternateImpl = class_getMethodImplementation(cls, alternateSelector) else {
        fatalError("Alternate implementation for \(alternateSelector) not found for class \(cls)")
    }

    class_addMethod(cls, originalSelector, originalImpl, method_getTypeEncoding(originalMethod))
    class_addMethod(cls, alternateSelector, alternateImpl, method_getTypeEncoding(alternateMethod))

    method_exchangeImplementations(originalMethod, alternateMethod)
}