Search code examples
arraysswiftcollectionsiterator

How to return a generic collection or iterator in Swift?


I'm trying to do something like this:

func fullOrSubarray(array: [Int]) -> ReturnType {
    return array.count < 10 ? array : array[0..<10]
}

What should go I write into ReturnType as I'm trying to return Array<Int> or ArraySlice<Int>?

I know I can simply convert ArraySlice<Int> to Array<Int> (or vice versa), but I don't want that for performance reasons. I want to return some generic collection or iterator type, but Swift doesn't seem to have one. Do I have to write a custom iterator for a thing like this? Other languages have all sorts of generic iterators so I can return Array, Set, etc interchangeably.


Solution

  • It's possible to do what you're describing, but first, you should see if you can avoid it. The specific example you've given is exactly prefix and should definitely be implemented as:

    func fullOrSubarray(array: [Int]) -> ArraySlice<Int> {
        array.prefix(10)
    }
    

    The correct thing to return here is an ArraySlice. There is no reason to return something else. Take a look at the implementation of prefix to see how stdlib handles this problem.

    That said, in some rare cases it may be useful to do what you're describing. For example, the underlying data may be in a Set or an Array, and you want to avoid copies. In that case, the correct tool is AnySequence or AnyCollection:

    func fullOrSubarray(array: [Int]) -> AnySequence<Int> {
        array.count < 10 ? AnySequence(array) : AnySequence(array[0..<10])
    }
    

    But this is a very unusual situation, and would never be used for Array and ArraySlice. For that, just use ArraySlice, or convert back to Array. Do not assume that converting from a slice to an Array is expensive. The stdlib is very clever about internal copy-on-write Array buffers and can often avoid making any copies.

    Even in cases where the data might currently be a Set, unless it is quite large or accessed quite often, it is generally better to just convert it to an Array before returning it. If it's accessed as an Array quite often, I'd lean towards storing it as an Array rather than a Set in the first place.

    But if all those fail you, and you need to return an unknown sequence type, that's what AnySequence (or AnyCollection) is for.