情報系人間のブログ

プログラミング、開発に関することを書いていきます。

Swiftで関数型データ構造

Swiftで関数型のデータ構造を作ってみます。 関数型データの構造は永続的であるという特徴があります。なので更新する際には既存のデータを破壊するのではなく新しいオブジェクトを作ります。Enumを使ってこのリストを作ってみます。 consはリストの先頭、nilはリストの終端を表します。

enum List<T: Equatable> {
    case `nil`
    indirect case cons(T, List<T>)
}

またリストは比較できると便利なのでEquatableを満たすようにします。

extension List: Equatable {

    static func ==(lhs: List, rhs: List) -> Bool {
        switch (lhs, rhs) {
        case (.`nil`, .`nil`):
            return true
        case (.cons(let h, let t), .cons(let h1, let t1)):
            return h == h1 && t == t1
        default:
            return false
        }
    }
}

let list = List<Int>.cons(1, .cons(2, .cons(3, .`nil`)))
let list1 = List<Int>.cons(1, .cons(2, .cons(3, .`nil`)))
let list2 = List<Int>.cons(1, .cons(2, .cons(3, .cons(4, .`nil`))))
list == list1 // true
list == list2 // false

CustomStringConvertible を満たすようにすると出力が見やすくなるのでやっておきましょう

extension List: CustomStringConvertible {

    var description: String {
        switch self {
        case .`nil`:
            return "nil"
        case .cons(let head, let tail):
            return "\(head),\(tail)"
        }
    }
}

let list = List<Int>.cons(1, .cons(2, .cons(3, .`nil`)))
print(list) //1,2,3,nil

リストなら以下のような要素の数を取得できるようにしたいですね。やってみましょう。

extension List {

    var isEmpty: Bool {
        switch self {
        case .`nil`:
            return true
        case .cons:
            return false
        }
    }

    var head: T? {
        switch self {
        case .`nil`:
            return nil
        case .cons(let head, _):
            return head
        }
    }

    var tail: List<T>? {
        switch self {
        case .`nil`:
            return nil
        case .cons(_, let tail):
            return tail
        }
    }

    var length: Int {
        switch self {
        case .`nil`:
            return 0
        case .cons(_, let tail):
            return 1 + tail.length
        }
    }

    var last: T? {
        switch self {
        case .`nil`:
            return nil
        case .cons(let head, .`nil`):
            return head
        case .cons(_, let tail):
            return tail.last
        }
    }

}

let list = List<Int>.cons(1, .cons(2, .cons(3, .`nil`)))
list.isEmpty // false
list.head // 1
list.tail // 2,3,nil
list.length // 3
list.last // 3

要素の追加、削除はできて当然ですね。これも実装していきます。 だいたいやっている事は同じです。リストの要素を辿っていき該当の要素に対して処理を行っています。

extension List {

    func add(value: T) -> List<T> {
        switch self {
        case .`nil`:
            return List<T>.cons(value, .`nil`)
        case .cons(let head, let tail):
            return .cons(head, tail.add(value: value))
        }
    }

    func remove(index: Int) -> List<T> {
        switch (self, index) {
        case (.`nil`, _):
            return self
        case (.cons(_, let tail), 0):
            return tail
        case (.cons(let head, let tail), let index):
            return .cons(head, tail.remove(index: index - 1))
        }
    }

    func update(index: Int, value: T) -> List<T> {
        switch (self, index) {
        case (.`nil`, _):
            return self
        case (.cons(_, let tail), 0):
            return .cons(value, tail)
        case (.cons(let head, let tail), let index):
            return .cons(head, tail.update(index: index - 1, value: value))
        }
    }

    func append(list: List<T>) -> List<T> {
        switch self {
        case .`nil`:
            return list
        case .cons(let head, let tail):
            return List.cons(head, tail.append(list: list))
        }
    }

}

let list = List<Int>.cons(1, .cons(2, .cons(3, .`nil`)))
list.add(value: 100) // 1,2,3,100,nil
list.remove(index: 1) // 1,3,nil
list.update(index: 0, value: 100) // 100,2,3,nil
list.append(list: list) // 1,2,3,1,2,3,nil

最後に高階関数を実装します。

extension List {

  func foldr<R>(acc: R, f: (T,R) -> R) -> R {
      switch self {
      case .`nil`:
          return acc
      case .cons(let head, let tail):
          return f(head, tail.foldr(acc: acc, f: f))
      }
  }

  func foldl<R>(acc: R, f: (R,T) -> R) -> R {
      switch self {
      case .`nil`:
          return acc
      case .cons(let head, let tail):
          return tail.foldl(acc: f(acc, head), f: f)
      }
  }

  func map<R>(f: @escaping (T) -> R) -> List<R> {
      return foldr(acc: List<R>.`nil`) { acc, list in
          return List<R>.cons(f(acc), list)
      }
  }

  func filter(f: @escaping (T) -> Bool) -> List<T> {
      return foldr(acc: List<T>.`nil`, f: { acc, list in
          if f(acc) {
              return List<T>.cons(acc, list)
          }
          return list
      })
  }

  func flatMap<R>(f: @escaping (T) -> List<R>) -> List<R> {
      return foldr(acc: List<R>.`nil`, f: { acc, list in
          return f(acc).append(list: list)
      })
  }

}

let list = List<Int>.cons(1, .cons(2, .cons(3, .`nil`)))
list.map { $0 * 10 } // "10,20,30,nil
list.filter { $0 % 2 == 0 } // 2,nil
list.flatMap { .cons($0, .cons($0, .`nil`)) } // 1,1,2,2,3,3,nil