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>)
}