How can I translate this utility function into an extension function?

Question

I wrote this utility function in Swift 4:

func insert<Key, Element>(_ value: Element, into dictionary: inout [Key : Set<Element>], at key: Key) {
    if let _ = dictionary[key] {
        dictionary[key]?.insert(value)
    }
    else {
        var set = Set<Element>()
        set.insert(value)
        dictionary[key] = set
    }
}

This is used like this:

insert("foo", into: &myDictionary, at: "bar")

... but I want to use it like this:

myDictionary.insert("foo", at: "bar")

I tried declaring it like this:

extension Dictionary where Value == Set<AnyHashable> {
    mutating func insert(_ value: Value.Element, at key: Key) { // Error here
        if let _ = self[key] {
            self[key]?.insert(value)
        } else {
            var set = Set<Value.Element>() // Error here
            set.insert(value)
            self[key] = set
        }
    }
}

... but I get the following errors:

/path/to/Sequence Extensions.swift:2:41: error: 'Element' is not a member type of 'Dictionary.Value'
    mutating func insert(_ value: Value.Element, at key: Key) {
                                  ~~~~~ ^
Swift.Set:608:22: note: did you mean 'Element'?
    public typealias Element = Element
                     ^
Swift._IndexableBase:3:22: note: did you mean '_Element'?
    public typealias _Element = Self.Element

/path/to/Sequence Extensions.swift:6:23: error: type 'Value.Element' does not conform to protocol 'Hashable'
            var set = Set<Value.Element>()
                      ^

Answer 1

You could do it using a protocol to identify Sets:

protocol SetType
{
   associatedtype Element:Hashable
   init()
   mutating func insert(_ : Element) ->  (inserted: Bool, memberAfterInsert: Element)
}

extension Set:SetType 
{}

extension Dictionary where Value : SetType 
{
   mutating func insert(_ value:Value.Element, at key:Key)
   {
      var valueSet:Value = self[key] ?? Value()
      valueSet.insert(value)
      self[key] = valueSet
   }
}

var oneToMany:[String:Set<String>] = [:]

oneToMany.insert("Dog", at: "Animal")
oneToMany.insert("Cat", at: "Animal")
oneToMany.insert("Tomato", at: "Vegetable")

This will produce a dictionary of sets:

["Animal": Set(["Dog", "Cat"]), "Vegetable": Set(["Tomato"])]

A more appropriate implementation would use the same return value as a Set's insert() function however:

extension Dictionary where Value : SetType 
{
   @discardableResult
   mutating func insert(_ value:Value.Element, at key:Key) ->  (inserted: Bool, memberAfterInsert: Value.Element)
   {
      var valueSet:Value = self[key] ?? Value()
      let result = valueSet.insert(value)
      if result.inserted 
      { self[key] = valueSet }
      return result
   }
}

[EDIT] I just read all of Hamish's response and realized that he had already given the same answer (essentially) and made use of SetAlgebra ( which I wasn't aware of) that does the same thing as the SetType I "reinvented". You should accept Hamish's answer.

Answer 2

Unfortunately, Swift doesn't currently support parameterised extensions (the ability to introduce type variables in extension declarations), so you cannot currently directly express the notion of "an extension with a constraint to some Set<T>". However, it is a part of the generics manifesto, so hopefully it's something that makes its way into a future version of the language.

Even if your extension with Value constrained to Set<AnyHashable> compiled, it wouldn't be terribly useful. You would need to first convert your desired dictionary to a temporary [Key: Set<AnyHashable>], then call the mutating method on it, and then convert it back to its original type (using as!).

This is because the extension is on a Dictionary with heterogenous Set values. It would've been perfectly legal for the extension method to insert arbitrary Hashable elements into one of the values of the dictionary. But that's not what you wanted to express.

In simple cases, I would argue that there's no need for an extension in the first place. You can just say:

var dict = [String: Set<String>]()
dict["key", default: []].insert("someValue")

using Dictionary's subscript overload that takes a default value, as introduced in SE-0165.

If you still want an extension, I would advise simply making it more generic. For example, instead of constraining Value to Set; constrain it to the protocol SetAlgebra (which Set conforms to).

It represents types that can perform set-like operations, and also derives from ExpressibleByArrayLiteral meaning that you can implement your method using the exact syntax as above:

extension Dictionary where Value : SetAlgebra {

    mutating func insert(_ value: Value.Element, at key: Key) {
        self[key, default: []].insert(value)
    }
}

Although one additional thing to consider here is the copy-on-write behaviour of Swift's collection types such as Set. In the above method, the dictionary will be queried for a given key, giving back either an existing set for that key, or a new empty one. Your value will then be inserted into this temporary set, and it will be re-inserted back into the dictionary.

The use of a temporary here means that if the set is already in the dictionary, the value will not be inserted into it in-place, the set's buffer will be copied first in order to preserve value semantics; which could be a performance concern (this is explored in more detail in this Q&A and this Q&A).

However that being said, I am currently looking to fix this for Dictionary's subscript(_:default:) in this pull request, such that the set can be mutated in-place.

Until fixed though, the solution is to first remove the set from the dictionary before mutating:

extension Dictionary where Value : SetAlgebra {

    mutating func insert(_ value: Value.Element, at key: Key) {
        var set = removeValue(forKey: key) ?? []
        set.insert(value)
        self[key] = set
    }
}

In which case, the use of an extension is fully justified.

It's worth noting that the use of a protocol constraint here is the general solution (or workaround in some cases) to the problem of not having parameterised extensions. It allows you to realise the placeholders you need as associated types of that protocol. See this Q&A for an example of how you can create your own protocol to serve that purpose.