Implementation of a flexible size bit set in Swift.

BitSet.swift

extension BidirectionalCollection {

    public func last(where predicate: (Element) throws -> Bool) rethrows -> Element? {
        for element in self.reversed() where try predicate(element) {
            return element
        }
        return nil
    }
}

public struct WordBitSet {
    
    public static let capacity = MemoryLayout<UInt>.size * 8
    
    public var rawValue: UInt
    
    public init(rawValue: UInt) {
        self.rawValue = rawValue
    }
}

extension WordBitSet: SetAlgebra {
    
    public typealias Element = Int
    
    public init() {
        self.rawValue = 0
    }
    
    public var isEmpty: Bool {
        return 0 == self.rawValue
    }
    
    public static func ==(lhs: WordBitSet, rhs: WordBitSet) -> Bool {
        return lhs.rawValue == rhs.rawValue
    }
    
    public func contains(_ member: Element) -> Bool {
        precondition(
            member < WordBitSet.capacity,
            "WordBitSet can only contain values in the range 0 ..< \(WordBitSet.capacity)"
        )
        return 0 != (rawValue & (1 << member))
    }
    
    public mutating func formIntersection(_ other: WordBitSet) {
        self.rawValue &= other.rawValue
    }
    
    public mutating func formSymmetricDifference(_ other: WordBitSet) {
        self.rawValue ^= other.rawValue
    }
    
    public mutating func formUnion(_ other: WordBitSet) {
        self.rawValue |= other.rawValue
    }
    
    @discardableResult
    public mutating func insert(_ newMember: Element) -> (inserted: Bool, memberAfterInsert: Element) {
        let inserted = self.contains(newMember)
        self.rawValue |= 1 << newMember
        return (inserted, newMember)
    }
    
    public func intersection(_ other: WordBitSet) -> WordBitSet {
        return WordBitSet(rawValue: self.rawValue & other.rawValue)
    }
    
    public func symmetricDifference(_ other: WordBitSet) -> WordBitSet {
        return WordBitSet(rawValue: self.rawValue ^ other.rawValue)
    }
    
    @discardableResult
    public mutating func remove(_ member: Element) -> Element? {
        guard self.contains(member) else { return nil }
        self.rawValue ^= 1 << member
        return member
    }
    
    public func union(_ other: WordBitSet) -> WordBitSet {
        return WordBitSet(rawValue: self.rawValue | other.rawValue)
    }
    
    @discardableResult
    public mutating func update(with newMember: Element) -> Element? {
        let inserted = self.contains(newMember)
        self.rawValue |= 1 << newMember
        return inserted ? newMember : nil
    }
}

extension WordBitSet {
    
    init(lessThan upperBound: Int) {
        precondition(upperBound <= WordBitSet.capacity)
        self.rawValue = (1 << upperBound) &- 1
    }
    
    public init(_ range: Range<Int>) {
        self.init(lessThan: range.upperBound)
        self.formSymmetricDifference(WordBitSet(lessThan: range.lowerBound))
    }
}

extension WordBitSet: BidirectionalCollection {

    public typealias Index = Int
    public typealias SubSequence = WordBitSet

    public var startIndex: Index {
        guard !self.isEmpty else { return endIndex }
        return (0 ..< WordBitSet.capacity).first(where: self.contains)!
    }

    public var endIndex: Index {
        return WordBitSet.capacity
    }

    public subscript(_ position: Index) -> Element {
        precondition(self.contains(position))
        return position
    }

    public subscript(bounds: Range<Index>) -> SubSequence {
        return self.intersection(WordBitSet(bounds))
    }

    public func index(after i: Index) -> Index {
        precondition(i < endIndex)
        let nextMember = (i + 1 ..< WordBitSet.capacity).first(where: self.contains)
        return nextMember ?? endIndex
    }

    public func index(before i: Index) -> Index {
        let startIndex = self.startIndex
        precondition(startIndex < i)
        let previousMember = (startIndex ..< i).last(where: self.contains)
        return previousMember ?? startIndex
    }
}

extension WordBitSet : CustomStringConvertible {

    public var description: String {
        return "{" + (self.lazy.map { $0.description }).joined(separator: ", ") + "}"
    }
}

extension WordBitSet: CustomPlaygroundQuickLookable {

    public var customPlaygroundQuickLook: PlaygroundQuickLook {
        return .text(description)
    }
}

public struct BitSet {

    public var words: [WordBitSet]

    public init() {
        self.words = []
    }
}

extension BitSet {
    public typealias Index = Int
}

private extension BitSet {

    mutating func extend(newWordCount: Int) {
        guard words.count < newWordCount else { return }
        let wordsToAdd = newWordCount - words.count
        words.append(contentsOf: repeatElement(WordBitSet(), count: wordsToAdd))
    }

