DispatchQueue+throttle+debounce.swift

import Dispatch

private var throttleItems = Atomic([AnyHashable: ThrottleItem]())
private var debounceItems = Atomic([AnyHashable: DebounceItem]())

extension DispatchQueue {
  
  /// First action is executed, then other actions will be performed at most once per specified interval.
  /// - Note:
  ///     Only the first call's `interval` is taken into account and updated every interval.
  /// - Parameters:
  ///   - interval: The interval for which multiple calls will be ignored
  ///   - context: The context in which the throttle should be executed. Defaults to queue's `label`
  ///   - onThrottle: The closure to be executed when the call is throttled
  ///   - action: The closure to be executed
  public func throttle(for interval: DispatchTimeInterval,
                       context: AnyHashable? = nil,
                       onThrottle: (() -> ())? = nil,
                       action: @escaping () -> ()) {
    let context = context ?? label as AnyHashable
    
    defer {
      // Cleanup & release context
      debounce(for: interval * 10) {
        _ = throttleItems.modify { value in
          value.removeValue(forKey: context)
        }
      }
    }
    
    let throttleItem = throttleItems.modify { value -> ThrottleItem? in
      guard let item = value[context] else {
        // FIRST interval info
        let throttleItem = ThrottleItem(lastCallTime: .now(), interval: interval, action: {}, onThrottle: nil)
        value[context] = throttleItem
        async {
          action()
        }
        return throttleItem
      }
      if item.lastCallTime + item.interval > .now() {
        // same interval, cancel previous action, update with new one
        item.onThrottle?()
        item.action = action
        item.onThrottle = onThrottle
        return nil
      }
      
      // new interval info
      let throttleItem = ThrottleItem(lastCallTime: .now(), interval: interval, action: action, onThrottle: onThrottle)
      value[context] = throttleItem
      return throttleItem
    }
    guard let newIntervalItem = throttleItem else {
      return
    }
    
    // new interval will start
    asyncAfter(deadline: .now() + interval) {
      throttleItems.withValue { _ in
        newIntervalItem.action()
      }
    }
  }
  
  /// Last action will be performed after the caller stops calling `debounce` function after a specified interval (cool-down).
  /// - Parameters:
  ///   - interval: The interval for which a calls needs to cool-down
  ///   - context: The context in which the debounce should be executed. Defaults to queue's `label`
  ///   - onDebounce: The closure to be executed when the call is debounced
  ///   - action: The closure to be executed
  public func debounce(for interval: DispatchTimeInterval,
                       context: AnyHashable? = nil,
                       onDebounce: (() -> ())? = nil,
                       action: @escaping () -> ()) {
    let context = context ?? label as AnyHashable
    let worker = DispatchWorkItem {
      defer {
        debounceItems.modify { value in
          value.removeValue(forKey: context)
        }
      }
      action()
    }
    
    asyncAfter(deadline: .now() + interval, execute: worker)
    
    debounceItems.modify { value in
      value[context]?.workItem.cancel()
      value[context]?.onDebounce?()
      value[context] = DebounceItem(workItem: worker, onDebounce: onDebounce)
    }
  }
}

private class ThrottleItem {
  let lastCallTime: DispatchTime
  let interval: DispatchTimeInterval
  var action: () -> ()
  var onThrottle: (() -> ())?
  
  init(lastCallTime: DispatchTime,
       interval: DispatchTimeInterval,
       action: @escaping () -> (),
       onThrottle: (() -> ())?) {
    self.lastCallTime = lastCallTime
    self.interval = interval
    self.action = action
    self.onThrottle = onThrottle
  }
}

private struct DebounceItem {
  let workItem: DispatchWorkItem
  let onDebounce: (() -> ())?
}

private extension DispatchTimeInterval {
  
  static func + (lhs: DispatchTimeInterval, rhs: DispatchTimeInterval) -> DispatchTimeInterval {
    guard let lhsNanos = lhs.nanoseconds,
          let rhsNanos = rhs.nanoseconds else { return .never }
    return .nanoseconds(lhsNanos + rhsNanos)
  }
  
  public static func - (lhs: DispatchTimeInterval, rhs: DispatchTimeInterval) -> DispatchTimeInterval {
    guard let lhsNanos = lhs.nanoseconds,
          let rhsNanos = rhs.nanoseconds else { return .never }
    return .nanoseconds(lhsNanos - rhsNanos)
  }
  
  public static func * (lhs: DispatchTimeInterval, rhs: Int) -> DispatchTimeInterval {
    guard let lhsNanos = lhs.nanoseconds else { return .never }
    return .nanoseconds(lhsNanos * rhs)
  }
  
  static func / (lhs: DispatchTimeInterval, rhs: Int) -> DispatchTimeInterval {
    guard let lhsNanos = lhs.nanoseconds else { return .never }
    return .nanoseconds(lhsNanos / rhs)
  }
  
  var nanoseconds: Int? {
    switch self {
    case .seconds(let value):       return Int(value * 1_000_000_000)
    case .milliseconds(let value):  return Int(value * 1_000_000)
    case .microseconds(let value):  return Int(value * 1_000)
    case .nanoseconds(let value):   return value
    case .never:                    return nil
    @unknown default:               fatalError("Unknown DispatchTimeInterval type")
    }
  }
}

private class Atomic<Value> {
  private var _value: Value
  
  private let queue = DispatchQueue(label: "Atomic", attributes: [.concurrent])
  
  public var value: Value {
    get { queue.sync { _value } }
    set { swap(newValue) }
  }
  
  public init(_ value: Value) {
    _value = value
  }
  
  @discardableResult
  public func modify<Result>(_ action: (inout Value) throws -> Result) rethrows -> Result {
    try queue.sync(flags: .barrier) {
      try action(&_value)
    }
  }
  
  @discardableResult
  public func withValue<Result>(_ action: (Value) throws -> Result) rethrows -> Result {
    try queue.sync(flags: .barrier) {
      try action(_value)
    }
  }
  
  /// Atomically replace the contents of the variable.
  ///
  /// - parameters:
  ///   - newValue: A new value for the variable.
  ///
  /// - returns: The old value.
  @discardableResult
  public func swap(_ newValue: Value) -> Value {
    modify { (value: inout Value) in
      let oldValue = value
      value = newValue
      return oldValue
    }
  }
}