io.datakernel.di.util.Trie Maven / Gradle / Ivy

Go to download
package io.datakernel.di.util;

import org.jetbrains.annotations.Nullable;

import java.util.Arrays;
import java.util.Collection;
import java.util.HashMap;
import java.util.Map;
import java.util.function.BiConsumer;
import java.util.function.Consumer;
import java.util.function.Function;
import java.util.stream.Stream;

import static io.datakernel.di.util.Utils.next;
import static java.util.Collections.emptyMap;

/**
 * Completely generic and abstract simple Java implementation
 * of the prefixed tree (or trie) data structure.
 */
public final class Trie {
	private final V payload;
	private final Map> children;

	public Trie(V payload, Map> children) {
		this.payload = payload;
		this.children = children;
	}

	public static  Trie leaf(V value) {
		return new Trie<>(value, new HashMap<>());
	}

	public static  Trie of(V payload, Map> children) {
		return new Trie<>(payload, children);
	}

	public V get() {
		return payload;
	}

	public Map> getChildren() {
		return children;
	}

	public Trie get(K key) {
		return children.get(key);
	}

	public Trie getOrDefault(K key, V defaultValue) {
		return children.getOrDefault(key, new Trie<>(defaultValue, emptyMap()));
	}

	public Trie computeIfAbsent(K key, Function f) {
		return children.computeIfAbsent(key, k -> leaf(f.apply(k)));
	}

	@Nullable
	public Trie get(K[] path) {
		Trie subtree = this;
		for (K key : path) {
			subtree = subtree.get(key);
			if (subtree == null) {
				return null;
			}
		}
		return subtree;
	}

	public Trie computeIfAbsent(K[] path, Function f) {
		Trie subtree = this;
		for (K key : path) {
			subtree = subtree.computeIfAbsent(key, f);
		}
		return subtree;
	}

	public void addAll(Trie other, BiConsumer merger) {
		mergeInto(this, other, merger);
	}

	public  Trie map(Function fn) {
		Trie root = leaf(fn.apply(payload));
		children.forEach((k, sub) -> root.children.put(k, sub.map(fn)));
		return root;
	}

	public void dfs(K[] path, BiConsumer consumer) {
		Trie sub = get(path);
		if (sub != null) {
			sub.dfsImpl(path, consumer);
		}
	}

	private void dfsImpl(K[] path, BiConsumer consumer) {
		children.forEach((key, child) -> child.dfsImpl(next(path, key), consumer));
		consumer.accept(path, payload);
	}

	public void dfs(Consumer consumer) {
		children.forEach((key, child) -> child.dfs(consumer));
		consumer.accept(payload);
	}

	private static  void mergeInto(Trie into, Trie from, BiConsumer merger) {
		if (into == from) {
			return;
		}
		merger.accept(into.get(), from.get());
		from.children.forEach((scope, child) -> mergeInto(into.children.computeIfAbsent(scope, $ -> child), child, merger));
	}

	public static  Trie merge(BiConsumer merger, V rootPayload, Trie first, Trie second) {
		Trie combined = leaf(rootPayload);
		mergeInto(combined, first, merger);
		mergeInto(combined, second, merger);
		return combined;
	}

	@SafeVarargs
	public static  Trie merge(BiConsumer merger, V rootPayload, Trie first, Trie second, Trie... rest) {
		return merge(merger, rootPayload, Stream.concat(Stream.of(first, second), Arrays.stream(rest)));
	}

	public static  Trie merge(BiConsumer merger, V rootPayload, Collection> bindings) {
		return merge(merger, rootPayload, bindings.stream());
	}

	public static  Trie merge(BiConsumer merger, V rootPayload, Stream> bindings) {
		Trie combined = leaf(rootPayload);
		bindings.forEach(sb -> mergeInto(combined, sb, merger));
		return combined;
	}

	public String prettyPrint() {
		return prettyPrint(0);
	}

	private String prettyPrint(int indent) {
		String indentStr = new String(new char[indent]).replace('\0', '\t');

		StringBuilder sb = new StringBuilder()
				.append("(")
				.append(payload)
				.append(") {");

		if (!children.isEmpty()) {
			sb.append('\n').append(indentStr);
			children.forEach((key, child) -> sb
					.append(indentStr)
					.append('\t')
					.append(key)
					.append(" -> ")
					.append(child.prettyPrint(indent + 1)));
		}
		return sb.append("}\n").toString();
	}

	@Override
	public String toString() {
		return "Trie{payload=" + payload + ", children=" + children + '}';
	}
}