TreeZipper.java

package com.googlecode.totallylazy.collections;

import com.googlecode.totallylazy.functions.Function1;
import com.googlecode.totallylazy.functions.Functions;
import com.googlecode.totallylazy.Option;
import com.googlecode.totallylazy.Pair;
import com.googlecode.totallylazy.annotations.tailrec;

import static com.googlecode.totallylazy.Option.none;
import static com.googlecode.totallylazy.Option.some;
import static com.googlecode.totallylazy.predicates.Predicates.is;
import static com.googlecode.totallylazy.predicates.Predicates.where;
import static com.googlecode.totallylazy.Sequences.sequence;
import static com.googlecode.totallylazy.collections.TreeZipper.Breadcrumb.breadcrumb;
import static com.googlecode.totallylazy.collections.TreeZipper.Direction.left;
import static com.googlecode.totallylazy.collections.TreeZipper.Direction.right;
import static com.googlecode.totallylazy.collections.TreeZipper.functions.direction;
import static java.lang.String.format;

public class TreeZipper<K, V> implements Zipper<Pair<K, V>> {
    public final TreeMap<K, V> focus;
    public final PersistentList<Breadcrumb<K, V>> breadcrumbs;

    private TreeZipper(TreeMap<K, V> focus, PersistentList<Breadcrumb<K, V>> breadcrumbs) {
        this.focus = focus;
        this.breadcrumbs = breadcrumbs;
    }

    public static <K, V> TreeZipper<K, V> zipper(TreeMap<K, V> focus) {
        return new TreeZipper<K, V>(focus, PersistentList.constructors.<Breadcrumb<K, V>>empty());
    }

    private static <K, V> TreeZipper<K, V> zipper(TreeMap<K, V> focus, PersistentList<Breadcrumb<K, V>> crumbs) {
        return new TreeZipper<K, V>(focus, crumbs);
    }

    public TreeZipper<K, V> left() {
        return zipper(focus.left(), breadcrumbs.cons(breadcrumb(left, focus.head(), focus.right())));
    }

    public TreeZipper<K, V> right() {
        return zipper(focus.right(), breadcrumbs.cons(breadcrumb(right, focus.head(), focus.left())));
    }

    public TreeZipper<K, V> up() {
        final Breadcrumb<K, V> breadcrumb = breadcrumbs.head();
        final TreeMap<K, V> left = breadcrumb.direction == Direction.left ? focus : breadcrumb.other;
        final TreeMap<K, V> right = breadcrumb.direction == Direction.left ? breadcrumb.other : focus;
        return zipper(focus.factory().create(focus.comparator(), breadcrumb.parent.first(), breadcrumb.parent.second(), left, right), breadcrumbs.tail());
    }

    @tailrec
    public TreeZipper<K, V> top() {
        if (breadcrumbs.isEmpty()) return this;
        return up().top();
    }

    public TreeMap<K, V> toTreeMap() {
        return top().focus;
    }

    public TreeZipper<K, V> modify(Function1<? super TreeMap<K, V>, ? extends TreeMap<K, V>> callable) {
        TreeMap<K, V> result = Functions.call(callable, focus);
        TreeZipper<K, V> newZipper = zipper(result, breadcrumbs);
        if (newZipper.focus.isEmpty()) return newZipper.up();
        return newZipper;
    }

    public TreeZipper<K, V> replace(K key, V value) {
        return modify(functions.replace(key, value));
    }

    public TreeZipper<K, V> delete() {
        return remove();
    }

    public TreeZipper<K, V> remove() {
        return modify(functions.<K, V>remove());
    }

    @tailrec
    public TreeZipper<K, V> first() {
        if (isFirst()) return this;
        return left().first();
    }

    @tailrec
    public TreeZipper<K, V> last() {
        if (isLast()) return this;
        return right().last();
    }

    @Override
    public boolean isFirst() {
        return focus.left().isEmpty();
    }

    @Override
    public boolean isLast() {
        return focus.right().isEmpty();
    }

    @Override
    public int index() {
        return focus.indexOf(value()) + breadcrumbs.toSequence().
                filter(where(direction, is(right))).
                fold(0, (integer, breadcrumb) -> integer + breadcrumb.other.size() + 1);
    }

    @Override
    public TreeZipper<K, V> index(int index) {
        int position = index();
        if (position == index) return this;
        if (position < index) return next().index(index);
        return previous().index(index);
    }

    public boolean isTop() {
        return breadcrumbs.isEmpty();
    }

    public TreeZipper<K, V> next() {
        if (focus.right().isEmpty()) return backtrack(right).up();
        return right().first();
    }

    public Option<TreeZipper<K, V>> nextOption() {
        try {
            return some(next());
        } catch (Exception e) {
            return none();
        }
    }

    public TreeZipper<K, V> previous() {
        if (focus.left().isEmpty()) return backtrack(left).up();
        return left().last();
    }

    public Option<TreeZipper<K, V>> previousOption() {
        try {
            return some(previous());
        } catch (Exception e) {
            return none();
        }
    }

    @tailrec
    private TreeZipper<K, V> backtrack(final Direction direction) {
        if (breadcrumbs.head().direction.equals(direction)) return up().backtrack(direction);
        return this;
    }

    @Override
    public String toString() {
        return "TreeZipper{" +
                "focus=" + focus +
                ", breadcrumbs=" + breadcrumbs +
                '}';
    }

    public Pair<K, V> pair() {
        return Pair.pair(focus.key(), focus.value());
    }

    @Override
    public Pair<K, V> value() {
        return focus.head();
    }

    public enum Direction {
        left, right
    }

    public static final class Breadcrumb<K, V> {
        public final Direction direction;
        public final Pair<K, V> parent;
        public final TreeMap<K, V> other;

        private Breadcrumb(Direction direction, Pair<K, V> parent, TreeMap<K, V> other) {
            this.parent = parent;
            this.other = other;
            this.direction = direction;
        }

        public static <K, V> Breadcrumb<K, V> breadcrumb(Direction direction, Pair<K, V> parent, TreeMap<K, V> other) {
            return new Breadcrumb<K, V>(direction, parent, other);
        }

        @Override
        public String toString() {
            return format("direction(%s), parent(%s), other(%s)", direction, parent, other);
        }

        @Override
        public int hashCode() {
            return sequence(direction, parent, other).hashCode();
        }

        @Override
        public boolean equals(Object obj) {
            return obj instanceof Breadcrumb &&
                    ((Breadcrumb) obj).direction.equals(direction) &&
                    ((Breadcrumb) obj).parent.equals(parent) &&
                    ((Breadcrumb) obj).other.equals(other);
        }
    }

    public static class functions extends TreeMap.functions {
        public static Function1<Breadcrumb<?, ?>, Direction> direction = breadcrumb -> breadcrumb.direction;
    }

}