HuffmanCompressionTable.java

/*
 * This file is part of Waarp Project (named also Waarp or GG).
 *
 *  Copyright (c) 2019, Waarp SAS, and individual contributors by the @author
 *  tags. See the COPYRIGHT.txt in the distribution for a full listing of
 * individual contributors.
 *
 *  All Waarp Project is free software: you can redistribute it and/or
 * modify it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or (at your
 * option) any later version.
 *
 * Waarp is distributed in the hope that it will be useful, but WITHOUT ANY
 * WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR
 * A PARTICULAR PURPOSE. See the GNU General Public License for more details.
 *
 *  You should have received a copy of the GNU General Public License along with
 * Waarp . If not, see <http://www.gnu.org/licenses/>.
 */

/*
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.waarp.compress.zstdunsafe;

import java.util.Arrays;

final class HuffmanCompressionTable {
  private final short[] values;
  private final byte[] numberOfBits;

  private int maxSymbol;
  private int maxNumberOfBits;

  public HuffmanCompressionTable(final int capacity) {
    this.values = new short[capacity];
    this.numberOfBits = new byte[capacity];
  }

  public static int optimalNumberOfBits(final int maxNumberOfBits,
                                        final int inputSize,
                                        final int maxSymbol) {
    if (inputSize <= 1) {
      throw new IllegalArgumentException(); // not supported. Use RLE instead
    }

    int result = maxNumberOfBits;

    result = Math.min(result, Util.highestBit((inputSize - 1)) -
                              1); // we may be able to reduce accuracy if input is small

    // Need a minimum to safely represent all symbol values
    result = Math.max(result, Util.minTableLog(inputSize, maxSymbol));

    result =
        Math.max(result, Huffman.MIN_TABLE_LOG); // absolute minimum for Huffman
    result =
        Math.min(result, Huffman.MAX_TABLE_LOG); // absolute maximum for Huffman

    return result;
  }

  public void initialize(final int[] counts, final int maxSymbol,
                         int maxNumberOfBits,
                         final HuffmanCompressionTableWorkspace workspace) {
    Util.checkArgument(maxSymbol <= Huffman.MAX_SYMBOL,
                       "Max symbol value too large");

    workspace.reset();

    final NodeTable nodeTable = workspace.nodeTable;
    nodeTable.reset();

    final int lastNonZero = buildTree(counts, maxSymbol, nodeTable);

    // enforce max table log
    maxNumberOfBits =
        setMaxHeight(nodeTable, lastNonZero, maxNumberOfBits, workspace);
    Util.checkArgument(maxNumberOfBits <= Huffman.MAX_TABLE_LOG,
                       "Max number of bits larger than max table size");

    // populate table
    final int symbolCount = maxSymbol + 1;
    for (int node = 0; node < symbolCount; node++) {
      final int symbol = nodeTable.symbols[node];
      numberOfBits[symbol] = nodeTable.numberOfBits[node];
    }

    final short[] entriesPerRank = workspace.entriesPerRank;
    final short[] valuesPerRank = workspace.valuesPerRank;

    for (int n = 0; n <= lastNonZero; n++) {
      entriesPerRank[nodeTable.numberOfBits[n]]++;
    }

    // determine starting value per rank
    short startingValue = 0;
    for (int rank = maxNumberOfBits; rank > 0; rank--) {
      valuesPerRank[rank] =
          startingValue; // get starting value within each rank
      startingValue += entriesPerRank[rank];
      startingValue >>>= 1;
    }

    for (int n = 0; n <= maxSymbol; n++) {
      values[n] =
          valuesPerRank[numberOfBits[n]]++; // assign value within rank, symbol order
    }

    this.maxSymbol = maxSymbol;
    this.maxNumberOfBits = maxNumberOfBits;
  }

  private int buildTree(final int[] counts, final int maxSymbol,
                        final NodeTable nodeTable) {
    // populate the leaves of the node table from the histogram of counts
    // in descending order by count, ascending by symbol value.
    short current = 0;

    for (int symbol = 0; symbol <= maxSymbol; symbol++) {
      final int count = counts[symbol];

      // simple insertion sort
      int position = current;
      while (position > 1 && count > nodeTable.count[position - 1]) {
        nodeTable.copyNode(position - 1, position);
        position--;
      }

      nodeTable.count[position] = count;
      nodeTable.symbols[position] = symbol;

      current++;
    }

    int lastNonZero = maxSymbol;
    while (nodeTable.count[lastNonZero] == 0) {
      lastNonZero--;
    }

    // populate the non-leaf nodes
    final short nonLeafStart = Huffman.MAX_SYMBOL_COUNT;
    current = nonLeafStart;

    int currentLeaf = lastNonZero;

    // combine the two smallest leaves to create the first intermediate node
    int currentNonLeaf = current;
    nodeTable.count[current] =
        nodeTable.count[currentLeaf] + nodeTable.count[currentLeaf - 1];
    nodeTable.parents[currentLeaf] = current;
    nodeTable.parents[currentLeaf - 1] = current;
    current++;
    currentLeaf -= 2;

    final int root = Huffman.MAX_SYMBOL_COUNT + lastNonZero - 1;

    // fill in sentinels
    for (int n = current; n <= root; n++) {
      nodeTable.count[n] = 1 << 30;
    }

    // create parents
    while (current <= root) {
      final int child1;
      if (currentLeaf >= 0 &&
          nodeTable.count[currentLeaf] < nodeTable.count[currentNonLeaf]) {
        child1 = currentLeaf--;
      } else {
        child1 = currentNonLeaf++;
      }

      final int child2;
      if (currentLeaf >= 0 &&
          nodeTable.count[currentLeaf] < nodeTable.count[currentNonLeaf]) {
        child2 = currentLeaf--;
      } else {
        child2 = currentNonLeaf++;
      }

      nodeTable.count[current] =
          nodeTable.count[child1] + nodeTable.count[child2];
      nodeTable.parents[child1] = current;
      nodeTable.parents[child2] = current;
      current++;
    }

    // distribute weights
    nodeTable.numberOfBits[root] = 0;
    for (int n = root - 1; n >= nonLeafStart; n--) {
      final short parent = nodeTable.parents[n];
      nodeTable.numberOfBits[n] = (byte) (nodeTable.numberOfBits[parent] + 1);
    }

    for (int n = 0; n <= lastNonZero; n++) {
      final short parent = nodeTable.parents[n];
      nodeTable.numberOfBits[n] = (byte) (nodeTable.numberOfBits[parent] + 1);
    }

    return lastNonZero;
  }

  // TODO: consider encoding 2 symbols at a time
  //   - need a table with 256x256 entries with
  //      - the concatenated bits for the corresponding pair of symbols
  //      - the sum of bits for the corresponding pair of symbols
  //   - read 2 symbols at a time from the input
  public void encodeSymbol(final BitOutputStream output, final int symbol) {
    output.addBitsFast(values[symbol], numberOfBits[symbol]);
  }

  public int write(final Object outputBase, final long outputAddress,
                   final int outputSize,
                   final HuffmanTableWriterWorkspace workspace) {
    final byte[] weights = workspace.weights;

    long output = outputAddress;

    final int numberOfBits1 = this.maxNumberOfBits;
    final int maxSymbol1 = this.maxSymbol;

    // convert to weights per RFC 8478 section 4.2.1
    for (int symbol = 0; symbol < maxSymbol1; symbol++) {
      final int bits = numberOfBits[symbol];

      if (bits == 0) {
        weights[symbol] = 0;
      } else {
        weights[symbol] = (byte) (numberOfBits1 + 1 - bits);
      }
    }

    // attempt weights compression by FSE
    int size = compressWeights(outputBase, output + 1, outputSize - 1, weights,
                               maxSymbol1, workspace);

    if (maxSymbol1 > 127 && size > 127) {
      // This should never happen. Since weights are in the range [0, 12], they can be compressed optimally to ~3.7 bits per symbol for a uniform distribution.
      // Since maxSymbol has to be <= MAX_SYMBOL (255), this is 119 bytes + FSE headers.
      throw new AssertionError();
    }

    if (size != 0 && size != 1 && size < maxSymbol1 / 2) {
      // Go with FSE only if:
      //   - the weights are compressible
      //   - the compressed size is better than what we'd get with the raw encoding below
      //   - the compressed size is <= 127 bytes, which is the most that the encoding can hold for FSE-compressed weights (see RFC 8478 section 4.2.1.1). This is implied
      //     by the maxSymbol / 2 check, since maxSymbol must be <= 255
      UnsafeUtil.UNSAFE.putByte(outputBase, output, (byte) size);
      return size + 1; // header + size
    } else {
      // Use raw encoding (4 bits per entry)

      // #entries = #symbols - 1 since last symbol is implicit. Thus, #entries = (maxSymbol + 1) - 1 = maxSymbol

      size = (maxSymbol1 + 1) / 2;  // ceil(#entries / 2)
      Util.checkArgument(size + 1 /* header */ <= outputSize,
                         "Output size too small"); // 2 entries per byte

      // encode number of symbols
      // header = #entries + 127 per RFC
      UnsafeUtil.UNSAFE.putByte(outputBase, output, (byte) (127 + maxSymbol1));
      output++;

      weights[maxSymbol1] =
          0; // last weight is implicit, so set to 0 so that it doesn't get encoded below
      for (int i = 0; i < maxSymbol1; i += 2) {
        UnsafeUtil.UNSAFE.putByte(outputBase, output,
                                  (byte) ((weights[i] << 4) +
                                          (weights[i + 1] & 0xFF)));
        output++;
      }

      return (int) (output - outputAddress);
    }
  }

  /**
   * Can this table encode all symbols with non-zero count?
   */
  public boolean isValid(final int[] counts, final int maxSymbol) {
    if (maxSymbol > this.maxSymbol) {
      // some non-zero count symbols cannot be encoded by the current table
      return false;
    }

    for (int symbol = 0; symbol <= maxSymbol; ++symbol) {
      if (counts[symbol] != 0 && numberOfBits[symbol] == 0) {
        return false;
      }
    }
    return true;
  }

  public int estimateCompressedSize(final int[] counts, final int maxSymbol) {
    int numberOfBits1 = 0;
    for (int symbol = 0; symbol <= Math.min(maxSymbol, this.maxSymbol);
         symbol++) {
      numberOfBits1 += this.numberOfBits[symbol] * counts[symbol];
    }

    return numberOfBits1 >>> 3; // convert to bytes
  }

  // http://fastcompression.blogspot.com/2015/07/huffman-revisited-part-3-depth-limited.html
  private static int setMaxHeight(final NodeTable nodeTable,
                                  final int lastNonZero,
                                  final int maxNumberOfBits,
                                  final HuffmanCompressionTableWorkspace workspace) {
    final int largestBits = nodeTable.numberOfBits[lastNonZero];

    if (largestBits <= maxNumberOfBits) {
      return largestBits;   // early exit: no elements > maxNumberOfBits
    }

    // there are several too large elements (at least >= 2)
    int totalCost = 0;
    final int baseCost = 1 << (largestBits - maxNumberOfBits);
    int n = lastNonZero;

    while (nodeTable.numberOfBits[n] > maxNumberOfBits) {
      totalCost += baseCost - (1 << (largestBits - nodeTable.numberOfBits[n]));
      nodeTable.numberOfBits[n] = (byte) maxNumberOfBits;
      n--;
    }  // n stops at nodeTable.numberOfBits[n + offset] <= maxNumberOfBits

    while (nodeTable.numberOfBits[n] == maxNumberOfBits) {
      n--;   // n ends at index of smallest symbol using < maxNumberOfBits
    }

    // renormalize totalCost
    totalCost >>>= (largestBits -
                    maxNumberOfBits);  // note: totalCost is necessarily a multiple of baseCost

    // repay normalized cost
    final int noSymbol = 0xF0F0F0F0;
    final int[] rankLast = workspace.rankLast;
    Arrays.fill(rankLast, noSymbol);

    // Get pos of last (smallest) symbol per rank
    int currentNbBits = maxNumberOfBits;
    for (int pos = n; pos >= 0; pos--) {
      if (nodeTable.numberOfBits[pos] >= currentNbBits) {
        continue;
      }
      currentNbBits = nodeTable.numberOfBits[pos];   // < maxNumberOfBits
      rankLast[maxNumberOfBits - currentNbBits] = pos;
    }

    while (totalCost > 0) {
      int numberOfBitsToDecrease = Util.highestBit(totalCost) + 1;
      for (; numberOfBitsToDecrease > 1; numberOfBitsToDecrease--) {
        final int highPosition = rankLast[numberOfBitsToDecrease];
        final int lowPosition = rankLast[numberOfBitsToDecrease - 1];
        if (highPosition == noSymbol) {
          continue;
        }
        if (lowPosition == noSymbol) {
          break;
        }
        final int highTotal = nodeTable.count[highPosition];
        final int lowTotal = 2 * nodeTable.count[lowPosition];
        if (highTotal <= lowTotal) {
          break;
        }
      }

      // only triggered when no more rank 1 symbol left => find closest one (note : there is necessarily at least one !)
      // HUF_MAX_TABLELOG test just to please gcc 5+; but it should not be necessary
      while ((numberOfBitsToDecrease <= Huffman.MAX_TABLE_LOG) &&
             (rankLast[numberOfBitsToDecrease] == noSymbol)) {
        numberOfBitsToDecrease++;
      }
      totalCost -= 1 << (numberOfBitsToDecrease - 1);
      if (rankLast[numberOfBitsToDecrease - 1] == noSymbol) {
        rankLast[numberOfBitsToDecrease - 1] =
            rankLast[numberOfBitsToDecrease];   // this rank is no longer empty
      }
      nodeTable.numberOfBits[rankLast[numberOfBitsToDecrease]]++;
      if (rankLast[numberOfBitsToDecrease] ==
          0) {   /* special case, reached largest symbol */
        rankLast[numberOfBitsToDecrease] = noSymbol;
      } else {
        rankLast[numberOfBitsToDecrease]--;
        if (nodeTable.numberOfBits[rankLast[numberOfBitsToDecrease]] !=
            maxNumberOfBits - numberOfBitsToDecrease) {
          rankLast[numberOfBitsToDecrease] =
              noSymbol;   // this rank is now empty
        }
      }
    }

    while (totalCost < 0) {  // Sometimes, cost correction overshoot
      if (rankLast[1] ==
          noSymbol) {  /* special case : no rank 1 symbol (using maxNumberOfBits-1); let's create one from largest rank 0 (using maxNumberOfBits) */
        while (nodeTable.numberOfBits[n] == maxNumberOfBits) {
          n--;
        }
        nodeTable.numberOfBits[n + 1]--;
        rankLast[1] = n + 1;
        totalCost++;
        continue;
      }
      nodeTable.numberOfBits[rankLast[1] + 1]--;
      rankLast[1]++;
      totalCost++;
    }

    return maxNumberOfBits;
  }

  /**
   * All elements within weightTable must be <= Huffman.MAX_TABLE_LOG
   */
  private static int compressWeights(final Object outputBase,
                                     final long outputAddress,
                                     final int outputSize, final byte[] weights,
                                     final int weightsLength,
                                     final HuffmanTableWriterWorkspace workspace) {
    if (weightsLength <= 1) {
      return 0; // Not compressible
    }

    // Scan input and build symbol stats
    final int[] counts = workspace.counts;
    Histogram.count(weights, weightsLength, counts);
    final int maxSymbol =
        Histogram.findMaxSymbol(counts, Huffman.MAX_TABLE_LOG);
    final int maxCount = Histogram.findLargestCount(counts, maxSymbol);

    if (maxCount == weightsLength) {
      return 1; // only a single symbol in source
    }
    if (maxCount == 1) {
      return 0; // each symbol present maximum once => not compressible
    }

    final short[] normalizedCounts = workspace.normalizedCounts;

    final int tableLog =
        FiniteStateEntropy.optimalTableLog(Huffman.MAX_FSE_TABLE_LOG,
                                           weightsLength, maxSymbol);
    FiniteStateEntropy.normalizeCounts(normalizedCounts, tableLog, counts,
                                       weightsLength, maxSymbol);

    long output = outputAddress;
    final long outputLimit = outputAddress + outputSize;

    // Write table description header
    final int headerSize =
        FiniteStateEntropy.writeNormalizedCounts(outputBase, output, outputSize,
                                                 normalizedCounts, maxSymbol,
                                                 tableLog);
    output += headerSize;

    // Compress
    final FseCompressionTable compressionTable = workspace.fseTable;
    compressionTable.initialize(normalizedCounts, maxSymbol, tableLog);
    final int compressedSize = FiniteStateEntropy.compress(outputBase, output,
                                                           (int) (outputLimit -
                                                                  output),
                                                           weights,
                                                           weightsLength,
                                                           compressionTable);
    if (compressedSize == 0) {
      return 0;
    }
    output += compressedSize;

    return (int) (output - outputAddress);
  }
}