Dictionary subscript or operator which adds a specified "default" value if key doesn't exist is Swift?

Question

I need to write something like the following:

var dict = [SomeEnum: SomeObject]()
var someValue: SomeObject! = dict[.case1]
if someValue == nil {
  someValue = someDefaultValue
  dict[.case1] = someValue
}

The Nil-coalescing operator already exists but it doesn't mutate the original dictionary. The same is for dict[.case1, default: someDefaultValue.

Could you please explain how to write a new subscript/operator to wrap the code above?

Answer 1

You can create custom subscripts by declaring a new method using the subscript keyword instead of the func keyword.

Based on your requirements, it seems you only need your custom subscript to have a getter, you don't need a setter for it - you can use the default setter to set values.

extension Dictionary {
  /// Subscript which returns the previously set `Value` for `key` in case such value exists.
  /// Otherwise, sets `defaultValue` as the value for `key` and returns `defaultValue`.
  subscript(_ key: Key, defaultValue defaultValue: Value) -> Value {
    mutating get {
      // In case the dictionary already contained a value for key, return the existing value
      if let value = self[key] {
        return value
      } else { // Otherwise, set the key to have defaultValue as its value and return defaultValue
        self[key] = defaultValue
        return defaultValue
      }
    }
  }
}

Usage example:

let defaultValue = "default"
var dict = [Int:String]()
dict[0] = "0"
dict[2] = "2"
let valueFor1 = dict[1, defaultValue: defaultValue] // "default"
print(dict) // [0: "0", 1: "default", 2: "2"]
dict[2, defaultValue: defaultValue] // this returns 2 as the key had a previously set value