UIButton.swift

import UIKit

extension UIButton {
    
    /// Sets the background color to use for the specified button state.
    /// - Parameters:
    ///   - color: The background color to use for the specified state.
    ///   - forState: The state that uses the specified background color. The values are described in `UIControl.State`.
    public func setBackgroundColor(_ color: UIColor, forState: UIControl.State) {
        if let colorImage = self.imageWithColor(color) {
            self.setBackgroundImage(colorImage, for: forState)
        }
    }
    
    /// Sets the image tint color to use for the specified button state.
    /// - Parameters:
    ///   - color: The tint color to use for the specified state.
    ///   - forState: The state that uses the specified tint color. The values are described in `UIControl.State`.
    public func setImageTintColor(_ color: UIColor, forUIControlState state: UIControl.State) {
        if let image = self.image(for: state),
            let tintedImage = self.tintedImageWithColor(color, image: image)
        {
            self.setImage(tintedImage, for: state)
        }
    }
    
    /// Sets the background tint color to use for the specified button state.
    /// - Parameters:
    ///   - color: The tint color to use for the specified state.
    ///   - forState: The state that uses the specified background tint color. The values are described in `UIControl.State`.
    public func setBackgroundTintColor(_ color: UIColor, forUIControlState state: UIControl.State) {
        if let backgroundImage = self.backgroundImage(for: state),
            let tintedImage = self.tintedImageWithColor(color, image: backgroundImage)
        {
            self.setBackgroundImage(tintedImage, for: state)
        }
    }
    
    private func imageWithColor(_ color: UIColor) -> UIImage? {
        UIGraphicsBeginImageContext(CGSize(width: 1, height: 1))
        
        defer { UIGraphicsEndImageContext() }
        
        guard let context = UIGraphicsGetCurrentContext() else {
            return nil
        }
        
        context.setFillColor(color.cgColor)
        context.fill(CGRect(x: 0, y: 0, width: 1, height: 1))
        
        return UIGraphicsGetImageFromCurrentImageContext()
    }
    
    private func tintedImageWithColor(_ color: UIColor, image: UIImage) -> UIImage? {
        UIGraphicsBeginImageContextWithOptions(image.size, false, UIScreen.main.scale)
        
        defer { UIGraphicsEndImageContext() }
        
        guard let context = UIGraphicsGetCurrentContext(), let cgImage = image.cgImage else {
            return nil
        }
        
        context.translateBy(x: 0, y: image.size.height)
        context.scaleBy(x: 1.0, y: -1.0)
        let rect = CGRect(x: 0, y: 0, width: image.size.width, height: image.size.height)
        context.setBlendMode(.normal)
        context.draw(cgImage, in: rect)
        context.setBlendMode(.sourceIn)
        color.setFill()
        context.fill(rect)
        
        return UIGraphicsGetImageFromCurrentImageContext()
    }
}