UIViewController+LifecyclePublishers.swift

import Combine
import ObjectiveC
import UIKit

private var eventSubjectKey = "viewcontroller_lifecycle_subject"
private var swizzlingPerformed = false

private final class LifecycleSubjectBag: NSObject {
    let viewDidLoadSubject = PassthroughSubject<Void, Never>()
    let viewWillAppearSubject = PassthroughSubject<Bool, Never>()
    let viewDidAppearSubject = PassthroughSubject<Bool, Never>()
    let viewWillDisappearSubject = PassthroughSubject<Bool, Never>()
    let viewDidDisappearSubject = PassthroughSubject<Bool, Never>()
}

extension UIViewController {
    func viewDidLoadPublisher() -> AnyPublisher<Void, Never> {
        objc_sync_enter(self)
        defer { objc_sync_exit(self) }
        objc_sync_enter(UIViewController.self)
        defer { objc_sync_exit(UIViewController.self) }

        let subjects = getOrCreateSubjects()
        return subjects.viewDidLoadSubject.eraseToAnyPublisher()
    }

    func viewWillAppearPublisher() -> AnyPublisher<Bool, Never> {
        objc_sync_enter(self)
        defer { objc_sync_exit(self) }
        objc_sync_enter(UIViewController.self)
        defer { objc_sync_exit(UIViewController.self) }

        let subjects = getOrCreateSubjects()
        return subjects.viewWillAppearSubject.eraseToAnyPublisher()
    }

    func viewDidAppearPublisher() -> AnyPublisher<Bool, Never> {
        objc_sync_enter(self)
        defer { objc_sync_exit(self) }
        objc_sync_enter(UIViewController.self)
        defer { objc_sync_exit(UIViewController.self) }

        let subjects = getOrCreateSubjects()
        return subjects.viewDidAppearSubject.eraseToAnyPublisher()
    }

    func viewWillDisappearPublisher() -> AnyPublisher<Bool, Never> {
        objc_sync_enter(self)
        defer { objc_sync_exit(self) }
        objc_sync_enter(UIViewController.self)
        defer { objc_sync_exit(UIViewController.self) }

        let subjects = getOrCreateSubjects()
        return subjects.viewWillDisappearSubject.eraseToAnyPublisher()
    }

    func viewDidDisappearPublisher() -> AnyPublisher<Bool, Never> {
        objc_sync_enter(self)
        defer { objc_sync_exit(self) }
        objc_sync_enter(UIViewController.self)
        defer { objc_sync_exit(UIViewController.self) }

        let subjects = getOrCreateSubjects()
        return subjects.viewDidDisappearSubject.eraseToAnyPublisher()
    }

    private func getOrCreateSubjects() -> LifecycleSubjectBag {
        if let subjects = objc_getAssociatedObject(self, &eventSubjectKey) as? LifecycleSubjectBag {
            return subjects
        }

        swizzleMethodsIfNeeded()

        let subjects = LifecycleSubjectBag()

        objc_setAssociatedObject(
            self,
            &eventSubjectKey,
            subjects,
            .OBJC_ASSOCIATION_RETAIN_NONATOMIC)

        return subjects
    }

    private func swizzleMethodsIfNeeded() {
        guard !swizzlingPerformed else {
            return
        }

        swizzlingPerformed = true

        let pairs: [(Selector, Selector)] = [
            (#selector(UIViewController.viewDidLoad), #selector(UIViewController._swizzled_viewDidLoad)),
            (#selector(UIViewController.viewWillAppear), #selector(UIViewController._swizzled_viewWillAppear)),
            (#selector(UIViewController.viewDidAppear), #selector(UIViewController._swizzled_viewDidAppear)),
            (#selector(UIViewController.viewWillDisappear), #selector(UIViewController._swizzled_viewWillDisappear)),
            (#selector(UIViewController.viewDidDisappear), #selector(UIViewController._swizzled_viewDidDisappear))
        ]

        pairs.forEach { originalSelector, swizzledSelector in
            guard let originalMethod = class_getInstanceMethod(UIViewController.self, originalSelector),
                  let swizzledMethod = class_getInstanceMethod(UIViewController.self, swizzledSelector) else {
                return
            }

            method_exchangeImplementations(originalMethod, swizzledMethod)
        }
    }

    @objc private dynamic func _swizzled_viewDidLoad() {
        _swizzled_viewDidLoad()

        guard let subjects = objc_getAssociatedObject(self, &eventSubjectKey) as? LifecycleSubjectBag else {
            return
        }

        subjects.viewDidLoadSubject.send()
    }

    @objc private dynamic func _swizzled_viewWillAppear(_ animated: Bool) {
        _swizzled_viewWillAppear(animated)

        guard let subjects = objc_getAssociatedObject(self, &eventSubjectKey) as? LifecycleSubjectBag else {
            return
        }

        subjects.viewWillAppearSubject.send(animated)
    }

    @objc private dynamic func _swizzled_viewDidAppear(_ animated: Bool) {
        _swizzled_viewDidAppear(animated)

        guard let subjects = objc_getAssociatedObject(self, &eventSubjectKey) as? LifecycleSubjectBag else {
            return
        }

        subjects.viewDidAppearSubject.send(animated)
    }

    @objc private dynamic func _swizzled_viewWillDisappear(_ animated: Bool) {
        _swizzled_viewWillDisappear(animated)

        guard let subjects = objc_getAssociatedObject(self, &eventSubjectKey) as? LifecycleSubjectBag else {
            return
        }

        subjects.viewWillDisappearSubject.send(animated)
    }

    @objc private dynamic func _swizzled_viewDidDisappear(_ animated: Bool) {
        _swizzled_viewDidDisappear(animated)

        guard let subjects = objc_getAssociatedObject(self, &eventSubjectKey) as? LifecycleSubjectBag else {
            return
        }

        subjects.viewDidDisappearSubject.send(animated)
    }
}