Swizzling in Swift

Deallocation.swift

// MARK: - check dealloc

private extension UIViewController {
    // Add below to _swizzled_viewDidDisappear(_:)
    func _checkDeallocation(afterDelay delay: TimeInterval = 1.5) {
        if isMovingFromParent || rootParent.isBeingDismissed {
            let disappearance = isMovingFromParent ? "removed from its parent" : "dismissed"
            DispatchQueue.main.asyncAfter(deadline: .now() + delay) { [weak self] in
                guard let self = self else { return }
                let message = "\(String(describing: type(of: self))) not deallocated after being \(disappearance)"
                print("###########################################################################")
                print("[Swizzled] -> [\(message)]")
                print("###########################################################################")
            }
        }
    }

    private var rootParent: UIViewController {
        var root = self
        while let parent = root.parent {
            root = parent
        }
        return root
    }
}

Swizzle.swift

import UIKit

extension UIViewController {
    static var isSwizzling: Bool = false

    class func swizzle() {
        guard !isSwizzling else { return }

        let methods: [Selector] = [
            #selector(UIViewController.viewDidLoad),
            #selector(UIViewController.viewWillAppear(_:)),
            #selector(UIViewController.viewDidDisappear(_:)),
        ]
        methods.forEach { (sel) in
            self._swizzleMethod(sel)
        }
        isSwizzling = true
    }

    private class func _swizzleMethod(_ selector: Selector) {
        let selectorString = NSStringFromSelector(selector)
        _swizzleMethod(selectorString)
    }

    private class func _swizzleMethod(_ selector: String) {
        let originalSelector = NSSelectorFromString(selector)
        let swizzledSelector = NSSelectorFromString("_swizzled_" + selector)

        guard
            let originalMethod = class_getInstanceMethod(self, originalSelector),
            let swizzledMethod = class_getInstanceMethod(self, swizzledSelector)
            else { return }

        if class_addMethod(self, originalSelector, method_getImplementation(swizzledMethod), method_getTypeEncoding(swizzledMethod)) {
            class_replaceMethod(self, swizzledSelector, method_getImplementation(originalMethod), method_getTypeEncoding(originalMethod))
        } else {
            method_exchangeImplementations(originalMethod, swizzledMethod)
        }
    }

    @objc private func _swizzled_viewDidLoad() {
        print("[Swizzled] -> [viewDidLoad : \(_className)]")
        _swizzled_viewDidLoad()
    }

    @objc private func _swizzled_viewWillAppear(_ animated: Bool) {
        print("[Swizzled] -> [viewWillAppear : \(_className)]")
        _swizzled_viewWillAppear(animated)
    }

    @objc private func _swizzled_viewDidDisappear(_ animated: Bool) {
        print("[Swizzled] -> [viewDidDisappear : \(_className)]")
        _swizzled_viewDidDisappear(animated)
    }

    private var _className: String {
        return String(describing: type(of: self))
    }
}


// MARK: - uninstall

extension UIViewController {
    class func unswizzle() {
        guard isSwizzling else { return }

        let methods: [Selector] = [
            #selector(UIViewController.viewDidLoad),
            #selector(UIViewController.viewWillAppear(_:)),
            #selector(UIViewController.viewDidDisappear(_:)),
        ]
        methods.forEach { (sel) in
            self._unswizzleMethod(sel)
        }
        isSwizzling = false
    }

    private class func _unswizzleMethod(_ selector: Selector) {
        let selectorString = NSStringFromSelector(selector)
        _unswizzleMethod(selectorString)
    }

    private class func _unswizzleMethod(_ selector: String) {
        let originalSelector = NSSelectorFromString(selector)
        let swizzledSelector = NSSelectorFromString("_swizzled_" + selector)

        guard
            let originalMethod = class_getInstanceMethod(self, originalSelector),
            let swizzledMethod = class_getInstanceMethod(self, swizzledSelector)
            else { return }

        if class_addMethod(self, swizzledSelector, method_getImplementation(originalMethod), method_getTypeEncoding(swizzledMethod)) {
            class_replaceMethod(self, originalSelector, method_getImplementation(swizzledMethod), method_getTypeEncoding(swizzledMethod))
        } else {
            method_exchangeImplementations(swizzledMethod, originalMethod)
        }
    }
}