Swiftで多次元配列

SwiftのenumはSwiftの言語仕様の中で最高

SwiftのEnumは強力だ。
Swiftのenumとswitch文はSwiftの言語仕様の中でも最も気に入っている。
このenumの存在がなければSwiftなんかとうの昔に飽きている。というくらい強力。
今日はこれを使って複数次元を混合した配列を実装してみた。

MultiDimensionalArray.swift

// MARK: Definition

public enum MultiDimensionalArray<T> {
    case val(T)
    case ary([Self])
}

// MARK: Convinient Initializers

public extension MultiDimensionalArray {

    static func ary(_ values: T...) -> Self {
        .ary(values.map(Self.val))
    }

    static func ary(_ values: Self...) -> Self {
        .ary(values)
    }
}

// MARK: Functional

public extension MultiDimensionalArray {

    func map<U>(_ transform: (T) throws -> U) rethrows -> MultiDimensionalArray<U> {
        switch self {
        case .val(let v):
            return try .val(transform(v))
        case .ary(let a):
            return try .ary(a.map { try $0.map(transform) })
        }
    }

    func flatMap<U>(_ transform: (T) throws -> MultiDimensionalArray<U>) rethrows -> MultiDimensionalArray<U> {
        switch self {
        case .val(let v):
            return try transform(v)
        case .ary(let a):
            return try a
                .map { try $0.flatMap(transform) }
                .reduce(.empty, +)
        }
    }

    func reduce<Result>(_ initialResult: Result, _ nextPartialResult: (Result, T) throws -> Result) rethrows -> Result {
        switch self {
        case .val(let v):
            return try nextPartialResult(initialResult, v)
        case .ary(let a):
            return try a.reduce(initialResult) { accum, next in
                try next.reduce(accum, nextPartialResult)
            }
        }
    }
}

// MARK: Monoid

public extension MultiDimensionalArray {

    static func + (lhs: Self, rhs: Self) -> Self {
        switch (lhs, rhs) {
        case (.ary(let a), .ary(let b)):
            return .ary(a + b)
        case (_, .val(_)):
            return lhs + .ary([rhs])
        case (.val(_), _):
            return .ary([lhs]) + rhs
        }
    }

    static var empty: Self {
        .ary([])
    }
}

// MARK: Properties

public extension MultiDimensionalArray {

    var count: Int {
        switch self {
        case .val(_):
            return 1
        case .ary(let a):
            return a.count
        }
    }

    var flatCount: Int {
        switch self {
        case .val(_):
            return 1
        case .ary(let a):
            return a.map(\.flatCount).reduce(0, +)
        }
    }

    var depth: Int {
        switch self {
        case .val(_):
            return 0
        case .ary(let a):
            return (a.map(\.depth).max() ?? 0) + 1
        }
    }

    var flatten: Self {
        flatMap(Self.val)
    }

    var flattenArray: [T] {
        flatten.map { [$0] }.reduce([], +)
    }
}

// MARK: Index Accessibility

extension MultiDimensionalArray {

    subscript(_ index: Int) -> Self {
        self[safe: index]!
    }

    subscript(_ indices: [Int]) -> Self {
        self[safe: indices]!
    }

    subscript(safe index: Int) -> Self? {
        self[safe: [index]]
    }

    subscript(safe indices: [Int]) -> Self? {
        switch (self, indices.splitted) {
        case (.ary(let a), .some(let t)):
            return a[t.head][safe: t.tail]
        case (_, .none):
            return self
        default:
            return .none
        }
    }

    subscript(_ indices: Int...) -> Self {
        self[indices]
    }

    subscript(safe indices: Int...) -> Self? {
        self[safe: indices]
    }
}

// MARK: Description

extension MultiDimensionalArray: CustomStringConvertible {

    public var description: String {
        switch self {
        case .val(let v):
            return "\(v)"
        case .ary(let a):
            return "[\(a.map(String.init).joined(separator: ", "))]"
        }
    }
}

// MARK: Equatable

extension MultiDimensionalArray: Equatable where T: Equatable {}

private extension Array {

    var splitted: (head: Element, tail: [Element])? {
        guard let f = first else { return .none }
        return (f, dropFirst().map { $0 })
    }
}

使い方を説明するのも面倒なのでサンプルコードを書いてみた。

TEST

let a: MultiDimensionalArray<Int> = .ary([.val(1), .val(2), .ary([.val(3), .val(4), .ary([.val(5), .val(6), .ary([]), .val(7)])]), .val(8)])
let b: MultiDimensionalArray<String> = .ary(.ary("apple", "banana"), .ary(.ary("gorilla", "ziraph")), .val("gorilla-gorilla-gorilla"))
let c: MultiDimensionalArray<String> = .val("zebra")

