Extend prefix code alphabet size to 68.

This commit also adds the tool that was used to extend the prefix
codes to the new alphabet size.

encoder/CMakeLists.txt

@@ -31,3 +31,6 @@ target_link_libraries(jxl_tiny hwy pthread)
 
 add_executable(cjxl_tiny cjxl_main.cc)
 target_link_libraries(cjxl_tiny jxl_tiny)
+
+add_executable(update_static_entropy_codes update_static_entropy_codes_main.cc)
+target_link_libraries(update_static_entropy_codes jxl_tiny)

encoder/config.h

@@ -7,7 +7,7 @@
 #ifndef ENCODER_CONFIG_H_
 #define ENCODER_CONFIG_H_
 
-#define OPTIMIZE_CODE 1
+#define OPTIMIZE_CODE 0
 #define OPTIMIZE_CHROMA_FROM_LUMA 1
 #define OPTIMIZE_BLOCK_SIZES 1
 

encoder/enc_entropy_code.cc

@@ -274,8 +274,6 @@ void WriteHuffmanTree(const uint8_t* depth, size_t length, size_t* tree_size,
   }
 }
 
-namespace {
-
 uint16_t ReverseBits(int num_bits, uint16_t bits) {
   static const size_t kLut[16] = {// Pre-reversed 4-bit values.
                                   0x0, 0x8, 0x4, 0xc, 0x2, 0xa, 0x6, 0xe,
@@ -321,6 +319,8 @@ void ConvertBitDepthsToSymbols(const uint8_t* depth, size_t len,
   }
 }
 
+namespace {
+
 // num = alphabet size
 // depths = symbol depths
 void StoreHuffmanTree(const uint8_t* depths, size_t num, BitWriter* writer) {

encoder/enc_entropy_code.h

@@ -17,6 +17,9 @@
 
 namespace jxl {
 
+void ConvertBitDepthsToSymbols(const uint8_t* depth, size_t len,
+                               uint16_t* bits);
+
 void OptimizePrefixCodes(const std::vector<Token>& tokens, EntropyCode* code);
 
 void OptimizeEntropyCode(const std::vector<Token>& tokens, EntropyCode* code);

encoder/entropy_code.h

@@ -13,7 +13,7 @@
 
 namespace jxl {
 
-static constexpr size_t kAlphabetSize = 64;
+static constexpr size_t kAlphabetSize = 68;
 static constexpr uint8_t kMaxContexts = 128;
 
 struct PrefixCode {

encoder/read_pfm.cc

@@ -202,7 +202,7 @@ bool ReadPFM(const char* fn, Image3F* image) {
     for (size_t c = 0; c < 3; ++c) {
       float* row_out = img.PlaneRow(c, y);
       for (size_t x = 0; x < xsize; ++x) {
-        row_out[x] = row_in[x * 3 + c];
+        memcpy(&row_out[x], &row_in[x * 3 + c], sizeof(*row_out));
         if (big_endian) row_out[x] = BSwapFloat(row_out[x]);
       }
     }

encoder/update_static_entropy_codes_main.cc

@@ -0,0 +1,162 @@
+// Copyright (c) the JPEG XL Project Authors.
+//
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file or at
+// https://developers.google.com/open-source/licenses/bsd
+
+#include <stdio.h>
+#include <stdlib.h>
+#include <string.h>
+
+#include <string>
+#include <vector>
+
+#include "encoder/base/printf_macros.h"
+#include "encoder/enc_entropy_code.h"
+#include "encoder/enc_huffman_tree.h"
+#include "encoder/entropy_code.h"
+#include "encoder/static_entropy_codes.h"
+
+namespace jxl {
+
+struct DynamicPrefixCode {
+  std::vector<uint8_t> depths;
+  std::vector<uint16_t> bits;
+};
+
+void OutputCodes(const char* type,
+                 const std::vector<DynamicPrefixCode>& prefix_codes) {
+  printf("static constexpr size_t kNum%sPrefixCodes = %" PRIuS ";\n", type,
+         prefix_codes.size());
+  printf("static constexpr PrefixCode k%sPrefixCodes[kNum%sPrefixCodes] = {\n",
+         type, type);
+  for (const auto& prefix_code : prefix_codes) {
+    size_t alphabet_size = prefix_code.depths.size();
+    printf("    {{\n");
+    for (size_t j = 0; j < alphabet_size; ++j) {
+      printf("%s%2d,%s", j % 16 == 0 ? "         " : " ", prefix_code.depths[j],
+             (j % 16 == 15 || j + 1 == alphabet_size) ? "\n" : "");
+    }
+    printf("     },\n");
+    printf("     {\n");
+    for (size_t j = 0; j < alphabet_size; ++j) {
+      printf("%s0x%04x,%s", j % 8 == 0 ? "         " : " ", prefix_code.bits[j],
+             (j % 8 == 7 || j + 1 == alphabet_size) ? "\n" : "");
+    }
+    printf("     }},\n");
+  }
+  printf("};\n");
+}
+
+bool ExtendPrefixCode(DynamicPrefixCode& prefix_code,
+                      size_t new_alphabet_size) {
+  static const int kTreeLimit = 15;
+
+  const size_t alphabet_size = prefix_code.depths.size();
+  if (prefix_code.bits.size() != alphabet_size) {
+    return false;
+  }
+
+  // Step 1. Convert bit depths to population counts.
+  std::vector<uint32_t> counts(new_alphabet_size);
+  uint32_t total_count = 0;
+  for (size_t i = 0; i < alphabet_size; ++i) {
+    counts[i] = 1 << (kTreeLimit - prefix_code.depths[i]);
+    total_count += counts[i];
+  }
+  if (total_count != 1 << kTreeLimit) {
+    return false;
+  }
+
+  // Step 2. Extend population counts by 1s.
+  for (size_t i = alphabet_size; i < new_alphabet_size; ++i) {
+    counts[i] = 1;
+  }
+
+  // Step 3. Regenerate depths and bits from new population counts.
+  prefix_code.depths.resize(new_alphabet_size);
+  prefix_code.bits.resize(new_alphabet_size);
+  CreateHuffmanTree(&counts[0], new_alphabet_size, kTreeLimit,
+                    &prefix_code.depths[0]);
+  ConvertBitDepthsToSymbols(&prefix_code.depths[0], new_alphabet_size,
+                            &prefix_code.bits[0]);
+  return true;
+}
+
+enum PrefixCodeType {
+  DC = 0,
+  AC = 1,
+};
+
+std::vector<DynamicPrefixCode> ConvertToDynamicCodes(
+    const PrefixCode* prefix_codes, size_t num_codes) {
+  std::vector<DynamicPrefixCode> out;
+  for (size_t i = 0; i < num_codes; ++i) {
+    const uint8_t* depths = prefix_codes[i].depths;
+    const uint16_t* bits = prefix_codes[i].bits;
+    DynamicPrefixCode code = {
+        {depths, depths + kAlphabetSize},
+        {bits, bits + kAlphabetSize},
+    };
+    out.emplace_back(std::move(code));
+  }
+  return out;
+}
+
+bool GenerateNewPrefixCodes(PrefixCodeType type, size_t new_alphabet_size) {
+  const char* type_name;
+  std::vector<DynamicPrefixCode> prefix_codes;
+  if (type == PrefixCodeType::DC) {
+    type_name = "DC";
+    prefix_codes = ConvertToDynamicCodes(kDCPrefixCodes, kNumDCPrefixCodes);
+  } else if (type == PrefixCodeType::AC) {
+    type_name = "AC";
+    prefix_codes = ConvertToDynamicCodes(kACPrefixCodes, kNumACPrefixCodes);
+  }
+  for (auto& prefix_code : prefix_codes) {
+    if (!ExtendPrefixCode(prefix_code, new_alphabet_size)) {
+      return false;
+    }
+  }
+  OutputCodes(type_name, prefix_codes);
+  return true;
+}
+
+}  // namespace jxl
+
+void PrintHelp(char* arg0) {
+  fprintf(stderr,
+          "Usage: %s <type> <new alphabet size>\n\n"
+          "Prints the updated entropy codes of the given type to stdout.\n"
+          "  <type> can be either 'DC' or 'AC'\n",
+          arg0);
+};
+
+int main(int argc, char** argv) {
+  if (argc != 3) {
+    PrintHelp(argv[0]);
+    return EXIT_FAILURE;
+  }
+  jxl::PrefixCodeType type = jxl::PrefixCodeType::DC;
+  if (strcmp(argv[1], "DC") == 0) {
+    type = jxl::PrefixCodeType::DC;
+  } else if (strcmp(argv[1], "AC") == 0) {
+    type = jxl::PrefixCodeType::AC;
+  } else {
+    PrintHelp(argv[0]);
+    return EXIT_FAILURE;
+  }
+  size_t new_alphabet_size = std::stoi(argv[2]);
+  if (new_alphabet_size <= jxl::kAlphabetSize) {
+    fprintf(stderr,
+            "New alphabet size must be greater than current alphabet size, "
+            "which is %" PRIuS ".\n",
+            jxl::kAlphabetSize);
+    return EXIT_FAILURE;
+  }
+  if (!jxl::GenerateNewPrefixCodes(type, new_alphabet_size)) {
+    fprintf(stderr, "Failed to extend prefix codes (internal error)\n");
+    return EXIT_FAILURE;
+  }
+  return EXIT_SUCCESS;
+}