001/*
002 * SPDX-License-Identifier: Apache-2.0
003 *
004 * Copyright 2024-2026 The Enola <https://enola.dev> Authors
005 *
006 * Licensed under the Apache License, Version 2.0 (the "License");
007 * you may not use this file except in compliance with the License.
008 * You may obtain a copy of the License at
009 *
010 *     https://www.apache.org/licenses/LICENSE-2.0
011 *
012 * Unless required by applicable law or agreed to in writing, software
013 * distributed under the License is distributed on an "AS IS" BASIS,
014 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
015 * See the License for the specific language governing permissions and
016 * limitations under the License.
017 */
018package dev.enola.common.tree;
019
020import static java.util.Objects.requireNonNull;
021
022import com.google.common.collect.ImmutableList;
023import com.google.common.collect.ImmutableMap;
024import com.google.errorprone.annotations.CanIgnoreReturnValue;
025import com.google.errorprone.annotations.ImmutableTypeParameter;
026
027import org.jspecify.annotations.NonNull;
028import org.jspecify.annotations.Nullable;
029
030import java.util.*;
031
032public class ImmutableTreeBuilder<@ImmutableTypeParameter N> implements TreeBuilder<N>, Tree<N> {
033
034    private @Nullable N root;
035    private final Set<N> nodes = new HashSet<>();
036    private final Map<@NonNull N, ImmutableList.Builder<N>> map = new HashMap<>();
037
038    @Override
039    @CanIgnoreReturnValue
040    public ImmutableTreeBuilder<N> root(N root) {
041        if (this.root != null) throw new IllegalStateException("Root already set: " + root);
042        this.root = requireNonNull(root, "root");
043        map.put(root, ImmutableList.builder());
044        nodes.add(root);
045        return this;
046    }
047
048    @Override
049    public N root() {
050        if (this.root == null) throw new IllegalStateException("Root not set");
051        return root;
052    }
053
054    @Override
055    @CanIgnoreReturnValue
056    public ImmutableTreeBuilder<N> addChild(N parent, N child) {
057        if (nodes.contains(child))
058            throw new IllegalStateException("Tree already contains child: " + child);
059        var entry = map.get(parent);
060        if (entry == null) throw new IllegalStateException("Parent Node not found: " + parent);
061        entry.add(child);
062        map.put(child, ImmutableList.builder());
063        nodes.add(child);
064        return this;
065    }
066
067    @Override
068    public Tree<N> build() {
069        if (this.root == null) throw new IllegalStateException("Root not set");
070        var builder =
071                ImmutableMap.<@NonNull N, ImmutableList<N>>builderWithExpectedSize(map.size());
072        map.forEach((node, children) -> builder.put(node, children.build()));
073        return new ImmutableListTree<>(root, builder.build());
074    }
075
076    @Override
077    @SuppressWarnings("UnstableApiUsage")
078    public Iterable<? extends N> successors(N node) {
079        var successorsBuilder = map.get(node);
080        if (successorsBuilder != null) return successorsBuilder.build();
081        else return ImmutableList.of();
082    }
083
084    private record ImmutableListTree<@ImmutableTypeParameter N>(
085            N root, ImmutableMap<@NonNull N, ImmutableList<N>> map) implements ImmutableTree<N> {
086
087        @Override
088        @SuppressWarnings("UnstableApiUsage")
089        public Iterable<? extends N> successors(N node) {
090            return requireNonNull(map.getOrDefault(node, ImmutableList.of()), "TODO Annotate?");
091        }
092    }
093}