Search code examples
arraysswiftprotocols

Conform a custom type to RandomAccessCollection with only a subscript


I usually implement types that behave like arrays, like this one:

struct Dataset: RandomAccessCollection {
    let ids: [Int]
    // Other properties and methods...

    // Boilerplate
    var startIndex: Int { ids.startIndex }
    var endIndex: Int { ids.endIndex }
    func formIndex(after i: inout Int) { i += 1 }
    func formIndex(before i: inout Int) { i -= 1 }

    subscript(index: Int) -> Int {
        // Dummy example, could be more complex and return a different type
        return ids[index]
    }
}

The problem is that I need to write each time lot of boilerplate code for the RandomAccessCollection conformance. I'd like a protocol or a mechanism reduce the boilerplate to just one or two requirements:

  • An underlying RandomAccessCollection (like the ids property in my example) from which to infer the protocol requirements (startIndex, endIndex, formIndex)
  • A subscript that would complete the remaining requirement

This mechanism would resemble the way Dataset inheritance is done in Pytorch currently: only with a __len__ and __getitem__ requirements.

I come up with a draft like this:

protocol ArrayProtocol: RandomAccessCollection where Index == BaseCollection.Index {
    associatedtype BaseCollection: RandomAccessCollection
    
    var baseCollection: BaseCollection { get set }
    subscript(index: Index) -> Element { get set }
}

// Provide the default implementation of the RandomAccessCollection protocol
extension ArrayProtocol {
    var startIndex: Index { baseCollection.startIndex }
    var endIndex: Index { baseCollection.endIndex }
    func formIndex(after i: inout Index) { baseCollection.index(after: i) }
    func formIndex(before i: inout Index) { baseCollection.index(before: i) }
}

This protocol would be used like that:

struct Dataset: ArrayProtocol {
    let ids: [Int]
    // Other properties and methods...

    // No more boilerplate
    var baseCollection: [Int] { ids }

    subscript(index: Int) -> Int {
        // Dummy example, could be more complex and return a different type
        return ids[index]
    }
}

But I can't find a way to make it working and I feel like the associated type is not a very good design pattern.

Any idea to solve this?

EDIT: the where Index == BaseCollection.Index clause is not necessary, the subscript could have a different Index type than the underlying collection.


Solution

  • Associated types are generally used for the generic type (usually the element of the collection). Add the associated type to it typealias BaseCollection = [Int] and remove the subscript requirement for set.

    protocol ArrayProtocol: RandomAccessCollection where Index == BaseCollection.Index {
        associatedtype BaseCollection: RandomAccessCollection
        var baseCollection: BaseCollection { get set }
        subscript(index: Index) -> Element { get }
    }
    

    extension ArrayProtocol {
        var startIndex: Index { baseCollection.startIndex }
        var endIndex: Index { baseCollection.endIndex }
        func formIndex(after i: inout Index) { baseCollection.index(after: i) }
        func formIndex(before i: inout Index) { baseCollection.index(before: i) }
    }
    

    struct Dataset: ArrayProtocol {
        typealias BaseCollection = [Int]
        var baseCollection: BaseCollection = [ ]
        subscript(index: Int) -> Int { baseCollection[index] }
    }
    


    Note that if you would like to keep the set requirement subscript(index: Index) -> Element { get set } you would need also to make sure BaseCollection conforms to MutableCollection as well.

    subscript(index: Int) -> BaseCollection.Element {
        get { baseCollection[index] }
        set { baseCollection[index] = newValue } 
    }