A Better `UIStackView`

UIModernStackView.swift

//
// Copyright (c) Vatsal Manot
//

#if canImport(UIKit)

import Swift
import UIKit

///
/// A modern `UIStackView` replacement.
///
open class UIModernStackView<ArrangedSubview: UIView>: UIStackView {
    override open var axis: NSLayoutConstraint.Axis {
        didSet {
            super.invalidateIntrinsicContentSize()
        }
    }

    override open var distribution: Distribution {
        didSet {
            super.invalidateIntrinsicContentSize()
        }
    }

    override open var alignment: Alignment {
        didSet {
            super.invalidateIntrinsicContentSize()
        }
    }

    override open var spacing: CGFloat {
        didSet {
            super.invalidateIntrinsicContentSize()
        }
    }

    override open var intrinsicContentSize: CGSize {
        var size = CGSize.zero

        var hasSubviewWithNoIntrinsicWidth: Bool = false
        var hasSubviewWithNoIntrinsicHeight: Bool = false

        for view in arrangedSubviews {
            let intrinsicContentSize = view.intrinsicContentSize

            if intrinsicContentSize.width == UIView.noIntrinsicMetric {
                hasSubviewWithNoIntrinsicWidth = true
            }

            if intrinsicContentSize.height == UIView.noIntrinsicMetric {
                hasSubviewWithNoIntrinsicHeight = true
            }

            if hasSubviewWithNoIntrinsicWidth && hasSubviewWithNoIntrinsicHeight {
                return super.intrinsicContentSize
            }

            switch axis {
            case .horizontal:
                size.width += intrinsicContentSize.width
                size.height = max(size.height, intrinsicContentSize.height)
            case .vertical:
                size.width = max(size.width, intrinsicContentSize.width)
                size.height += intrinsicContentSize.height
            }
        }

        if arrangedSubviews.count != 0 {
            let totalSpacing = .init(arrangedSubviews.count - 1) * spacing

            switch axis {
            case .horizontal:
                size.width += totalSpacing
            case .vertical:
                size.height += totalSpacing
            }
        }

        if hasSubviewWithNoIntrinsicWidth {
            size.width = UIView.noIntrinsicMetric
        }

        if hasSubviewWithNoIntrinsicHeight {
            size.height = UIView.noIntrinsicMetric
        }

        return size
    }

    init() {
        super.init(frame: .zero)
    }

    public required init(coder: NSCoder) {
        fatalError("init(coder:) has not been implemented")
    }

    override open func invalidateIntrinsicContentSize() {
        super.invalidateIntrinsicContentSize()

        for view in arrangedSubviews {
            view.invalidateIntrinsicContentSize()
        }
    }

    override open func layoutSubviews() {
        super.layoutSubviews()
        super.invalidateIntrinsicContentSize()
    }

    override open func setContentHuggingPriority(_ priority: UILayoutPriority, for axis: NSLayoutConstraint.Axis) {
        super.setContentHuggingPriority(priority, for: axis)

        for view in arrangedSubviews {
            view.setContentHuggingPriority(priority, for: axis)
        }

        super.invalidateIntrinsicContentSize()
    }

    open var contentHuggingPriority: UILayoutPriority {
        get {
            return contentHuggingPriority(for: axis)
        } set {
            setContentHuggingPriority(newValue, for: axis)
        }
    }

    override open func setContentCompressionResistancePriority(_ priority: UILayoutPriority, for axis: NSLayoutConstraint.Axis) {
        super.setContentCompressionResistancePriority(priority, for: axis)

        for view in arrangedSubviews {
            view.setContentCompressionResistancePriority(priority, for: axis)
        }

        super.invalidateIntrinsicContentSize()
    }

    open var contentCompressionResistancePriority: UILayoutPriority {
        get {
            return contentCompressionResistancePriority(for: axis)
        } set {
            setContentCompressionResistancePriority(newValue, for: axis)
        }
    }

    override open func didAddSubview(_ subview: UIView) {
        super.didAddSubview(subview)
        super.invalidateIntrinsicContentSize()

        setNeedsLayout()
    }

    @available(*, deprecated)
    override open func insertArrangedSubview(_ view: UIView, at stackIndex: Int) {
        precondition(view is ArrangedSubview)

        super.insertArrangedSubview(view, at: stackIndex)
    }

    open func insertArrangedSubview(_ view: ArrangedSubview, at stackIndex: Int) {
        super.insertArrangedSubview(view, at: stackIndex)
    }

    @available(*, deprecated)
    override open func addArrangedSubview(_ view: UIView) {
        precondition(view is ArrangedSubview)

        super.addArrangedSubview(view)
    }

    open func addArrangedSubview(_ view: ArrangedSubview) {
        super.addArrangedSubview(view)
    }

    @available(*, deprecated)
    override open func removeArrangedSubview(_ view: UIView) {
        precondition(view is ArrangedSubview)

        super.removeArrangedSubview(view)
    }

    open func removeArrangedSubview(_ view: ArrangedSubview) {
        super.removeArrangedSubview(view)
    }
}

extension UIModernStackView {
    public static func horizontal(
        _ arrangedSubviews: [ArrangedSubview],
        distribution: UIStackView.Distribution = .equalSpacing,
        alignment: UIStackView.Alignment = .center,
        spacing: CGFloat = 0) -> UIModernStackView<ArrangedSubview> {
        let result = UIModernStackView<ArrangedSubview>()

        result.axis = .horizontal
        result.distribution = distribution
        result.alignment = alignment
        result.spacing  = spacing
        arrangedSubviews.forEach(result.addArrangedSubview)

        return result
    }

    /// Returns a vertical stack view whose parameters can either be specified or left to defaults.
    public static func vertical(
        _ arrangedSubviews: [ArrangedSubview],
        distribution: UIStackView.Distribution = .equalSpacing,
        alignment: UIStackView.Alignment = .center,
        spacing: CGFloat = 0) -> UIModernStackView<ArrangedSubview> {
        let result = UIModernStackView<ArrangedSubview>()

        result.axis = .vertical
        result.distribution = distribution
        result.alignment = alignment
        result.spacing  = spacing
        arrangedSubviews.forEach(result.addArrangedSubview)

        return result
    }
}

#endif