Swift: Adding multiple-index array access

This morning, #swift-lang buzzed about multiple subscripts. “zOMG, you can have subscripts take more than one argument?” Why yes, yes you can. It’s a neat Swift trick, but a (pardon me) a tricky one.

For example, you might create an array and want to index it with foo[3, 5] or foo[7, 9]. Turns out that the general case is slightly trickier than I anticipated. The following demonstrates the desired indexing behavior.

let foo = [1, 2, 3, 4, 5, 6, 7]
println(foo[2]) // prints 3
println(foo[2, 4]) // prints [3, 5]
println(foo[2, 4, 1]) // prints [3, 5, 2]
println(foo[2, 4, 1, 5]) // prints [3, 5, 2, 6]

However, you don’t get this behavior out of the box.

Screen Shot 2014-07-10 at 11.47.55

Backtracking a little, you can easily build a custom subscript when you know a priori the number of elements you’re aiming for. For example, this extension returns two index items at a time. This is easy to implement and Swift can easily use parameter matching to figure out the overloaded results you’re aiming for.

// Two at a time
extension Array {
    typealias ArrayType = Element
    subscript(i1: Int, i2:Int) -> [ArrayType] {
        return [self[i1], self[i2]]
    }
}

It’s harder to extend this behavior for an arbitrary number of arguments. Specifically, you must take care with variadic arguments. The following approach won’t work.

subscript(i1: Int, i2:Int...)

Use that approach and you end up with infinite loops. That’s because Swift cannot distinguish this parameter declaration from an override of the standard single-item indexing. Variadic parameters accept zero or more values of a specified type, so  [2], [2, 4], [2, 4, 6], and [2,4, 6, 8] all match this declaration.

(Int, Int…) is virtually identical to (Int) at runtime, so Swift chooses this override to the original implementation for single-parameter lookups. That’s where the infinite loops come from. To make this work, you need a signature that won’t confuse swift. The proper solution uses two non-variadic arguments followed by a third variadic one, creating a “use this implementation for two or more parameters” scenario.

The following extension provides a final working solution. When you provide at least two subscript arguments, Swift knows to use this subscripting implementation.

extension Array {
    typealias ArrayType = Element
    subscript(i1: Int, i2: Int, rest: Int...) ->  [ArrayType] {
        var result = [self[i1], self[i2]]
            for index in rest {
                result += self[index]
            }
            return result
    }
}

 

Comments are closed.