path: root/morfologik-fsa/src/main/java/morfologik/fsa/FSA5Serializer.java

package morfologik.fsa;

import java.io.IOException;
import java.io.OutputStream;
import java.nio.ByteBuffer;
import java.util.*;

import com.carrotsearch.hppc.*;
import com.carrotsearch.hppc.BitSet;

import static morfologik.fsa.FSAFlags.*;

/**
 * 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;
    }

    /**
     * {@inheritDoc}
     */
    @Override
    public FSASerializer withLogger(IMessageLogger logger) {
        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.
         */
        os.write(new byte[] { '\\', 'f', 's', 'a' });
        os.write(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;
    }
}