class BinaryTree : IMutableSet { private class TreeNode( var value : T, var parent : TreeNode ) { var left : TreeNode var right : TreeNode } private val compare : MatchableComparison private var root : TreeNode private var version = 0 // override var size : Int { get; private set; } this(compare : Comparison) { this.compare = asMatchableComparison(comparison) } this() : this(naturalOrder()) { } private [operator] fun T.compareTo(other : T) : Int = compare(this, other) override fun contains(item : T) : Boolean { return contains(root, item) fun contains(node : TreeNode, item : T) : Boolean { if (node == null) return false when(compare(item, node.value)) { EQ -> true LS -> contains(node.left, item) GT -> contains(node.right, item) } } } override fun add(item : T) : Boolean { if (add(Ref(root), null)) { size++ version++ return true } return false // In principle, this has access to item anyway, but then it's unreachable code // BAD: the naive implementation of ref will create H(T) ref objects, but can be optimized to create only one fun add(node : Ref, parent : TreeNode) : Boolean { if (node.value == null) { node.value = TreeNode(item, parent) return true } when (compare(item, node.value.value)) { EQ -> false LS -> add(ref node.value.left, node) GT -> add(ref node.value.right, node) } } // In principle, this has access to item anyway fun addNoRef(node : TreeNode) : Boolean { if (node == null) { root = TreeNode(item, null) return true } when (compare(item, node.value)) { EQ -> return false LS -> if (node.left == null) { node.left = TreeNode(item, node) return true } else return add(node.left) GT -> if (node.right == null) { node.right = TreeNode(item, node) return true } else return add(node.right) } } } override fun remove(item : T) : Boolean { val toRemove = find(root, item) if (toRemove == null) return false remove(toRemove) size-- version++ return true fun find(node : TreeNode) : TreeNode { if (node == null) return null when (compare(item, node.value)) { EQ -> node LS -> find(node.left) GT -> find(node.right) } } } private fun remove(node : TreeNode) { when (node) { is TreeNode #(null, null) -> replace(node, null) is TreeNode #(null, right) -> replace(node, right) is TreeNode #(left, null) -> replace(node, left) is TreeNode #(left, right) -> { val min = min(node.right) node.value = min.value remove(min) } } fun replace(node : TreeNode, replace : TreeNode) { if (node == root) { root = replace return } if (node.parent.left == node) { node.parent.left = replace } else if (node.parent.right == node) { node.parent.right = replace } else Assert(false) } } // Relies on tail-recursion optimization private fun min(node : TreeNode) { if (node.left == null) node else min(node.left) } override fun iterator() : IIterator = mutableIterator() override fun mutableIterator() : IMutableIterator = object : IMutableIterator { val version = BinaryTree.this.version val down = Stack() val up = Stack() var lastNode : TreeNode this() { if (root != null) down.push(root) } override fun next() : T { if (!hasNext) throw NoMoreElementsException() checkVersion() lastNode = nextNode() return lastNode.value } private fun nextNode() : TreeNode { while (true) { if (!down.isEmpty) { val curNode = down.pop() if (curNode.left != null) { up.push(curNode) down.push(curNode.left) } else { if (curNode.right != null) { down.push(curNode.right) } return curNode; } } else { val curNode = up.pop() if (curNode.right != null) { down.push(curNode.right) } return curNode } } } override val hasNext : Boolean get() = !down.isEmpty || !up.isEmpty override fun remove() { checkVersion() if (lastNode == null) throw IterationException("Nothing to remove") remove(lastNode) version++ BinaryTree.this.version = version } private fun checkVersion() { if (version != BinaryTree.this.version) { throw ConcurrentModificationException() } } } }