ImmutableTree.java

import com.google.common.collect.ImmutableList;
import java.util.Collection;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.Optional;
import java.util.function.Function;
import java.util.function.Predicate;
import java.util.function.UnaryOperator;
import java.util.stream.Collectors;
import java.util.stream.Stream;
import org.immutables.value.Value;
import org.immutables.value.Value.Immutable;

@Immutable
public interface Tree<T> {
  @Value.Parameter(order = 1)
  T value();

  @Value.Parameter(order = 2)
  List<Tree<T>> _children();
  
  @Value.Parameter(order = 3)
  @Value.Auxiliary // DO NOT INCLUDE in equals, hashcode, tostring to avoid stackoverflow
  Optional<Tree<T>> parent();

  @Value.Derived
    // Compute children only once to add this as parent
  default List<Tree<T>> children() {
    final UnaryOperator<Tree<T>> addThisParent =
        t -> ImmutableTree.copyOf(t).withParent(this);
    final List<Tree<T>> childrenWithThisParent =
        _children().stream().map(addThisParent).collect(Collectors.toList());
    return childrenWithThisParent;
  }  

  static <T> Tree<T> of(final T value) {
    return ImmutableTree.of(value, Collections.emptyList(), Optional.empty());
  }

  static <T> Tree<T> of(final T value, final List<Tree<T>> children) {
    return ImmutableTree.of(value, children, Optional.empty());
  }

  static <T> Tree<T> of(
      final T value, final List<Tree<T>> children, final Optional<Tree<T>> parent) {
    return ImmutableTree.of(value, children, parent);
  }

  default Optional<Tree<T>> findNode(final T value) {
    if (value == null) {
      return Optional.empty();
    }
    if (value.equals(value())) {
      return Optional.of(this);
    }
    final Optional<Tree<T>> node =
        children().stream()
            .map(t -> t.findNode(value))
            .filter(Optional::isPresent)
            .findAny()
            .orElse(Optional.empty());
    return node;
  }

  default List<T> parentValues(final T value) {
    return findNode(value).map(t -> parentValues(t)).orElse(Collections.emptyList());
  }

  private static <T> List<T> parentValues(final Tree<T> node) {
    if (node.parent().isEmpty()) {
      return Collections.emptyList();
    }
    final Optional<Tree<T>> seed = node.parent();
    final Predicate<Optional<Tree<T>>> hasNext = Optional::isPresent;
    final UnaryOperator<Optional<Tree<T>>> next =
        optionalTree -> optionalTree.flatMap(Tree::parent);
    final Stream<Tree<T>> parentStream =
        Stream.iterate(seed, hasNext, next).map(Optional::get);
    return collectValues(parentStream);
  }

  private static <T> List<T> collectValues(final Stream<Tree<T>> iterate) {
    return iterate.map(Tree::value).collect(Collectors.toList());
  }

  default List<T> childrenValues() {
    final Tree<T> node = this;
    return childrenValues(node);
  }

  default List<T> childrenValues(final T value) {
    return findNode(value).map(t -> childrenValues(t)).orElse(Collections.emptyList());
  }

  private static <T> List<T> childrenValues(final Tree<T> node) {
    final List<Tree<T>> seed = node.children();
    final Predicate<List<Tree<T>>> hasNext = Predicate.not(List::isEmpty);
    final UnaryOperator<List<Tree<T>>> next = Tree::flattenChildrenList;
    final Stream<Tree<T>> childrenStream =
        Stream.iterate(seed, hasNext, next).flatMap(List::stream);
    return collectValues(childrenStream);
  }

  private static <T> List<Tree<T>> flattenChildrenList(final List<Tree<T>> children) {
    return children.stream()
        .map(Tree::children)
        .flatMap(List::stream)
        .collect(Collectors.toList());
  }

  static <T, ID> List<Tree<T>> build(
      final Collection<T> source,
      final Function<T, ID> idMapper,
      final Function<T, Optional<ID>> parentMapper) {
    Objects.requireNonNull(source, "source cannot be null");
    Objects.requireNonNull(source, "idMapper cannot be null");
    Objects.requireNonNull(source, "parentMapper cannot be null");
    if (source.isEmpty()) {
      return Collections.emptyList();
    }
    final Map<Optional<ID>, List<T>> byParent =
        source.stream().collect(Collectors.groupingBy(parentMapper));
    final List<T> roots =
        Optional.ofNullable(byParent.get(Optional.empty())).orElse(Collections.emptyList());
    return build(roots, idMapper, byParent);
  }

  private static <T, ID> List<Tree<T>> build(
      final Collection<T> roots,
      final Function<T, ID> idMapper,
      final Map<Optional<ID>, List<T>> byParent) {
    Objects.requireNonNull(roots, "roots cannot be null");
    Objects.requireNonNull(byParent, "byParent cannot be null");
    if (roots.isEmpty()) {
      return Collections.emptyList();
    }
    final Function<T, List<T>> getChildren =
        root ->
            Optional.ofNullable(byParent.get(Optional.ofNullable(idMapper.apply(root))))
                .orElse(Collections.emptyList());
    final List<Tree<T>> myTrees =
        roots.stream()
            .map(root -> Tree.of(root, build(getChildren.apply(root), idMapper, byParent)))
            .collect(ImmutableList.toImmutableList());
    return myTrees;
  }

  default String draw() {
    final StringBuilder builder = new StringBuilder();
    draw("", builder);
    return builder.toString();
  }

  private void draw(final String indent, final StringBuilder builder) {
    builder.append(value());
    final List<Tree<T>> seed = children();
    final Predicate<List<Tree<T>>> hasNext = Predicate.not(List::isEmpty);
    final UnaryOperator<List<Tree<T>>> next = list -> list.subList(1, list.size());
    Stream.iterate(seed, hasNext, next)
        .forEach(
            index -> {
              final boolean isLast = index.subList(1, index.size()).isEmpty();
              builder.append('\n').append(indent).append(isLast ? "└──" : "├──");
              index.get(0).draw(indent + (isLast ? "   " : "│  "), builder);
            });
  }
}