code: package morfologik.fsa.builders;
  
  import static morfologik.fsa.CFSA2.*;
  import static morfologik.fsa.FSAFlags.*;
  
  import java.io.IOException;
  import java.io.OutputStream;
  import java.util.BitSet;
  import java.util.Comparator;
  import java.util.EnumSet;
  import java.util.Locale;
  import java.util.PriorityQueue;
  import java.util.Set;
  import java.util.TreeSet;
  import java.util.logging.Level;
  import java.util.logging.Logger;
  
  import com.carrotsearch.hppc.BoundedProportionalArraySizingStrategy;
  import com.carrotsearch.hppc.IntArrayList;
  import com.carrotsearch.hppc.IntIntHashMap;
  import com.carrotsearch.hppc.IntStack;
  import com.carrotsearch.hppc.cursors.IntCursor;
  import com.carrotsearch.hppc.cursors.IntIntCursor;
  
  import morfologik.fsa.CFSA2;
  import morfologik.fsa.FSA;
  import morfologik.fsa.FSAFlags;
  import morfologik.fsa.FSAHeader;
  import morfologik.fsa.StateVisitor;
  import morfologik.fsa.builders.FSAUtils.IntIntHolder;
  
  /**
   * Serializes in-memory {@link FSA} graphs to {@link CFSA2}.
   * 
   * <p>
   * It is possible to serialize the automaton with numbers required for perfect
   * hashing. See {@link #withNumbers()} method.
   * </p>
   * 
   * @see CFSA2
   */
  public final class CFSA2Serializer implements FSASerializer {
    private final Logger logger = Logger.getLogger(getClass().getName());
  
    /**
     * Supported flags.
     */
    private final static EnumSet<FSAFlags> flags = EnumSet.of(NUMBERS, FLEXIBLE, STOPBIT, NEXTBIT);
  
    /**
     * No-state id.
     */
    private final static int NO_STATE = -1;
  
    /**
     * <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();
  
    /**
     * Scratch array for serializing vints.
     */
    private final byte[] scratch = new byte[5];
  
    /**
     * The most frequent labels for integrating with the flags field.
     */
    private byte[] labelsIndex;
  
    /**
     * Inverted index of labels to be integrated with flags field. A label at
     * index <code>i<code> has the index or zero (no integration).
     */
    private int[] labelsInvIndex;
  
    /**
     * 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 CFSA2Serializer withNumbers() {
      withNumbers = true;
      return this;
    }
  
    /**
     * Serializes any {@link FSA} to {@link CFSA2} stream.
     * 
     * @see #withNumbers()
     * @return Returns <code>os</code> for chaining.
     */
    @Override
    public <T extends OutputStream> T serialize(final FSA fsa, T os) throws IOException {
      /*
       * Calculate the most frequent labels and build indexed labels dictionary.
       */
      computeLabelsIndex(fsa);
  
      /*
       * Calculate the number of bytes required for the node data, if
       * serializing with numbers.
       */
      if (withNumbers) {
        this.numbers = FSAUtils.rightLanguageForAllStates(fsa);
      }
  
      /*
       * Linearize all the states, optimizing their layout.
       */
      IntArrayList linearized = linearize(fsa);
  
      /*
       * Emit the header.
       */
      FSAHeader.write(os, CFSA2.VERSION);
  
      EnumSet<FSAFlags> fsaFlags = EnumSet.of(FLEXIBLE, STOPBIT, NEXTBIT);
      if (withNumbers) {
        fsaFlags.add(NUMBERS);
      }
  
      final short sflags = FSAFlags.asShort(fsaFlags);
      os.write((sflags >> 8) & 0xFF);
      os.write((sflags) & 0xFF);
  
      /*
       * Emit labels index.
       */
      os.write(labelsIndex.length);
      os.write(labelsIndex);
  
      /*
       * Emit the automaton.
       */
      int size = emitNodes(fsa, os, linearized);
      assert size == 0 : "Size changed in the final pass?";
  
      return os;
    }
  
    /**
     * Compute a set of labels to be integrated with the flags field.
     */
    private void computeLabelsIndex(final FSA fsa) {
      // Compute labels count.
      final int[] countByValue = new int[256];
  
      fsa.visitAllStates(new StateVisitor() {
        public boolean accept(int state) {
          for (int arc = fsa.getFirstArc(state); arc != 0; arc = fsa.getNextArc(arc))
            countByValue[fsa.getArcLabel(arc) & 0xff]++;
          return true;
        }
      });
  
      // Order by descending frequency of counts and increasing label value.
      Comparator<IntIntHolder> comparator = new Comparator<IntIntHolder>() {
        public int compare(IntIntHolder o1, IntIntHolder o2) {
          int countDiff = o2.b - o1.b;
          if (countDiff == 0) {
            countDiff = o1.a - o2.a;
          }
          return countDiff;
        }
      };
  
      TreeSet<IntIntHolder> labelAndCount = new TreeSet<IntIntHolder>(comparator);
      for (int label = 0; label < countByValue.length; label++) {
        if (countByValue[label] > 0) {
          labelAndCount.add(new IntIntHolder(label, countByValue[label]));
        }
      }
  
      labelsIndex = new byte[1 + Math.min(labelAndCount.size(), CFSA2.LABEL_INDEX_SIZE)];
      labelsInvIndex = new int[256];
      for (int i = labelsIndex.length - 1; i > 0 && !labelAndCount.isEmpty(); i--) {
        IntIntHolder p = labelAndCount.first();
        labelAndCount.remove(p);
        labelsInvIndex[p.a] = i;
        labelsIndex[i] = (byte) p.a;
      }
    }
  
    /**
     * Return supported flags.
     */
    @Override
    public Set<FSAFlags> getFlags() {
      return flags;
    }
  
    /**
     * Linearization of states.
     */
    private IntArrayList linearize(final FSA fsa) throws IOException {
      /*
       * Compute the states with most inlinks. These should be placed as close to the 
       * start of the automaton, as possible so that v-coded addresses are tiny.  
       */
      final IntIntHashMap inlinkCount = computeInlinkCount(fsa);
  
      /*
       * An array of ordered states for serialization.
       */
      final IntArrayList linearized = new IntArrayList(0, new BoundedProportionalArraySizingStrategy(1000, 10000, 1.5f));
  
      /*
       * Determine which states should be linearized first (at fixed positions) so as to
       * minimize the place occupied by goto fields.
       */
      int maxStates = Integer.MAX_VALUE;
      int minInlinkCount = 2;
      int [] states = computeFirstStates(inlinkCount, maxStates, minInlinkCount);
  
      /*
       * Compute initial addresses, without node rearrangements.
       */
      int serializedSize = linearizeAndCalculateOffsets(fsa, new IntArrayList(), linearized, offsets);
  
      /*
       * Probe for better node arrangements by selecting between [lower, upper]
       * nodes from the potential candidate nodes list. 
       */
      IntArrayList sublist = new IntArrayList();
      sublist.buffer = states;
      sublist.elementsCount = states.length;
  
      /*
       * Probe the initial region a little bit, looking for optimal cut. It can't be binary search
       * because the result isn't monotonic.
       */
      log(Level.FINE, "Compacting, initial output size: %,d", serializedSize);
      int cutAt = 0;
      for (int cut = Math.min(25, states.length); cut <= Math.min(150, states.length); cut += 25) {
        sublist.elementsCount = cut;
        int newSize = linearizeAndCalculateOffsets(fsa, sublist, linearized, offsets);
        log(Level.FINE, "Moved %,d states, output size: %,d", sublist.size(), newSize);
        if (newSize >= serializedSize) {
          break;
        }
        cutAt = cut;
      }
  
      /*
       * Cut at the calculated point and repeat linearization.
       */
      sublist.elementsCount = cutAt;
      int size = linearizeAndCalculateOffsets(fsa, sublist, linearized, offsets);
      log(Level.FINE, "%,d states moved, final size: %,d", sublist.size(), size);
      return linearized;
    }
  
    private void log(Level level, String msg, Object... args) {
      logger.log(level, String.format(Locale.ROOT, msg, args));
    }
  
    /**
     * Linearize all states, putting <code>states</code> in front of the automaton
     * and calculating stable state offsets.
     */
    private int linearizeAndCalculateOffsets(FSA fsa, IntArrayList states, IntArrayList linearized,
        IntIntHashMap offsets) throws IOException {
      final BitSet visited = new BitSet();
      final IntStack nodes = new IntStack();
      linearized.clear();
  
      /*
       * Linearize states with most inlinks first.
       */
      for (int i = 0; i < states.size(); i++) {
        linearizeState(fsa, nodes, linearized, visited, states.get(i));
      }
  
      /*
       * Linearize the remaining states by chaining them one after another, in depth-order.
       */
      nodes.push(fsa.getRootNode());
      while (!nodes.isEmpty()) {
        final int node = nodes.pop();
        if (visited.get(node))
          continue;
  
        linearizeState(fsa, nodes, linearized, visited, node);
      }
  
      /*
       * Calculate new state offsets. This is iterative. We start with 
       * maximum potential offsets and recalculate until converged.
       */
      int MAX_OFFSET = Integer.MAX_VALUE;
      for (IntCursor c : linearized) {
        offsets.put(c.value, MAX_OFFSET);
      }
  
      int i, j = 0;
      while ((i = emitNodes(fsa, null, linearized)) > 0) {
        j = i;
      }
      return j;
    }
  
    /**
     * Add a state to linearized list.
     */
    private void linearizeState(final FSA fsa, IntStack nodes, IntArrayList linearized, BitSet visited, int node) {
      linearized.add(node);
      visited.set(node);
      for (int arc = fsa.getFirstArc(node); arc != 0; arc = fsa.getNextArc(arc)) {
        if (!fsa.isArcTerminal(arc)) {
          final int target = fsa.getEndNode(arc);
          if (!visited.get(target))
            nodes.push(target);
        }
      }
    }
  
    /**
     * Compute the set of states that should be linearized first to minimize other
     * states goto length.
     */
    private int[] computeFirstStates(IntIntHashMap inlinkCount, int maxStates, int minInlinkCount) {
      Comparator<IntIntHolder> comparator = new Comparator<FSAUtils.IntIntHolder>() {
        public int compare(IntIntHolder o1, IntIntHolder o2) {
          int v = o1.a - o2.a;
          return v == 0 ? (o1.b - o2.b) : v;
        } 
      };
  
      PriorityQueue<IntIntHolder> stateInlink = new PriorityQueue<IntIntHolder>(1, comparator);
      IntIntHolder scratch = new IntIntHolder();
      for (IntIntCursor c : inlinkCount) {
        if (c.value > minInlinkCount) {
          scratch.a = c.value;
          scratch.b = c.key;
  
          if (stateInlink.size() < maxStates || comparator.compare(scratch, stateInlink.peek()) > 0) {
            stateInlink.add(new IntIntHolder(c.value, c.key));
            if (stateInlink.size() > maxStates) {
              stateInlink.remove();
            }
          }
        }
      }
  
      int [] states = new int [stateInlink.size()];
      for (int position = states.length; !stateInlink.isEmpty();) {
        IntIntHolder i = stateInlink.remove();
        states[--position] = i.b;
      }
  
      return states;
    }
  
    /**
     * Compute in-link count for each state.
     */
    private IntIntHashMap computeInlinkCount(final FSA fsa) {
      IntIntHashMap inlinkCount = new IntIntHashMap();
      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;
  
        visited.set(node);
  
        for (int arc = fsa.getFirstArc(node); arc != 0; arc = fsa.getNextArc(arc)) {
          if (!fsa.isArcTerminal(arc)) {
            final int target = fsa.getEndNode(arc);
            inlinkCount.putOrAdd(target, 1, 1);
            if (!visited.get(target))
              nodes.push(target);
          }
        }
      }
  
      return inlinkCount;
    }
  
    /**
     * Update arc offsets assuming the given goto length.
     */
    private int emitNodes(FSA fsa, OutputStream os, IntArrayList linearized) throws IOException {
      int offset = 0;
  
      // Add epsilon state.
      offset += emitNodeData(os, 0);
      if (fsa.getRootNode() != 0)
        offset += emitArc(os, BIT_LAST_ARC, (byte) '^', offsets.get(fsa.getRootNode()));
      else
        offset += emitArc(os, BIT_LAST_ARC, (byte) '^', 0);
  
      boolean offsetsChanged = false;
      final int max = linearized.size();
      for (IntCursor c : linearized) {
        final int state = c.value;
        final int nextState = c.index + 1 < max ? linearized.get(c.index + 1) : NO_STATE;
  
        if (os == null) {
          offsetsChanged |= (offsets.get(state) != offset);
          offsets.put(state, offset);
        } else {
          assert offsets.get(state) == offset : state + " " + offsets.get(state) + " " + offset;
        }
  
        offset += emitNodeData(os, withNumbers ? numbers.get(state) : 0);
        offset += emitNodeArcs(fsa, os, state, nextState);
      }
  
      return offsetsChanged ? offset : 0;
    }
  
    /**
     * Emit all arcs of a single node.
     */
    private int emitNodeArcs(FSA fsa, OutputStream os, final int state, final int nextState) throws IOException {
      int offset = 0;
      for (int arc = fsa.getFirstArc(state); arc != 0; arc = fsa.getNextArc(arc)) {
        int targetOffset;
        final int target;
  
        if (fsa.isArcTerminal(arc)) {
          target = 0;
          targetOffset = 0;
        } else {
          target = fsa.getEndNode(arc);
          targetOffset = offsets.get(target);
        }
  
        int flags = 0;
  
        if (fsa.isArcFinal(arc)) {
          flags |= BIT_FINAL_ARC;
        }
  
        if (fsa.getNextArc(arc) == 0) {
          flags |= BIT_LAST_ARC;
        }
  
        if (targetOffset != 0 && target == nextState) {
          flags |= BIT_TARGET_NEXT;
          targetOffset = 0;
        }
  
        offset += emitArc(os, flags, fsa.getArcLabel(arc), targetOffset);
      }
  
      return offset;
    }
  
    /** */
    private int emitArc(OutputStream os, int flags, byte label, int targetOffset) throws IOException {
      int length = 0;
  
      int labelIndex = labelsInvIndex[label & 0xff];
      if (labelIndex > 0) {
        if (os != null)
          os.write(flags | labelIndex);
        length++;
      } else {
        if (os != null) {
          os.write(flags);
          os.write(label);
        }
        length += 2;
      }
  
      if ((flags & BIT_TARGET_NEXT) == 0) {
        int len = writeVInt(scratch, 0, targetOffset);
        if (os != null) {
          os.write(scratch, 0, len);
        }
        length += len;
      }
  
      return length;
    }
  
    /** */
    private int emitNodeData(OutputStream os, int number) throws IOException {
      int size = 0;
  
      if (withNumbers) {
        size = writeVInt(scratch, 0, number);
        if (os != null) {
          os.write(scratch, 0, size);
        }
      }
  
      return size;
    }
  
    /** */
    @Override
    public CFSA2Serializer withFiller(byte filler) {
      throw new UnsupportedOperationException("CFSA2 does not support filler. Use .info file.");
    }
  
    /** */
    @Override
    public CFSA2Serializer withAnnotationSeparator(byte annotationSeparator) {
      throw new UnsupportedOperationException("CFSA2 does not support separator. Use .info file.");
    }
  
    /**
     * Write a v-int to a byte array.
     */
    static int writeVInt(byte[] array, int offset, int value) {
      assert value >= 0 : "Can't v-code negative ints.";
  
      while (value > 0x7F) {
        array[offset++] = (byte) (0x80 | (value & 0x7F));
        value >>= 7;
      }
      array[offset++] = (byte) value;
  
      return offset;
    }
  }
  