CountedSet custom implementation in Swift

CountedSet.swift

struct CountedSet<Element: Hashable> {
  private var storage: [Element: Int] = [:]
}

/// Creating a CountedSet from a Dictionary or DictionaryLiteral
///
/// let example: CountedSet = ["a": 1, "b": 3, "c": 42]
extension CountedSet: ExpressibleByDictionaryLiteral {
  init(dictionaryLiteral elements: (Element, Int)...) {
    self.storage = Dictionary(uniqueKeysWithValues: elements)
  }

  init(_ dictionary: [Element: Int]) {
    self.storage = dictionary
  }
}

/// Creating a CountedSet from an Array or ArrayLiteral
///
/// let example: CountedSet = ["a", "b", "c", "a"]
extension CountedSet: ExpressibleByArrayLiteral {
  init(arrayLiteral elements: Element...) {
    self.init(elements)
  }

  init(_ elements: [Element]) {
    self.storage = elements.reduce(into: [:]) { acc, e in acc[e, default: 0] += 1 }
  }
}

/// Debug Description
extension CountedSet: CustomDebugStringConvertible {
  var debugDescription: String {
    storage.debugDescription
  }
}

/// SetAlgebra conformance
extension CountedSet: SetAlgebra {
  func contains(_ member: Element) -> Bool {
    storage[member, default: 0] > 0
  }

  mutating func insert(_ newMember: __owned Element) -> (inserted: Bool, memberAfterInsert: Element) {
    storage[newMember, default: 0] += 1
    return (true, newMember)
  }

  mutating func remove(_ member: Element) -> Element? {
    if let count = storage[member], count > 0 {
      storage[member] = count == 1 ? nil : count - 1
      return member
    } else {
      return nil
    }
  }

  mutating func update(with newMember: __owned Element) -> Element? {
    if let count = storage[newMember], count > 0 {
      storage[newMember] = count + 1
      return newMember
    } else {
      storage[newMember] = 1
      return nil
    }
  }

  __consuming func union(_ other: __owned CountedSet<Element>) -> CountedSet<Element> {
    var copy = self
    copy.formUnion(other)
    return copy
  }

  __consuming func intersection(_ other: CountedSet<Element>) -> CountedSet<Element> {
    var copy = self
    copy.formIntersection(other)
    return copy
  }

  __consuming func symmetricDifference(_ other: __owned CountedSet<Element>) -> CountedSet<Element> {
    var copy = self
    copy.formSymmetricDifference(other)
    return copy
  }

  mutating func formUnion(_ other: __owned CountedSet<Element>) {
    storage.merge(other.storage, uniquingKeysWith: +)
  }

  mutating func formIntersection(_ other: CountedSet<Element>) {
    for pair: (key: Element, value: Int) in storage {
      if let otherCount = other.storage[pair.key] {
        storage[pair.key] = pair.value + otherCount
      } else {
        storage[pair.key] = nil
      }
    }
  }

  mutating func formSymmetricDifference(_ other: __owned CountedSet<Element>) {
    for pair: (key: Element, value: Int) in other.storage {
      if storage.keys.contains(pair.key) {
        storage[pair.key] = nil
      } else {
        storage[pair.key] = pair.value
      }
    }
  }
}

/// Collection-like Views into the storage, e.g. for iteration
extension CountedSet {
  subscript(element: Element) -> Int {
    storage[element, default: 0]
  }

  /// Collection of unique Elements – each element listed once, regardless of their count in the Set
  var uniquedElements: some Collection /* where .Element = Self.Element */ {
    storage.keys
  }

  /// Sequence of Elements, each Element emited as many times as their count in the Set
  var allElements: some Sequence /* where .Element = Self.Element */ {
    return AnySequence<Element> { () -> AnyIterator<Element> in
      var pairsIterator = storage.makeIterator()
      var currentPair: (key: Element, value: Int)? = nil
      return AnyIterator {
        if currentPair?.value ?? 0 == 0 {
          currentPair = pairsIterator.next()
        }
        currentPair?.value -= 1
        return currentPair?.key
      }
    }
  }

  /// Lazy Collection of (element: Element, count: Int) pairs
  var pairs: some Collection /* where .Element = (element: Self.Element, count: Int) */ {
    storage.lazy.map { (element: $0.key, count: $0.value) }
  }

  var dictionaryRepresentation: [Element: Int] { storage }

  /// Total count of elements in the set, including multiple occurrences
  var count: Int {
    storage.reduce(0) { acc, pair in acc + pair.value }
  }
}

Demo.swift

var s1: CountedSet = ["a", "b", "c", "a"]

print("s1 =", s1)
print(#"s1.contains("x") ="#, s1.contains("x"))
print(#"s1["x"] ="#, s1["x"])
print(#"s1.contains("a") ="#, s1.contains("a"))
print(#"s1["a"] ="#, s1["a"])

s1.insert("d")
s1.insert("b")
s1.insert("b")
print(#"s1 after 3 inserts ("d","b","b"):"#, s1)

s1.update(with: "d")
print(#"s1 after update(with: "d"):"#, s1)

s1.insert("e")
s1.insert("e")
print(#"s1 after 2x insert("e"):"#, s1)
s1.remove("e")
print(#"s1 after one remove("e"):"#, s1)
s1.remove("e")
print(#"s1 after another remove("e"):"#, s1)

print("----")


let s2: CountedSet = ["b": 7, "x": 1, "y": 2, "z": 5]
print("s2 =", s2)

print("s1.union(s2) =", s1.union(s2))
assert(s1.union(s2) == s2.union(s1))
print("s1.intersection(s2) =", s1.intersection(s2))
assert(s1.intersection(s2) == s2.intersection(s1))
print("s1.symmetricDifference(s2) =", s1.symmetricDifference(s2))
assert(s1.symmetricDifference(s2) == s2.symmetricDifference(s1))

print("----")

print("s1.uniquedElements =", s1.uniquedElements)
print("s1.allElements =", Array(s1.allElements))

print("Iterating on s1.pairs:")
for e in s1.pairs {
  print("-", e)
}
print("s1.dictionaryRepresentation =", s1.dictionaryRepresentation)
print("s1.count =", s1.count)