    mutating func extend(newLowerBound: Element) {
        extend(newWordCount: BitSet.indexOfWord(newLowerBound) + 1)
    }

    mutating func truncate() {
        var newWordCount = words.count
        for index in words.indices.reversed() {
            guard words[index].isEmpty else { break }
            newWordCount = index
        }
        words.removeLast(words.count - newWordCount)
    }
}

extension BitSet {

    public static func indexOfWord(_ value: BitSet.Index) -> Array<WordBitSet>.Index {
        return value / WordBitSet.capacity
    }

    public static func indexInWord(_ value: BitSet.Index) -> WordBitSet.Index {
        return value % WordBitSet.capacity
    }

    public static func index(indexOfWord: Array<WordBitSet>.Index, indexInWord: WordBitSet.Index) -> BitSet.Index {
        return indexOfWord * WordBitSet.capacity + indexInWord
    }
}

extension BitSet: SetAlgebra {

    public typealias Element = Int

    public var isEmpty: Bool {
        return words.isEmpty
    }

    public static func ==(lhs: BitSet, rhs: BitSet) -> Bool {
        return lhs.words == rhs.words
    }

    public func contains(_ member: Element) -> Bool {
        guard BitSet.indexOfWord(member) < words.count else { return false }
        return words[BitSet.indexOfWord(member)].contains(BitSet.indexInWord(member))
    }

    public mutating func formIntersection(_ other: BitSet) {
        for index in 0 ..< min(self.words.count, other.words.count) {
            self.words[index].formIntersection(other.words[index])
        }
        truncate()
    }

    public mutating func formSymmetricDifference(_ other: BitSet) {
        let newWordCount = max(self.words.count, other.words.count)
        var _other = other
        _other.extend(newWordCount: newWordCount)
        self.extend(newWordCount: newWordCount)
        for index in 0 ..< newWordCount {
            self.words[index].formSymmetricDifference(other.words[index])
        }
        truncate()
    }

    public mutating func formUnion(_ other: BitSet) {
        for index in 0 ..< min(self.words.count, other.words.count) {
            self.words[index].formUnion(other.words[index])
        }
        guard self.words.count < other.words.count else { return }
        self.words.append(contentsOf: other.words.suffix(from: self.words.count))
    }

    @discardableResult
    public mutating func insert(_ newMember: Element) -> (inserted: Bool, memberAfterInsert: Element) {
        extend(newLowerBound: newMember)
        let inserted = self.contains(newMember)
        words[BitSet.indexOfWord(newMember)].insert(BitSet.indexInWord(newMember))
        return (inserted, newMember)
    }

    public func intersection(_ other: BitSet) -> BitSet {
        var new = self
        new.formIntersection(other)
        return new
    }

    public func symmetricDifference(_ other: BitSet) -> BitSet {
        var new = self
        new.formSymmetricDifference(other)
        return new
    }

    @discardableResult
    public mutating func remove(_ member: Element) -> Element? {
        guard self.contains(member) else { return nil }
        words[BitSet.indexOfWord(member)].remove(BitSet.indexInWord(member))
        truncate()
        return member
    }

    public func union(_ other: BitSet) -> BitSet {
        var new = self
        new.formUnion(other)
        return new
    }

    @discardableResult
    public mutating func update(with newMember: Element) -> Element? {
        extend(newLowerBound: newMember)
        let inserted = self.contains(newMember)
        guard inserted else { return nil }
        words[BitSet.indexOfWord(newMember)].update(with: BitSet.indexInWord(newMember))
        return newMember
    }
}

extension BitSet: CustomStringConvertible {

    public var description: String {

        let memberDescriptions = words.indices.lazy.flatMap { (indexOfWord: Int) in
            return self.words[indexOfWord].lazy.map { (indexInWord: WordBitSet.Index) -> String in
                return BitSet.index(indexOfWord: indexOfWord, indexInWord: indexInWord).description
            }
        }

        return "{" + memberDescriptions.joined(separator: ", ") + "}"
    }
}

// Playground


do {
    let set0: Set<Int> = [0, 4, 16]
    var set1: Set<Int> = set0
    
    set1.remove(0)
    set1.isSubset(of: set0)
    set1.insert(32)
    set0.symmetricDifference(set1)
}

do {
    let set0: WordBitSet = [0, 4, 16]
    var set1: WordBitSet = set0

    set1.remove(0)
    set1.isSubset(of: set0)
    set1.insert(32)
    set0.symmetricDifference(set1)
}

do {
    let set0: BitSet = [0, 4, 16]
    var set1: BitSet = set0

    set1.remove(0)
    set1.isSubset(of: set0)
    set1.insert(32)
    set0.symmetricDifference(set1)
}