//=== BinaryTreeNode.swift

public final class BinaryTreeNode<T>: BinaryTreeNodeProtocol {
    public let payload: T
    public let left: BinaryTreeNode<T>?
    public let right: BinaryTreeNode<T>?

    init(
        _ payload: T,
        _ left: BinaryTreeNode<T>? = nil,
        _ right: BinaryTreeNode<T>? = nil
        ) {
        self.payload = payload
        self.left = left
        self.right = right
    }
}

//=== BinaryTreeNodeProtocol.swift

public protocol BinaryTreeNodeProtocol {
    associatedtype Payload
    var payload: Payload { get }
    var left: Self? { get }
    var right: Self? { get }
}

//=== BinaryPreorderIterator.swift

public extension BinaryTreeNodeProtocol {
    var preorderTraversal: BinaryTreePreorderSequence<Self> { return BinaryTreePreorderSequence(self) }
}

public struct BinaryTreePreorderSequence<Tree: BinaryTreeNodeProtocol>: Sequence {
    private let tree: Tree

    public init(_ tree: Tree) { self.tree = tree }

    public func makeIterator() -> BinaryPreorderIterator<Tree> {
        return BinaryPreorderIterator(tree: self.tree)
    }
}

public struct BinaryPreorderIterator<Tree: BinaryTreeNodeProtocol>: IteratorProtocol {
    private var backtrackStack: [(node: Tree, nextAction: NextAction)]

    private enum NextAction { case visitRoot, visitLeft, visitRight, returnToParent }

    public init(tree: Tree) {
        backtrackStack = [(tree, .visitRoot)]
    }

    public mutating func next() -> Tree.Payload? {
        // Using a non-constant boolean expression circumvent's the compilers warnings about the value always being false.
        let debugPrinting = Bool("false")!

        if debugPrinting { print("next()") }

        while let popped = backtrackStack.last {

            if debugPrinting {
                print("tloop - backtrackStack:")
                for (node, nextAction) in backtrackStack.dropLast() {
                    print("tt| (node.payload) (nextAction)")
                }
                print("tt↳ (popped.node.payload) (popped.nextAction) <--- current staten")
            }

            advanceState()

            switch (popped.nextAction) {
            case .visitRoot:
                let returnValue = popped.node.payload
                if debugPrinting { print("treturning (returnValue)n") }
                return returnValue

            case .visitLeft:
                // This block could be more nicely expressed as the following code,
                // but that doesn't work, because the second method call is unconditional,
                // unlike an if/else-if pair.
                //
                // tryMarkingLeftToVisitNext()
                // tryMarkingRightToVisitNext()

                if let leftChild = lastState().node.left {
                    self.visitNext(leftChild)
                }
                else if let rightChild = lastState().node.right {
                    self.visitNext(rightChild)
                }

            case .visitRight:
                tryMarkingRightToVisitNext()

            case .returnToParent:
                break
            }
        }

        return nil
    }

    private func assertLastStateExists() {
        precondition(!backtrackStack.isEmpty, """
        lastState() should only be called when there IS a last state, i.e. the backtrace stack isn't empty.
        """)
    }

    private func lastState() -> (node: Tree, nextAction: NextAction) {
        assertLastStateExists()
        return backtrackStack.last!
    }

    private mutating func setNextActionOnLastNode(_ newNextAction: NextAction) {
        assertLastStateExists()
        backtrackStack[backtrackStack.count - 1].nextAction = newNextAction // gross, but Array.last is write-only
    }

    private mutating func advanceState() {
        switch lastState().nextAction {
        case .visitRoot:  setNextActionOnLastNode(.visitLeft)
        case .visitLeft:  setNextActionOnLastNode(.visitRight)
        case .visitRight: setNextActionOnLastNode(.returnToParent)
        case .returnToParent: backtrackStack.removeLast()
        }
    }

    private mutating func tryMarkingLeftToVisitNext() {
        if let leftChild = lastState().node.left {
            self.visitNext(leftChild)
        }
    }

    private mutating func tryMarkingRightToVisitNext() {
        if let rightChild = lastState().node.right {
            self.visitNext(rightChild)
        }
    }

    private mutating func visitNext(_ node: Tree) {
        backtrackStack.append((node: node, nextAction: .visitRoot))
    }
}

func sampleTree() -> BinaryTreeNode<Int> {
    typealias B = BinaryTreeNode

    //      4
    //    /
    //   2     6
    //  /    /
    // 1   3 5   7
    return B(4, B(2, B(1), B(3)), B(6, B(5), B(7)))
}

print(Array(sampleTree().preorderTraversal))