equalTrees.swift

//
//  NaiveTree.swift
//  immutable_nary_tree
//
//  Created by zephyz on 18.02.17.
//  Copyright © 2017 moe.zephyz. All rights reserved.
//

import Foundation

public indirect enum Tree<Element> {
    case leaf
    case node(Element, [Tree])
}

extension Tree {
    func map<T>(_ transform: (Element) -> T) -> Tree<T> {
        switch self {
        case .leaf: return .leaf
        case .node(let e, let content): return .node(transform(e), content.map {$0.map(transform)})
        }
    }
}

public func treeEqual<E: Equatable> (_ lhs: Tree<E>, _ rhs: Tree<E>) -> Bool {
    let zipped = lhs.zipUnion(rhs)
    let mappedEqual: Tree<Bool> = zipped.map { tuple in
        switch tuple {
        case let (l?, r?): return l == r
        case _: return false
        }
    }
    let combine: (Bool, Bool) -> Bool = { $0 && $1 }
    return mappedEqual.reduce(initial: true)(combine)
}


extension Tree {
    func reduce<T>(initial: T) -> ((T, Element) -> T) -> T {
        return { combine in
            switch self {
            case .leaf: return initial
            case .node(let elem, let content):
                let combined = combine(initial, elem)
                let reduced = content.reduce(combined, { (acc, tree) -> T in
                    tree.reduce(initial: acc)(combine)
                })
                return reduced
            }
        }
    }
}

extension Tree {
    func zipUnion<T>(_ tree: Tree<T>) -> Tree<(Element?, T?)> {
        switch (self, tree) {
        case (.leaf, .leaf): return .leaf
        case (.node(let elem, let content), .leaf):
            return .node((elem, nil), content.map {$0.zipUnion(.leaf)})
        case (.leaf, .node(let elem, let content)):
            return .node((nil, elem), content.map {(Tree<Element>.leaf).zipUnion($0)})
        case (.node(let leftElem, let leftContent), .node(let rightElem, let rightContent)):
            let zipped = leftContent.zipUnion(rightContent)
            let asTree = zipped.map { ($0 ?? .leaf, $1 ?? .leaf) }
            return .node((leftElem, rightElem), asTree.map {$0.zipUnion($1)})
        }
    }
}

extension Collection where Indices.Iterator.Element == Index {
    
    /// Returns the element at the specified index iff it is within bounds, otherwise nil.
    subscript (safe index: Index) -> Generator.Element? {
        return indices.contains(index) ? self[index] : nil
    }
}

extension Array {
    
    func zipUnion<T>(_ array: [T]) -> [(Element?, T?)] {
        var acc: [(Element?, T?)]  = []
        let m = self.count < array.count ? array.count : self.count
        for index in 0..<m {
            let left = self[safe: index]
            let right = array[safe: index]
            acc.append((left, right))
        }
        return acc
    }
}