Can I hook when a weakly-referenced object (of arbitrary type) is freed?

Question

I'm writing a container class in Swift, which works like as java.util.WeakHashMap in Java. My current implementation is here.

class WeakRefMap<Key: Hashable, Value: AnyObject> {

    private var mapping = [Key: WeakBox<Value>]()

    subscript(key: Key) -> Value? {
        get { return mapping[key]?.raw }
        set {
            if let o = newValue {
                mapping[key] = WeakBox(o)
            }
            else {
                mapping.removeValueForKey(key)
            }
        }
    }

    var count: Int { return mapping.count }
}

class WeakBox<E: AnyObject> {
    weak var raw: E!
    init(  _ raw: E) { self.raw = raw }
}

In this implementation, holded objects in the container are weakly-referenced via WeakBox, so holding values never prevents the objects from being released when not needed anymore.

But clearly there is a problem in this code; The entries remains even after the object of its entry is freed.

To solve this problem, I need to hook just before a holded object is released, and remove its (corresponding) entry. I know a solution only for NSObject, but it's not applicable to AnyObject.

Could anyone help me? Thanks. (^_^)

Try using the `deinit` method - see the [Apple Docs here](https://developer.apple.com/library/ios/documentation/Swift/Conceptual/Swift_Programming_Language/Deinitialization.html). — pbasdf, Feb 23 '15 at 09:49
See this: http://stackoverflow.com/questions/25497928/dealloc-in-swift — Ian, Feb 23 '15 at 09:59
Thanks. But that's not what I meant. The container class I made is a generic class, so it needs to work for all the subtypes of `AnyObject` as its `Value`. The `deinit` approach is applicable only when type is decidable. — findall, Feb 23 '15 at 10:22
It seems that there is no way for Swift classes, compare http://stackoverflow.com/questions/24317332/know-when-a-weak-var-becomes-nil-in-swift. — Martin R, Feb 23 '15 at 10:41

rintaro · Accepted Answer · 2015-02-23T11:11:54.560

Sad to say, didSet or willSet observer doesn't get called when weak var raw property value is deallocated.

So, you have to use objc_setAssociatedObject in this case:

// helper class to notify deallocation
class DeallocWatcher {
    let notify:()->Void
    init(_ notify:()->Void) { self.notify = notify }
    deinit { notify() }
}

class WeakRefMap<Key: Hashable, Value: AnyObject> {

    private var mapping = [Key: WeakBox<Value>]()

    subscript(key: Key) -> Value? {
        get { return mapping[key]?.raw }
        set {
            if let o = newValue {
                // Add helper to associated objects.
                // When `o` is deallocated, `watcher` is also deallocated.
                // So, `watcher.deinit()` will get called.
                let watcher = DeallocWatcher { [unowned self] in self.mapping[key] = nil }
                objc_setAssociatedObject(o, unsafeAddressOf(self), watcher, objc_AssociationPolicy(OBJC_ASSOCIATION_RETAIN_NONATOMIC))
                mapping[key] = WeakBox(o)
            }
            else {
                mapping[key] = nil
            }
        }
    }

    var count: Int { return mapping.count }

    deinit {
        // cleanup
        for e in self.mapping.values {
            objc_setAssociatedObject(e.raw, unsafeAddressOf(self), nil, 0)
        }
    }
}

NOTE: Before Swift 1.2. this solution does not work for arbitrary Swift classes.

Does this work for arbitrary Swift classes or only for subclasses of NSObject (for which OP claims to already have a solution) ? — Martin R, Feb 23 '15 at 10:36
Ah, it seems, it doesn't work for non-`NSObject` subclasses. :( — rintaro, Feb 23 '15 at 10:40
@MartinR It seems, it does work in Xcode 6.3 Beta, but not in Xcode 6.1.1. I don't know why. — rintaro, Feb 23 '15 at 11:07
Thanks a lot, @rintaro san! I've assumed mistakenly that `objc_setAssociatedObject` does not work for `AnyObject`, but now I tested and noticed it works in Swift 1.2. (But not in older versions.) — findall, Feb 23 '15 at 11:08

Vladimir Gorbenko · Answer 2 · 2015-09-16T12:01:46.283

Previous example has some bugs, for example:

Invalid size of dictionary bug: this example prints "1" instead of "2":

let dict = WeakRefMap<String, NSObject>()
autoreleasepool {
    let val = NSObject()
    dict["1"] = val
    dict["2"] = val
    print("dict size: \(dict.count)")
}

Fixed WeakRefMap:

private class DeallocWatcher<Key: Hashable> {

    let notify:(keys: Set<Key>)->Void

    private var keys = Set<Key>()

    func insertKey(key: Key) {
        keys.insert(key)
    }

    init(_ notify:(keys: Set<Key>)->Void) { self.notify = notify }
    deinit { notify(keys: keys) }
}

public class WeakRefMap<Key: Hashable, Value: AnyObject> {

    private var mapping = [Key: WeakBox<Value>]()

    public init() {}

    public subscript(key: Key) -> Value? {
        get { return mapping[key]?.raw }
        set {
            if let o = newValue {
                // Add helper to associated objects.
                // When `o` is deallocated, `watcher` is also deallocated.
                // So, `watcher.deinit()` will get called.

                if let watcher = objc_getAssociatedObject(o, unsafeAddressOf(self)) as? DeallocWatcher<Key> {

                    watcher.insertKey(key)
                } else {

                    let watcher = DeallocWatcher { [unowned self] (keys: Set<Key>) -> Void in
                        for key in keys {
                            self.mapping[key] = nil
                        }
                    }

                    watcher.insertKey(key)

                    objc_setAssociatedObject(o, unsafeAddressOf(self), watcher, .OBJC_ASSOCIATION_RETAIN_NONATOMIC)
                }

                mapping[key] = WeakBox(o)
            } else {
                if let index = mapping.indexForKey(key) {

                    let (_, value) = mapping[index]
                    objc_setAssociatedObject(value.raw, unsafeAddressOf(self), nil, .OBJC_ASSOCIATION_RETAIN_NONATOMIC)
                    mapping.removeAtIndex(index)
                }
            }
        }
    }

    public var count: Int { return mapping.count }

    deinit {
        // cleanup
        for e in self.mapping.values {
            objc_setAssociatedObject(e.raw, unsafeAddressOf(self), nil, .OBJC_ASSOCIATION_RETAIN_NONATOMIC)
        }
    }
}

Can I hook when a weakly-referenced object (of arbitrary type) is freed?

2 Answers2

Linked