[swift] Collectionの添字に演算子を渡す

検証環境

$ swift --version
Apple Swift version 4.1 (swiftlang-902.0.48 clang-902.0.37.1)
Target: x86_64-apple-darwin17.6.0

Collectionの添字に ... 演算子を渡す事ができる

一部のCollection型に、 (UnboundedRange_)->() という関数をサブスクリプトの引数に取るメソッドが実装されていることを知りました。

Array<Element>

Array

subscript(x: (UnboundedRange_) -> ()) -> ArraySlice<Element> { get }

String

String

subscript(x: (UnboundedRange_) -> ()) -> Substring { get }

これによって以下のように、添字に ... 演算子を指定することで、
簡単にすべての要素を含むサブシーケンスを取得することができます。

サブシーケンスの取得

func count(_ sub: Substring) -> Int {
    return sub.count;
}

print(type(of: "hello"[...])) // Substring
print(count("hello"[...]))    // 5

添字に ... 演算子を指定できる理由

これは　UnboundedRange_ 型内で、 ... 演算子がオーバーロードされているためです。

UnboundedRange

public enum UnboundedRange_ {
  public static postfix func ... (_: UnboundedRange_) -> () {
    fatalError("uncallable")
  }
}

上記の ... 演算子を呼び出すと、内部でfatalErrorが呼び出されてクラッシュします。
ここで重要なことは、... 演算子が (UnboundedRange_) -> () 型であることです。

これによって呼び出し側の型で、subscript (_: (UnboundedRange)->()) メソッドをオーバーロードすることによって、 ... 演算子が引数に指定されると、オーバーロードしたメソッドが呼び出されるという寸法です。

実際には、Collectionでオーバーロードされています。

github 実際の実装

Range.swift

public typealias UnboundedRange = (UnboundedRange_)->()

extension Collection {
    public subscript(x: UnboundedRange) -> SubSequence {
        return self[startIndex...]
    }
}

ということは以下のように (Int)->(Int) を受け取るsubscriptをArray<Element>でオーバーロードするだけで、 + や - 演算子を引き渡すことによって、オーバーロードしたメソッドを呼び出せるはずです。

Array+.swift

extension Array where Element == Int {
    subscript (_ op: (Int) -> (Int)) -> [Int] {
        return self.map {
            return op($0)
        }
    }
}

let arr: [Int] = Array(1..<4)
print(arr[+]) // OK [1, 2, 3]
print(arr[-]) // OK [-1, -2, -3] 

問題なくオーバーロードしたメソッドが呼び出されました。

カスタムの演算子を添え字に指定できるようにする

... 演算子と同様に、 添字に ..< 演算子を指定すると、内部でdropLastを呼び出すsubscriptメソッドをオーバーロードいたしました。

DroppedLastRange.swift

postfix operator ..< // ①

enum DroppedLastRange_ {
    public static postfix func ..<(_: DroppedLastRange_) -> () {
        fatalError("uncallable")
    }
}

// ②
typealias DroppedLastRange = (DroppedLastRange_) -> ()

extension Array {
    // ③ subscriptメソッドをオーバーロード
    subscript (_ index: DroppedLastRange) -> ArraySlice<Element> {
        return self.dropLast()
    }

}

print(([] as [Int])[..<]) // [] 
print([0][..<])           // []
print([0,1][..<])         // [0]
print([0,1,2][..<])       // [0,1]

Note:

① ... 演算子も postfix なので合わせました
② UnboundedRange_ -> () も同様に UnboundedRange と型名がつけられています
③ subscriptメソッドをオーバーロードすることによって、添字に ..< が指定されると呼び出されます

まとめ

改めて 型 の使い方を勉強させていただきました。

参考サイト

[Swift] UnboundedRangeを“呼ぶ”
https://qiita.com/YOCKOW/items/dd409b9588f4be72f58f

Unbounded Ranges in Swift 4
https://tonisuter.com/blog/2017/08/unbounded-ranges-swift-4/