Auxiliary Java class to organize FSM

AbstractFsm.java

package com.igormaznitsa.misc;

import java.util.Arrays;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CopyOnWriteArrayList;
import java.util.concurrent.atomic.AtomicReference;
import java.util.function.BiConsumer;
import java.util.function.BiPredicate;
import java.util.function.Consumer;

/**
 * Small auxiliary class to organize FSM with notifications about transitions.
 * The Code snippet is published under
 * <a href="https://www.apache.org/licenses/LICENSE-2.0">Apache License 2.0</a>
 *
 * @author Igor Maznitsa (http://www.igormaznitsa.com)
 *
 * @param <T> enum data type to be used as state class
 */
public abstract class AbstractFsm<T extends Enum<?>> {

  private final Map<T, Set<T>> nodes = new HashMap<>();
  private final Map<T, Consumer<T>> hooks = new ConcurrentHashMap<>();
  private final AtomicReference<T> state = new AtomicReference<>();

  private final Map<T, Map<T, BiPredicate<T, T>>> transitionPredicates = new ConcurrentHashMap<>();
  private final List<BiConsumer<T, T>> transitionListeners = new CopyOnWriteArrayList<>();
  private final List<BiConsumer<T, T>> illegalTransitionListeners = new CopyOnWriteArrayList<>();
  private final List<Consumer<T>> terminalListeners = new CopyOnWriteArrayList<>();

  public static final class Builder<T> {

    private final Map<T, Set<T>> nodes = new HashMap<>();

    private T curfrom;

    public Builder<T> from(final T state) {
      this.curfrom = notNull(state);
      this.nodes.put(this.curfrom, new HashSet<>());
      return this;
    }

    public Builder<T> to(final T... states) {
      this.nodes.get(notNull(this.curfrom)).addAll(Arrays.asList(states));
      return this;
    }

    public Builder<T> loop() {
      this.nodes.get(notNull(this.curfrom)).add(this.curfrom);
      return this;
    }

    private Builder<T> validate() {
      if (this.nodes.isEmpty()) {
        throw new IllegalStateException("Empty state map not allowed");
      }
      this.nodes.entrySet().forEach((e) -> {
        e.getValue().stream().filter((s) -> (!this.nodes.containsKey(s))).forEachOrdered((s) -> {
          throw new IllegalStateException(String.format("Detected wrong edge %s -> %s", e.getKey(), s));
        });
      });
      return this;
    }
  }

  public AbstractFsm() {
    this.nodes.putAll(make(new Builder<>()).validate().nodes);
    this.state.set(notNull(initial()));
  }

  /**
   * Template method to fill FSM by states and transitions.
   *
   * @param builder use methods of the object to add states
   */
  protected abstract Builder<T> make(Builder<T> builder);

  /**
   * Template method returns initial state for FSM, nust not be null.
   *
   * @return initial state, must not be null
   */
  protected abstract T initial();

  public AbstractFsm<T> hook(final T state, final Consumer<T> hook) {
    if (this.hooks.containsKey(state)) {
      throw new IllegalArgumentException(String.format("State %s already hooked", state));
    } else {
      this.hooks.put(state, hook);
    }
    return this;
  }

  public AbstractFsm<T> removeHook(final T state) {
    this.hooks.remove(notNull(state));
    return this;
  }

  public AbstractFsm<T> resetHooks() {
    this.hooks.clear();
    return this;
  }

  public AbstractFsm<T> reset(){
    this.hooks.clear();
    this.transitionPredicates.clear();
    this.illegalTransitionListeners.clear();
    this.terminalListeners.clear();
    this.transitionListeners.clear();
    force(initial());
    return this;
  }
  
  public AbstractFsm<T> setPredicate(T from, T to, BiPredicate<T, T> predicate) {
    Map<T, BiPredicate<T, T>> map = this.transitionPredicates.get(from);
    if (map == null) {
      map = new ConcurrentHashMap<>();
      this.transitionPredicates.put(from, map);
    }
    map.put(to, predicate);
    return this;
  }

  public AbstractFsm<T> resetPredicate(T from, T to) {
    Map<T, BiPredicate<T, T>> map = this.transitionPredicates.get(from);
    if (map != null) {
      map.remove(to);
    }
    return this;
  }

  public void addTerminalListener(final Consumer<T> listener) {
    this.terminalListeners.add(notNull(listener));
  }

  public void removeTerminalListener(final Consumer<T> listener) {
    this.terminalListeners.remove(notNull(listener));
  }

  public void addTransitionListener(final BiConsumer<T, T> listener) {
    this.transitionListeners.add(notNull(listener));
  }

  public void removeTransitionListener(final BiConsumer<T, T> listener) {
    this.transitionListeners.remove(notNull(listener));
  }

  public void addIllegalTransitionListener(final BiConsumer<T, T> listener) {
    this.illegalTransitionListeners.add(notNull(listener));
  }

  public void removeIllegalTransitionListener(final BiConsumer<T, T> listener) {
    this.illegalTransitionListeners.remove(notNull(listener));
  }

  private boolean changeInternalState(final T state, final boolean forced) {
    final T prevState = this.state.get();
    final Consumer<T> consumer = this.hooks.get(state);
    boolean result = false;

    if (!forced) {
      final Map<T, BiPredicate<T, T>> map = this.transitionPredicates.get(prevState);
      if (map != null) {
        final BiPredicate<T, T> predicate = map.get(state);
        if (predicate != null && !predicate.test(prevState, state)) {
          this.illegalTransitionListeners.forEach((l) -> {
            l.accept(prevState, state);
          });
          return false;
        }
      }
    }

    if (this.state.compareAndSet(prevState, state)) {
      if (!forced) {
        if (consumer != null) {
          consumer.accept(state);
        }
        this.transitionListeners.forEach((c) -> {
          c.accept(prevState, state);
        });

        if (this.nodes.get(state).isEmpty()) {
          this.terminalListeners.forEach((c) -> {
            c.accept(state);
          });
        }
      }
      result = true;
    }
    return result;
  }

  /**
   * Check that FSM in terminal state, state which doesn't have allowed next
   * states
   *
   * @return true if FSM in terminal state, false otherwise
   */
  public boolean isTerminated() {
    final T current = this.state.get();
    return current == null ? false : this.nodes.get(current).isEmpty();
  }

  /**
   * Get current state
   *
   * @return current state, must not be null
   */
  public T get() {
    return this.state.get();
  }

  /**
   * Try move FSM into new state.
   *
   * @param newState new state, must not be null
   * @return true if FSM is moved into new state, false otherwise
   */
  public boolean tryNext(final T newState) {
    final T current = this.state.get();
    final Set<T> possibleTransitions = current == null ? Collections.<T>emptySet() : this.nodes.get(current);
    if (possibleTransitions.contains(notNull(newState))) {
      return changeInternalState(newState, false);
    } else {
      this.illegalTransitionListeners.forEach((l) -> {
        l.accept(current, newState);
      });
    }
    return false;
  }

  /**
   * Set state for the FSM without notifications and checks.
   *
   * @param state state to be set, must not be null
   */
  public void force(final T state) {
    if (this.nodes.containsKey(notNull(state))) {
      changeInternalState(state, true);
    }
  }

  protected static <X> X notNull(final X value) {
    if (value == null) {
      throw new NullPointerException();
    }
    return value;
  }
}