MH175 MH175 - 22 days ago 4
Swift Question

Implementing a custom iterator that skips nil elements in a WeakSet

I am implementing a

WeakSet
, which wraps its elements weakly in a
WeakWrapper
so as to not increase their retain count.

My question is, how do I create an iterator so that I can iterate over the elements skipping those that have been deallocated (i.e. are
nil
). Note that I am trying to optimize over the iteration; it's ok if insertion/removal are relatively slower, but there should be little/no performance cost to setting up the iterator.

Here is my
WeakSet
in its basic form. I can call
clean()
to remove
WeakWrapper
s whose objects have been deallocated.

struct WeakSet<T> where T: AnyObject & Hashable {
private var set: Set<WeakWrapper<T>> = []

mutating func insert(_ elem: T) {
self.set.insert(WeakWrapper<T>(elem))
}

mutating func remove(_ elem: T) {
self.set.remove(WeakWrapper<T>(elem))
}

mutating func clean() {
for elem in set {
if elem.obj == nil {
self.set.remove(elem)
}
}
}
}

fileprivate class WeakWrapper<T>: Hashable where T: AnyObject {
weak var obj: T?
let hashValue: Int

init(_ obj: T) {
self.obj = obj
self.hashValue = ObjectIdentifier(obj).hashValue
}

static func ==(lhs: WeakWrapper, rhs: WeakWrapper) -> Bool {
return lhs.hashValue == rhs.hashValue
}
}


I want to be able to do something like this, where the generated elements are the underlying non-nil elements of type
T
, not the wrapped elements.

class MyObject: NSObject {
func doSomething() { }
}

var weakSet = WeakSet<MyObject>()
for myObject in weakSet {
myObject.doSomething()
}

Answer Source

A possible solution, using built-in methods from the Swift standard library:

extension WeakSet: Sequence {
    func makeIterator() -> AnyIterator<T> {
        return AnyIterator(self.set.lazy.flatMap { $0.obj }.makeIterator())
    }
}

Starting with the lazy view of of the set, a (lazy) collection of its non-nil objects is created using flatMap.

It works also without the lazy, but then an array with all non-nil objects is created eagerly as soon as makeIterator() is called.

Another solution, using a custom iterator type:

struct WeakSetIterator<T>: IteratorProtocol where T: AnyObject {
    fileprivate var iter: SetIterator<WeakWrapper<T>>

    mutating func next() -> T? {
        while let wrapper = iter.next() {
            if let obj = wrapper.obj { return obj }
        }
        return nil
    }
}

extension WeakSet: Sequence {
    func makeIterator() -> WeakSetIterator<T> {
        return WeakSetIterator(iter: self.set.makeIterator())
    }
}