Scoped extension of mutable state

This post contains the seed of an idea for an evolution post, but I'm posting this here as there are probably better ways to achieve what I want.

What I'm looking for is to be able to add mutable state to a class in a way that is private (or at least can be private) to a particular piece of code that uses that class.

For example, let's imagine I've got a game with a mutable Unit type that represents a unit, and I've got an immutable struct that implements the Perk protocol, so that a Unit gets a damage boost if they've been shot in the last two turns. That perk needs to know the number of turns since a unit was shot, but the perk can't count the turns (because it's immutable) and we don't want to add a 'turns since shot' property to Unit because it has too many properties already.

It would be nice if I could add a private property extension that only the perk can see, but extensions cannot add stored properties.

So I decided to go the route of dictionaries, and add some sort of dictionary to Unit. But now the problem is how to keep one perk's keys (and values) private from another perk (should it want to, at least.)

I can't use enum as the key, because enums can't be extended, so all the keys must be public. I can't use Type, because types can't be used as keys. I can use ObjectIdentifier, and with some wrapping, it's not too bad:

class Bag {
    var bag: [ObjectIdentifier: Any]
    
    init() {
        self.bag = [:]
    }
    
    func set(value: Any, for key: Any.Type) {
        bag[ObjectIdentifier(key)] = value
    }
    
    func get<T>(key: Any.Type) -> T {
        bag[ObjectIdentifier(key)] as! T
    }
}

and to use it...

fileprivate enum TurnsSinceShot { }

struct MoreDamageWhenShotPerk {
    func handle(event: Event) {
        switch event {
            ...
            case is UnitShot: 
                unit.bag.set(value: 0, for: TurnsSinceShot.self)
            ...
        }
    }

    func damageMultiplier() -> Float {
        let turnsSinceShot: Int =  unit.bag.get(key: TurnsSinceShot.self)
        return turnsSinceShot <= 2 
            ? 1.5
            : 1.0
    }
}

...and this lets me choose wether bag values are private to this perk or public, but the issue here s I need to create a type for every value I need to store. It also feels a bit gross and over-complex.

Is there a better way to do this? Am I missing something incredibly obvious? Could something enabling this become a language feature?

This sounds a lot like SwiftUI’s EnvironmentValues, which has these key features:

  • A subscript that accepts a value of an arbitrary type that conforms to…
  • An EnvironmentKey protocol that defines an associated Value type for type-safe storage, along with a default value for that property. But if you’re ok with making the return value of get(key:) always optional, it could be fine to avoid the default value property and just return nil like a dictionary does.
  • An @Entry macro:
    extension EnvironmentValues {
        @Entry var foo: Bool = true
    }
    // expands to:
    extension EnvironmentValues {
        var foo: Bool {
            get { self[__Key_foo.self] }
            set { self[__Key_foo.self] = newValue }
        }
        private struct __Key_foo: EnvironmentKey {
            typealias Value = Bool
            static let defaultValue: Value = true
        }
    }
    
  • Internal storage mapping keys to their associated values. A dictionary could work for this if you don’t expect the storage to be accessed extremely frequently, but it might make sense to file a request in Feedback Assistant (or on the swift-collections repository?) to ask if a higher-performance option could be added, potentially based on the SwiftUI implementation.
1 Like

An alternate approach could be to use the ObjectIdentifier of the containing object as a key into a separately-stored dictionary. This would allow you to have type safety without needing to use a more complicated data structure. Using the ObjectIdentifier as the key instead of the real object would avoid keeping the object alive longer than necessary.

class Bag<Value> {
    var bag: [ObjectIdentifier: Value]
    
    init() {
        self.bag = [:]
    }
    
    subscript(object: AnyObject) -> Value? {
        get {
            bag[ObjectIdentifier(object)]
        }
        set {
            bag[ObjectIdentifier(object)] = newValue
        }
    }
}

extension Unit {
    private static let turnsSinceShotBag = Bag<Int>()
    var turnsSinceShot: Int? {
        get { Self.turnsSinceShotBag[self] }
        set { Self.turnsSinceShotBag[self] = newValue } 
    }
}

One downside to this approach as opposed to my other suggestion is that you would need to manually garbage-collect the entries in the bag somehow. Without garbage collection, you would end up with abandoned entries in the dictionary that are no longer valid (increasing your memory footprint), but the more serious concern is that object identifiers can be reused so it is possible a future object would end up inheriting values stored by a past object. There are a couple of workarounds here.

  1. Use a custom key type that stores the ObjectIdentifier along with a weak reference. When accessing a value, ensure that the weak reference is not nil and if it is, remove the entry and return nil. (index(forKey:) might be useful here). But this doesn’t fix the problem of abandoned memory slowly growing over time as objects go away.
  2. Periodically scan the dictionary for values with a nil weak reference. This is garbage collection, though, and Swift avoids it for a reason (to ensure predictable performance).
  3. Require that the key object conform to a protocol, and then add code to deinit to remove the object’s value from the bag. This fixes both problems, but requires a change to the object that will be used as the key so it can report back to get cleaned up later.
Baggable protocol + implementation
protocol Baggable: AnyObject {
    func willBeStored(in bag: some BagProtocol)
}
protocol BagProtocol: AnyObject {
    func removeValue(for object: AnyObject)
}
class Bag<Value>: BagProtocol {
    var bag: [ObjectIdentifier: Value] = [:]
    subscript(object: some Baggable) -> Value? {
        get {
            bag[ObjectIdentifier(object)]
        }
        set {
            if bag[ObjectIdentifier(object)] == nil {
                object.willBeStored(in: self)
            }
            bag[ObjectIdentifier(object)] = newValue
        }
    }
    func removeValue(for object: AnyObject) {
        bag.removeValue(forKey: ObjectIdentifier(object))
    }
}

class Unit: Baggable {
    var bags: [any BagProtocol] = []
    func willBeStored(in bag: some BagProtocol) {
        bags.append(bag)
    }
    deinit {
        for bag in bags {
            bag.removeValue(for: self)
        }
    }
}