print(a)
print(a.map { $0 * 2 })
print(a.flatMap { .val($0 * 2) })
print(a.flatMap { .ary([.val($0), .val($0 * 2)]) })
print(a.flatMap { .ary([.ary([.val($0), .val($0 * 2)])]) })
print(MultiDimensionalArray.val(10).reduce(5, +))
print(a.count)
print(a.flatCount)
print(a.depth)
print(MultiDimensionalArray<Int>.ary([.ary([])]).depth)
print(MultiDimensionalArray.val(true).flatMap(MultiDimensionalArray.val))
print(MultiDimensionalArray.ary([.val(false)]).flatMap(MultiDimensionalArray.val))
print(b.map { $0.uppercased() }.flattenArray)
print(c.map { $0.uppercased() }.flattenArray)
print(b[1, 0, 1])
print(a == a)
print(b != (b + b))

ちなみに出力はこんな感じ。

Console

[1, 2, [3, 4, [5, 6, [], 7]], 8]
[2, 4, [6, 8, [10, 12, [], 14]], 16]
[2, 4, 6, 8, 10, 12, 14, 16]
[1, 2, 2, 4, 3, 6, 4, 8, 5, 10, 6, 12, 7, 14, 8, 16]
[[1, 2], [2, 4], [3, 6], [4, 8], [5, 10], [6, 12], [7, 14], [8, 16]]
15
4
8
4
2
true
[false]
["APPLE", "BANANA", "GORILLA", "ZIRAPH", "GORILLA-GORILLA-GORILLA"]
["ZEBRA"]
ziraph
true
true

おまけ

ちなみに以前、enumで普通に配列を実装してみた。
こんな感じだった。

List.swift

// MARK: List
public enum List<Element> {
    case empty
    indirect case cons(Element, List<Element>)
}

// MARK: - Initializers
public extension List {
    init(_ elements: Element...) {
        self = .init(from: elements)
    }

    init<T: Collection>(from collection: T) where Element == T.Element {
        self = collection.reversed().reduce(.empty) { .cons($1, $0) }
    }
}

// MARK: - Implement Collection (Essential)
extension List: Collection {
    public func index(after i: Int) -> Int { i + 1 }

    public var startIndex: Int { 0 }

    public var endIndex: Int { count }

    public var count: Int {
        switch self {
            case .cons(_, let xs):
                return xs.count + 1
            case .empty:
                return 0
        }
    }
}

// MARK: - Implement Collection (Additional)
extension List {
    public __consuming func dropFirst(_ k: Int = 1) -> List<Element> {
        return self[from: k]
    }

    public __consuming func reversed() -> List<Element> {
        reduce(.empty) { .cons($1, $0) }
    }

    public func map<T>(_ transform: (Element) throws -> T) rethrows -> List<T> {
        guard case let .cons(x, xs) = self else { return .empty }
        do { return .cons(try transform(x), try xs.map(transform)) }
        catch { return .empty }
    }

    public func flatMap<T>(_ transform: (Element) throws -> List<T>) rethrows -> List<T> {
        guard case let .cons(x, xs) = self else { return .empty }
        do { return try transform(x) + xs.flatMap(transform) }
        catch { return .empty }
    }
}

// MARK: - Implement subscript
public extension List {
    subscript(index: Int) -> Element {
        let elem = self[safe: index]
        precondition(elem != nil, "Out of bounds")
        return elem!
    }

    subscript(safe index: Int) -> Element? {
        guard case let .cons(x, xs) = self else { return .none }
        return index > 0 ? xs[safe: index - 1] : x
    }


    subscript(from index: Int) -> List<Element> {
        guard index > 0 else { return self }
        guard case let .cons(_, xs) = self else { return .empty }
        return xs[from: index - 1]
    }
}

// MARK: - Implement Equatable
extension List: Equatable where Element: Equatable {
    public static func == (lhs: List<Element>, rhs: List<Element>) -> Bool {
        lhs.count == rhs.count && zip(lhs, rhs).reduce(true) { $0 && $1.0 == $1.1 }
    }
}

// MARK: - Implement CustomStringConvertible
extension List: CustomStringConvertible {
    public var description: String {
        return map { "\($0)" }.joined(separator: ", ")
    }
}

// MARK: - Operator Support
public func +<T>(lhs: List<T>, rhs: List<T>) -> List<T> {
    guard case let .cons(x, xs) = lhs else { return rhs }
    if case .empty = xs { return .cons(x, rhs) }
    return .cons(x, xs + rhs)
}

複数の要素を入れておく箱としてなら最初のこの４行だけで十分完成している。

public enum List<Element> {
    case empty
    indirect case cons(Element, List<Element>)
}