code: package morfologik.fsa.builders;
  
  import static morfologik.fsa.FSAFlags.*;
  
  import java.io.IOException;
  import java.io.OutputStream;
  import java.nio.ByteBuffer;
  import java.util.Arrays;
  import java.util.BitSet;
  import java.util.EnumSet;
  import java.util.Set;
  
  import morfologik.fsa.FSA;
  import morfologik.fsa.FSA5;
  import morfologik.fsa.FSAFlags;
  import morfologik.fsa.FSAHeader;
  
  import com.carrotsearch.hppc.IntIntHashMap;
  import com.carrotsearch.hppc.IntStack;
  
  /**
   * Serializes in-memory {@link FSA} graphs to a binary format compatible with
   * Jan Daciuk's <code>fsa</code>'s package <code>FSA5</code> format.
   * 
   * <p>
   * It is possible to serialize the automaton with numbers required for perfect
   * hashing. See {@link #withNumbers()} method.
   * </p>
   * 
   * @see FSA5
   * @see FSA#read(java.io.InputStream)
   */
  public final class FSA5Serializer implements FSASerializer {
    /**
     * Maximum number of bytes for a serialized arc.
     */
    private final static int MAX_ARC_SIZE = 1 + 5;
  
    /**
     * Maximum number of bytes for per-node data.
     */
    private final static int MAX_NODE_DATA_SIZE = 16;
  
    /**
     * Number of bytes for the arc's flags header (arc representation without the
     * goto address).
     */
    private final static int SIZEOF_FLAGS = 1;
  
    /**
     * Supported flags.
     */
    private final static EnumSet<FSAFlags> flags = EnumSet.of(NUMBERS, SEPARATORS, FLEXIBLE, STOPBIT, NEXTBIT);
  
    /**
     * @see FSA5#filler
     */
    public byte fillerByte = FSA5.DEFAULT_FILLER;
  
    /**
     * @see FSA5#annotation
     */
    public byte annotationByte = FSA5.DEFAULT_ANNOTATION;
  
    /**
     * <code>true</code> if we should serialize with numbers.
     * 
     * @see #withNumbers()
     */
    private boolean withNumbers;
  
    /**
     * A hash map of [state, offset] pairs.
     */
    private IntIntHashMap offsets = new IntIntHashMap();
  
    /**
     * A hash map of [state, right-language-count] pairs.
     */
    private IntIntHashMap numbers = new IntIntHashMap();
  
    /**
     * Serialize the automaton with the number of right-language sequences in each
     * node. This is required to implement perfect hashing. The numbering also
     * preserves the order of input sequences.
     * 
     * @return Returns the same object for easier call chaining.
     */
    public FSA5Serializer withNumbers() {
      withNumbers = true;
      return this;
    }
  
    /**
     * {@inheritDoc}
     */
    @Override
    public FSA5Serializer withFiller(byte filler) {
      this.fillerByte = filler;
      return this;
    }
  
    /**
     * {@inheritDoc}
     */
    @Override
    public FSA5Serializer withAnnotationSeparator(byte annotationSeparator) {
      this.annotationByte = annotationSeparator;
      return this;
    }
  
    /**
     * Serialize root state <code>s</code> to an output stream in
     * <code>FSA5</code> format.
     * 
     * @see #withNumbers()
     * @return Returns <code>os</code> for chaining.
     */
    @Override
    public <T extends OutputStream> T serialize(final FSA fsa, T os) throws IOException {
  
      // Prepare space for arc offsets and linearize all the states.
      int[] linearized = linearize(fsa);
  
      /*
       * Calculate the number of bytes required for the node data, if
       * serializing with numbers.
       */
      int nodeDataLength = 0;
      if (withNumbers) {
        this.numbers = FSAUtils.rightLanguageForAllStates(fsa);
        int maxNumber = numbers.get(fsa.getRootNode());
        while (maxNumber > 0) {
          nodeDataLength++;
          maxNumber >>>= 8;
        }
      }
  
      // Calculate minimal goto length.
      int gtl = 1;
      while (true) {
        // First pass: calculate offsets of states.
        if (!emitArcs(fsa, null, linearized, gtl, nodeDataLength)) {
          gtl++;
          continue;
        }
  
        // Second pass: check if goto overflows anywhere.
        if (emitArcs(fsa, null, linearized, gtl, nodeDataLength))
          break;
  
        gtl++;
      }
  
      /*
       * Emit the header.
       */
      FSAHeader.write(os, FSA5.VERSION);
      os.write(fillerByte);
      os.write(annotationByte);
      os.write((nodeDataLength << 4) | gtl);
  
      /*
       * Emit the automaton.
       */
      boolean gtlUnchanged = emitArcs(fsa, os, linearized, gtl, nodeDataLength);
      assert gtlUnchanged : "gtl changed in the final pass.";
  
      return os;
    }
  
    /**
     * Return supported flags.
     */
    @Override
    public Set<FSAFlags> getFlags() {
      return flags;
    }
  
    /**
     * Linearization of states.
     */
    private int[] linearize(final FSA fsa) {
      int[] linearized = new int[0];
      int last = 0;
  
      BitSet visited = new BitSet();
      IntStack nodes = new IntStack();
      nodes.push(fsa.getRootNode());
  
      while (!nodes.isEmpty()) {
        final int node = nodes.pop();
        if (visited.get(node)) {
          continue;
        }
  
        if (last >= linearized.length) {
          linearized = Arrays.copyOf(linearized, linearized.length + 100000);
        }
  
        visited.set(node);
        linearized[last++] = node;
  
        for (int arc = fsa.getFirstArc(node); arc != 0; arc = fsa.getNextArc(arc)) {
          if (!fsa.isArcTerminal(arc)) {
            int target = fsa.getEndNode(arc);
            if (!visited.get(target))
              nodes.push(target);
          }
        }
      }
  
      return Arrays.copyOf(linearized, last);
    }
  
    /**
     * Update arc offsets assuming the given goto length.
     */
    private boolean emitArcs(FSA fsa, OutputStream os, int[] linearized, int gtl, int nodeDataLength) throws IOException {
      final ByteBuffer bb = ByteBuffer.allocate(Math.max(MAX_NODE_DATA_SIZE, MAX_ARC_SIZE));
  
      int offset = 0;
  
      // Add dummy terminal state.
      offset += emitNodeData(bb, os, nodeDataLength, 0);
      offset += emitArc(bb, os, gtl, 0, (byte) 0, 0);
  
      // Add epsilon state.
      offset += emitNodeData(bb, os, nodeDataLength, 0);
      if (fsa.getRootNode() != 0)
        offset += emitArc(bb, os, gtl, FSA5.BIT_LAST_ARC | FSA5.BIT_TARGET_NEXT, (byte) '^', 0);
      else
        offset += emitArc(bb, os, gtl, FSA5.BIT_LAST_ARC, (byte) '^', 0);
  
      int maxStates = linearized.length;
      for (int j = 0; j < maxStates; j++) {
        final int s = linearized[j];
  
        if (os == null) {
          offsets.put(s, offset);
        } else {
          assert offsets.get(s) == offset : s + " " + offsets.get(s) + " " + offset;
        }
  
        offset += emitNodeData(bb, os, nodeDataLength, withNumbers ? numbers.get(s) : 0);
  
        for (int arc = fsa.getFirstArc(s); arc != 0; arc = fsa.getNextArc(arc)) {
          int targetOffset;
          final int target;
          if (fsa.isArcTerminal(arc)) {
            targetOffset = 0;
            target = 0;
          } else {
            target = fsa.getEndNode(arc);
            targetOffset = offsets.get(target);
          }
  
          int flags = 0;
          if (fsa.isArcFinal(arc)) {
            flags |= FSA5.BIT_FINAL_ARC;
          }
  
          if (fsa.getNextArc(arc) == 0) {
            flags |= FSA5.BIT_LAST_ARC;
  
            if (j + 1 < maxStates && target == linearized[j + 1] && targetOffset != 0) {
              flags |= FSA5.BIT_TARGET_NEXT;
              targetOffset = 0;
            }
          }
  
          int bytes = emitArc(bb, os, gtl, flags, fsa.getArcLabel(arc), targetOffset);
          if (bytes < 0)
            // gtl too small. interrupt eagerly.
            return false;
  
          offset += bytes;
        }
      }
  
      return true;
    }
  
    /** */
    private int emitArc(ByteBuffer bb, OutputStream os, int gtl, int flags, byte label, int targetOffset)
        throws IOException {
      int arcBytes = (flags & FSA5.BIT_TARGET_NEXT) != 0 ? SIZEOF_FLAGS : gtl;
  
      flags |= (targetOffset << 3);
      bb.put(label);
      for (int b = 0; b < arcBytes; b++) {
        bb.put((byte) flags);
        flags >>>= 8;
      }
  
      if (flags != 0) {
        // gtl too small. interrupt eagerly.
        return -1;
      }
  
      bb.flip();
      int bytes = bb.remaining();
      if (os != null) {
        os.write(bb.array(), bb.position(), bb.remaining());
      }
      bb.clear();
  
      return bytes;
    }
  
    /** */
    private int emitNodeData(ByteBuffer bb, OutputStream os, int nodeDataLength, int number) throws IOException {
      if (nodeDataLength > 0 && os != null) {
        for (int i = 0; i < nodeDataLength; i++) {
          bb.put((byte) number);
          number >>>= 8;
        }
  
        bb.flip();
        os.write(bb.array(), bb.position(), bb.remaining());
        bb.clear();
      }
  
      return nodeDataLength;
    }
  }
  