Swift container that weakly holds references to objects.

WeakSet.swift

import Foundation

public struct WeakSet<Element>: Sequence, ExpressibleByArrayLiteral, CustomStringConvertible where Element: AnyObject {

  private var storage = Set<Box>()

  public init(_ members: [Element]) {
    members.forEach { insert($0) }
  }

  public init(arrayLiteral members: Element...) {
    self.init(members)
  }

  public var count: Int { allElements().count }

  public mutating func insert(_ newMember: Element) {
    storage.insert(Box(newMember))
  }

  public mutating func remove(_ member: Element) -> Element? {
    storage.remove(Box(member))?.element
  }

  public func contains(_ member: Element) -> Bool {
    storage.contains(Box(member))
  }

  public func makeIterator() -> AnyIterator<Element> {
    AnyIterator(allElements().makeIterator())
  }

  private func allElements() -> [Element] {
    storage.compactMap(\.element)
  }

  public var description: String {
    allElements().description
  }

  private final class Box: Hashable {

    let identifier: ObjectIdentifier
    weak var element: Element?

    init(_ element: Element) {
      self.identifier = ObjectIdentifier(element)
      self.element = element
    }

    func hash(into hasher: inout Hasher) {
      identifier.hash(into: &hasher)
    }

    static func == (lhs: WeakSet<Element>.Box, rhs: WeakSet<Element>.Box) -> Bool {
      guard let lhsElement = lhs.element, let rhsElement = rhs.element else {
        return false
      }
      return lhsElement === rhsElement
    }
  }
}