Trie.java

package com.googlecode.totallylazy.collections;

import com.googlecode.totallylazy.Option;
import com.googlecode.totallylazy.Segment;
import com.googlecode.totallylazy.Value;
import com.googlecode.totallylazy.annotations.tailrec;

import static com.googlecode.totallylazy.Option.none;
import static com.googlecode.totallylazy.Option.option;

public class Trie<K, V> implements Value<V> {
    private final Option<V> value;
    private final PersistentMap<K, Trie<K, V>> children;

    private Trie(Option<V> value, PersistentMap<K, Trie<K, V>> children) {
        this.value = value;
        this.children = children;
    }

    public static <K, V> Trie<K, V> trie() {
        return trie(Option.<V>none());
    }

    public static <K, V> Trie<K, V> trie(Option<V> value) {
        return trie(value, ListMap.<K, Trie<K, V>>emptyListMap());
    }

    public static <K, V> Trie<K, V> trie(Option<V> value, PersistentMap<K, Trie<K, V>> children) {
        return new Trie<K, V>(value, children);
    }

    public boolean contains(Segment<K> key){
        if(key.isEmpty()) return !value.isEmpty();
        Option<Trie<K, V>> child = childFor(key);
        return !child.isEmpty() && child.get().contains(key.tail());
    }

    @tailrec
    public Option<V> get(Segment<K> key) {
        if(key.isEmpty()) return value;
        Option<Trie<K, V>> child = childFor(key);
        if(child.isEmpty()) return none();
        return child.get().get(key.tail());
    }

    public Trie<K, V> put(Segment<K> key, V value) {
        if(key.isEmpty()) return trie(option(value), children);
        return trie(this.value, children.insert(key.head(), childFor(key).getOrElse(Trie.<K, V>trie()).put(key.tail(), value)));
    }

    public Trie<K, V> remove(Segment<K> key) {
        return put(key, null);
    }

    public V value() {
        return value.get();
    }

    public boolean isEmpty(){
        return value.isEmpty() && children.isEmpty();
    }

    private Option<Trie<K, V>> childFor(Segment<K> key) {return children.lookup(key.head());}
}