text: // Copyright (c) the JPEG XL Project Authors. All rights reserved.
  //
  // Use of this source code is governed by a BSD-style
  // license that can be found in the LICENSE file.
  
  #include "fast_lossless.h"
  
  #include <assert.h>
  #include <stdint.h>
  #include <stdio.h>
  #include <string.h>
  
  #include <algorithm>
  #include <array>
  #include <memory>
  #include <queue>
  #include <vector>
  
  #if (!defined(__BYTE_ORDER__) || (__BYTE_ORDER__ != __ORDER_LITTLE_ENDIAN__))
  #error "system not known to be little endian"
  #endif
  
  struct BitWriter {
    void Allocate(size_t maximum_bit_size) {
      assert(data == nullptr);
      // Leave some padding.
      data.reset((uint8_t*)malloc(maximum_bit_size / 8 + 32));
    }
  
    void Write(uint32_t count, uint64_t bits) {
      buffer |= bits << bits_in_buffer;
      bits_in_buffer += count;
      memcpy(data.get() + bytes_written, &buffer, 8);
      size_t bytes_in_buffer = bits_in_buffer / 8;
      bits_in_buffer -= bytes_in_buffer * 8;
      buffer >>= bytes_in_buffer * 8;
      bytes_written += bytes_in_buffer;
    }
  
    void ZeroPadToByte() {
      if (bits_in_buffer != 0) {
        Write(8 - bits_in_buffer, 0);
      }
    }
  
    std::unique_ptr<uint8_t[], void (*)(void*)> data = {nullptr, free};
    size_t bytes_written = 0;
    size_t bits_in_buffer = 0;
    uint64_t buffer = 0;
  };
  
  constexpr size_t kLZ77Offset = 224;
  constexpr size_t kLZ77MinLength = 16;
  
  struct PrefixCode {
    static constexpr size_t kNumLZ77 = 17;
    static constexpr size_t kNumRaw = 15;
  
    alignas(32) uint8_t raw_nbits[16] = {};
    alignas(32) uint8_t raw_bits[16] = {};
    uint8_t lz77_nbits[kNumLZ77] = {};
  
    uint16_t lz77_bits[kNumLZ77] = {};
  
    static uint16_t BitReverse(size_t nbits, uint16_t bits) {
      constexpr uint16_t kNibbleLookup[16] = {
          0b0000, 0b1000, 0b0100, 0b1100, 0b0010, 0b1010, 0b0110, 0b1110,
          0b0001, 0b1001, 0b0101, 0b1101, 0b0011, 0b1011, 0b0111, 0b1111,
      };
      uint16_t rev16 = (kNibbleLookup[bits & 0xF] << 12) |
                       (kNibbleLookup[(bits >> 4) & 0xF] << 8) |
                       (kNibbleLookup[(bits >> 8) & 0xF] << 4) |
                       (kNibbleLookup[bits >> 12]);
      return rev16 >> (16 - nbits);
    }
  
    // Create the prefix codes given the code lengths.
    // Supports the code lengths being split into two halves.
    static void ComputeCanonicalCode(const uint8_t* first_chunk_nbits,
                                     uint8_t* first_chunk_bits,
                                     size_t first_chunk_size,
                                     const uint8_t* second_chunk_nbits,
                                     uint16_t* second_chunk_bits,
                                     size_t second_chunk_size) {
      uint8_t code_length_counts[16] = {};
      for (size_t i = 0; i < first_chunk_size; i++) {
        code_length_counts[first_chunk_nbits[i]]++;
        assert(first_chunk_nbits[i] <= 7);
        assert(first_chunk_nbits[i] > 0);
      }
      for (size_t i = 0; i < second_chunk_size; i++) {
        code_length_counts[second_chunk_nbits[i]]++;
      }
  
      uint16_t next_code[16] = {};
  
      uint16_t code = 0;
      for (size_t i = 1; i < 16; i++) {
        code = (code + code_length_counts[i - 1]) << 1;
        next_code[i] = code;
      }
  
      for (size_t i = 0; i < first_chunk_size; i++) {
        first_chunk_bits[i] =
            BitReverse(first_chunk_nbits[i], next_code[first_chunk_nbits[i]]++);
      }
      for (size_t i = 0; i < second_chunk_size; i++) {
        second_chunk_bits[i] =
            BitReverse(second_chunk_nbits[i], next_code[second_chunk_nbits[i]]++);
      }
    }
  
    PrefixCode(uint64_t* raw_counts, uint64_t* lz77_counts) {
      // "merge" together all the lz77 counts in a single symbol for the level 1
      // table (containing just the raw symbols, up to length 7).
      uint64_t level1_counts[kNumRaw + 1];
      memcpy(level1_counts, raw_counts, kNumRaw * sizeof(uint64_t));
      size_t numraw = kNumRaw;
      while (numraw > 0 && level1_counts[numraw - 1] == 0) numraw--;
  
      level1_counts[numraw] = 0;
      for (size_t i = 0; i < kNumLZ77; i++) {
        level1_counts[numraw] += lz77_counts[i];
      }
      uint8_t level1_nbits[kNumRaw + 1] = {};
      ComputeCodeLengths(level1_counts, numraw + 1, 7, level1_nbits);
  
      uint8_t level2_nbits[kNumLZ77] = {};
      ComputeCodeLengths(lz77_counts, kNumLZ77, 15 - level1_nbits[numraw],
                         level2_nbits);
      for (size_t i = 0; i < numraw; i++) {
        raw_nbits[i] = level1_nbits[i];
      }
      for (size_t i = 0; i < kNumLZ77; i++) {
        lz77_nbits[i] =
            level2_nbits[i] ? level1_nbits[numraw] + level2_nbits[i] : 0;
      }
  
      ComputeCanonicalCode(raw_nbits, raw_bits, numraw, lz77_nbits, lz77_bits,
                           kNumLZ77);
    }
  
    static void ComputeCodeLengths(uint64_t* freqs, size_t n, size_t limit,
                                   uint8_t* nbits) {
      if (n <= 1) return;
      assert(n <= (1 << limit));
      assert(n <= 32);
      int parent[64] = {};
      int height[64] = {};
      using QElem = std::pair<uint64_t, size_t>;
      std::priority_queue<QElem, std::vector<QElem>, std::greater<QElem>> q;
      // Standard Huffman code construction. On failure (i.e. if going beyond the
      // length limit), try again with halved frequencies.
      while (true) {
        size_t num_nodes = 0;
        for (size_t i = 0; i < n; i++) {
          if (freqs[i] == 0) continue;
          q.emplace(freqs[i], num_nodes++);
        }
        if (num_nodes <= 1) return;
        while (q.size() > 1) {
          QElem n1 = q.top();
          q.pop();
          QElem n2 = q.top();
          q.pop();
          size_t next = num_nodes++;
          parent[n1.second] = next;
          parent[n2.second] = next;
          q.emplace(n1.first + n2.first, next);
        }
        assert(q.size() == 1);
        q.pop();
        bool is_ok = true;
        for (size_t i = num_nodes - 1; i-- > 0;) {
          height[i] = height[parent[i]] + 1;
          is_ok &= height[i] <= limit;
        }
        if (is_ok) {
          num_nodes = 0;
          for (size_t i = 0; i < n; i++) {
            if (freqs[i] == 0) continue;
            nbits[i] = height[num_nodes++];
          }
          break;
        } else {
          for (size_t i = 0; i < n; i++) {
            freqs[i] = (freqs[i] + 1) >> 1;
          }
        }
      }
    }
  
    void WriteTo(BitWriter* writer) const {
      uint64_t code_length_counts[18] = {};
      code_length_counts[17] = 3 + 2 * (kNumLZ77 - 1);
      for (size_t i = 0; i < kNumRaw; i++) {
        code_length_counts[raw_nbits[i]]++;
      }
      for (size_t i = 0; i < kNumLZ77; i++) {
        code_length_counts[lz77_nbits[i]]++;
      }
      uint8_t code_length_nbits[18] = {};
      ComputeCodeLengths(code_length_counts, 18, 5, code_length_nbits);
      writer->Write(2, 0b00);  // HSKIP = 0, i.e. don't skip code lengths.
  
      // As per Brotli RFC.
      uint8_t code_length_order[18] = {1, 2, 3, 4,  0,  5,  17, 6,  16,
                                       7, 8, 9, 10, 11, 12, 13, 14, 15};
      uint8_t code_length_length_nbits[] = {2, 4, 3, 2, 2, 4};
      uint8_t code_length_length_bits[] = {0, 7, 3, 2, 1, 15};
  
      // Encode lengths of code lengths.
      size_t num_code_lengths = 18;
      while (code_length_nbits[code_length_order[num_code_lengths - 1]] == 0) {
        num_code_lengths--;
      }
      for (size_t i = 0; i < num_code_lengths; i++) {
        int symbol = code_length_nbits[code_length_order[i]];
        writer->Write(code_length_length_nbits[symbol],
                      code_length_length_bits[symbol]);
      }
  
      // Compute the canonical codes for the codes that represent the lengths of
      // the actual codes for data.
      uint16_t code_length_bits[18] = {};
      ComputeCanonicalCode(nullptr, nullptr, 0, code_length_nbits,
                           code_length_bits, 18);
      // Encode raw bit code lengths.
      for (size_t i = 0; i < kNumRaw; i++) {
        writer->Write(code_length_nbits[raw_nbits[i]],
                      code_length_bits[raw_nbits[i]]);
      }
      size_t num_lz77 = kNumLZ77;
      while (lz77_nbits[num_lz77 - 1] == 0) {
        num_lz77--;
      }
      // Encode 0s until 224 (start of LZ77 symbols). This is in total 224-15 =
      // 209.
      static_assert(kLZ77Offset == 224, "");
      static_assert(kNumRaw == 15, "");
      writer->Write(code_length_nbits[17], code_length_bits[17]);
      writer->Write(3, 0b010);  // 5
      writer->Write(code_length_nbits[17], code_length_bits[17]);
      writer->Write(3, 0b000);  // (5-2)*8 + 3 = 27
      writer->Write(code_length_nbits[17], code_length_bits[17]);
      writer->Write(3, 0b110);  // (27-2)*8 + 9 = 209
      // Encode LZ77 symbols, with values 224+i*16.
      for (size_t i = 0; i < num_lz77; i++) {
        writer->Write(code_length_nbits[lz77_nbits[i]],
                      code_length_bits[lz77_nbits[i]]);
        if (i != num_lz77 - 1) {
          // Encode gap between LZ77 symbols: 15 zeros.
          writer->Write(code_length_nbits[17], code_length_bits[17]);
          writer->Write(3, 0b000);  // 3
          writer->Write(code_length_nbits[17], code_length_bits[17]);
          writer->Write(3, 0b100);  // (3-2)*8+7 = 15
        }
      }
    }
  };
  
  constexpr size_t kChunkSize = 16;
  
  void EncodeHybridUint000(uint32_t value, uint32_t* token, uint32_t* nbits,
                           uint32_t* bits) {
    uint32_t n = 31 - __builtin_clz(value);
    *token = value ? n + 1 : 0;
    *nbits = value ? n : 0;
    *bits = value ? value - (1 << n) : 0;
  }
  
  void AppendWriter(BitWriter* dest, const BitWriter* src) {
    if (dest->bits_in_buffer == 0) {
      memcpy(dest->data.get() + dest->bytes_written, src->data.get(),
             src->bytes_written);
      dest->bytes_written += src->bytes_written;
    } else {
      size_t i = 0;
      uint64_t buf = dest->buffer;
      uint64_t bits_in_buffer = dest->bits_in_buffer;
      uint8_t* dest_buf = dest->data.get() + dest->bytes_written;
      // Copy 8 bytes at a time until we reach the border.
      for (; i + 8 < src->bytes_written; i += 8) {
        uint64_t chunk;
        memcpy(&chunk, src->data.get() + i, 8);
        uint64_t out = buf | (chunk << bits_in_buffer);
        memcpy(dest_buf + i, &out, 8);
        buf = chunk >> (64 - bits_in_buffer);
      }
      dest->buffer = buf;
      dest->bytes_written += i;
      for (; i < src->bytes_written; i++) {
        dest->Write(8, src->data[i]);
      }
    }
    dest->Write(src->bits_in_buffer, src->buffer);
  }
  
  void AssembleFrame(size_t width, size_t height, size_t nb_chans,
                     size_t bitdepth,
                     const std::vector<std::array<BitWriter, 4>>& group_data,
                     BitWriter* output) {
    size_t total_size_groups = 0;
    std::vector<size_t> group_sizes(group_data.size());
    for (size_t i = 0; i < group_data.size(); i++) {
      size_t sz = 0;
      for (size_t j = 0; j < nb_chans; j++) {
        const auto& writer = group_data[i][j];
        sz += writer.bytes_written * 8 + writer.bits_in_buffer;
      }
      sz = (sz + 7) / 8;
      group_sizes[i] = sz;
      total_size_groups += sz * 8;
    }
    output->Allocate(1000 + group_data.size() * 32 + total_size_groups);
  
    // Signature
    output->Write(16, 0x0AFF);
  
    // Size header, hand-crafted.
    // Not small
    output->Write(1, 0);
  
    auto wsz = [output](size_t size) {
      if (size - 1 < (1 << 9)) {
        output->Write(2, 0b00);
        output->Write(9, size - 1);
      } else if (size - 1 < (1 << 13)) {
        output->Write(2, 0b01);
        output->Write(13, size - 1);
      } else if (size - 1 < (1 << 18)) {
        output->Write(2, 0b10);
        output->Write(18, size - 1);
      } else {
        output->Write(2, 0b11);
        output->Write(30, size - 1);
      }
    };
  
    wsz(height);
  
    // No special ratio.
    output->Write(3, 0);
  
    wsz(width);
  
    // Hand-crafted ImageMetadata.
    output->Write(1, 0);  // all_default
    output->Write(1, 0);  // extra_fields
    output->Write(1, 0);  // bit_depth.floating_point_sample
    if (bitdepth == 8) {
      output->Write(2, 0b00);  // bit_depth.bits_per_sample = 8
    } else if (bitdepth == 10) {
      output->Write(2, 0b01);  // bit_depth.bits_per_sample = 10
    } else if (bitdepth == 12) {
      output->Write(2, 0b10);  // bit_depth.bits_per_sample = 12
    } else {
      output->Write(2, 0b11);  // 1 + u(6)
      output->Write(6, bitdepth - 1);
    }
    output->Write(1, 1);  // 16-bit-buffer sufficient
    bool have_alpha = (nb_chans == 2 || nb_chans == 4);
    if (have_alpha) {
      output->Write(2, 0b01);  // One extra channel
      output->Write(1, 1);     // ... all_default (ie. 8-bit alpha)
    } else {
      output->Write(2, 0b00);  // No extra channel
    }
    output->Write(1, 0);  // Not XYB
    if (nb_chans > 1) {
      output->Write(1, 1);  // color_encoding.all_default (sRGB)
    } else {
      output->Write(1, 0);     // color_encoding.all_default false
      output->Write(1, 0);     // color_encoding.want_icc false
      output->Write(2, 1);     // grayscale
      output->Write(2, 1);     // D65
      output->Write(1, 0);     // no gamma transfer function
      output->Write(2, 0b10);  // tf: 2 + u(4)
      output->Write(4, 11);    // tf of sRGB
      output->Write(2, 1);     // relative rendering intent
    }
    output->Write(2, 0b00);  // No extensions.
  
    output->Write(1, 1);  // all_default transform data
  
    // No ICC, no preview. Frame should start at byte boundery.
    output->ZeroPadToByte();
  
    auto wsz_fh = [output](size_t size) {
      if (size < (1 << 8)) {
        output->Write(2, 0b00);
        output->Write(8, size);
      } else if (size - 256 < (1 << 11)) {
        output->Write(2, 0b01);
        output->Write(11, size - 256);
      } else if (size - 2304 < (1 << 14)) {
        output->Write(2, 0b10);
        output->Write(14, size - 2304);
      } else {
        output->Write(2, 0b11);
        output->Write(30, size - 18688);
      }
    };
  
    // Handcrafted frame header.
    output->Write(1, 0);     // all_default
    output->Write(2, 0b00);  // regular frame
    output->Write(1, 1);     // modular
    output->Write(2, 0b00);  // default flags
    output->Write(1, 0);     // not YCbCr
    output->Write(2, 0b00);  // no upsampling
    if (have_alpha) {
      output->Write(2, 0b00);  // no alpha upsampling
    }
    output->Write(2, 0b01);  // default group size
    output->Write(2, 0b00);  // exactly one pass
    if (width % kChunkSize == 0) {
      output->Write(1, 0);  // no custom size or origin
    } else {
      output->Write(1, 1);  // custom size
      wsz_fh(0);            // x0 = 0
      wsz_fh(0);            // y0 = 0
      wsz_fh((width + kChunkSize - 1) / kChunkSize *
             kChunkSize);  // xsize rounded up to chunk size
      wsz_fh(height);      // ysize same
    }
    output->Write(2, 0b00);  // kReplace blending mode
    if (have_alpha) {
      output->Write(2, 0b00);  // kReplace blending mode for alpha channel
    }
    output->Write(1, 1);     // is_last
    output->Write(2, 0b00);  // a frame has no name
    output->Write(1, 0);     // loop filter is not all_default
    output->Write(1, 0);     // no gaborish
    output->Write(2, 0);     // 0 EPF iters
    output->Write(2, 0b00);  // No LF extensions
    output->Write(2, 0b00);  // No FH extensions
  
    output->Write(1, 0);      // No TOC permutation
    output->ZeroPadToByte();  // TOC is byte-aligned.
    for (size_t i = 0; i < group_data.size(); i++) {
      size_t sz = group_sizes[i];
      if (sz < (1 << 10)) {
        output->Write(2, 0b00);
        output->Write(10, sz);
      } else if (sz - 1024 < (1 << 14)) {
        output->Write(2, 0b01);
        output->Write(14, sz - 1024);
      } else if (sz - 17408 < (1 << 22)) {
        output->Write(2, 0b10);
        output->Write(22, sz - 17408);
      } else {
        output->Write(2, 0b11);
        output->Write(30, sz - 4211712);
      }
    }
    output->ZeroPadToByte();  // Groups are byte-aligned.
  
    for (size_t i = 0; i < group_data.size(); i++) {
      for (size_t j = 0; j < nb_chans; j++) {
        AppendWriter(output, &group_data[i][j]);
      }
      output->ZeroPadToByte();
    }
  }
  
  void PrepareDCGlobalCommon(bool is_single_group, size_t width, size_t height,
                             const PrefixCode& code, BitWriter* output) {
    output->Allocate(100000 + (is_single_group ? width * height * 16 : 0));
    // No patches, spline or noise.
    output->Write(1, 1);  // default DC dequantization factors (?)
    output->Write(1, 1);  // use global tree / histograms
    output->Write(1, 0);  // no lz77 for the tree
  
    output->Write(1, 1);   // simple code for the tree's context map
    output->Write(2, 0);   // all contexts clustered together
    output->Write(1, 1);   // use prefix code for tree
    output->Write(4, 15);  // don't do hybriduint for tree - 2 symbols anyway
    output->Write(7, 0b0100101);  // Alphabet size is 6: we need 0 and 5 (var16)
    output->Write(2, 1);          // simple prefix code
    output->Write(2, 1);          // with two symbols
    output->Write(3, 0);          // 0
    output->Write(3, 5);          // 5
    output->Write(5, 0b00010);    // tree repr: predictor is 5, all else 0
  
    output->Write(1, 1);     // Enable lz77 for the main bitstream
    output->Write(2, 0b00);  // lz77 offset 224
    static_assert(kLZ77Offset == 224, "");
    output->Write(10, 0b0000011111);  // lz77 min length 16
    static_assert(kLZ77MinLength == 16, "");
    output->Write(4, 4);  // 404 hybrid uint config for lz77: 4
    output->Write(3, 0);  // 0
    output->Write(3, 4);  // 4
    output->Write(1, 1);  // simple code for the context map
    output->Write(2, 1);  // two clusters
    output->Write(1, 1);  // raw/lz77 length histogram last
    output->Write(1, 0);  // distance histogram first
    output->Write(1, 1);  // use prefix codes
    output->Write(4, 0);  // 000 hybrid uint config for distances (only need 0)
    output->Write(4, 0);  // 000 hybrid uint config for symbols (only <= 10)
    // Distance alphabet size:
    output->Write(5, 0b00001);  // 2: just need 1 for RLE (i.e. distance 1)
    // Symbol + LZ77 alphabet size:
    output->Write(1, 1);    // > 1
    output->Write(4, 8);    // <= 512
    output->Write(8, 255);  // == 512
  
    // Distance histogram:
    output->Write(2, 1);  // simple prefix code
    output->Write(2, 0);  // with one symbol
    output->Write(1, 1);  // 1
  
    // Symbol + lz77 histogram:
    code.WriteTo(output);
  
    // Group header for global modular image.
    output->Write(1, 1);  // Global tree
    output->Write(1, 1);  // All default wp
  }
  
  void PrepareDCGlobal(bool is_single_group, size_t width, size_t height,
                       size_t nb_chans, size_t bitdepth, const PrefixCode& code,
                       BitWriter* output) {
    PrepareDCGlobalCommon(is_single_group, width, height, code, output);
    if (nb_chans > 2) {
      output->Write(2, 0b01);     // 1 transform
      output->Write(2, 0b00);     // RCT
      output->Write(5, 0b00000);  // Starting from ch 0
      output->Write(2, 0b00);     // YCoCg
    } else {
      output->Write(2, 0b00);  // no transforms
    }
    if (!is_single_group) {
      output->ZeroPadToByte();
    }
  }
  
  void EncodeHybridUint404_Mul16(uint32_t value, uint32_t* token_div16,
                                 uint32_t* nbits, uint32_t* bits) {
    // NOTE: token in libjxl is actually << 4.
    uint32_t n = 31 - __builtin_clz(value);
    *token_div16 = value < 16 ? 0 : n - 3;
    *nbits = value < 16 ? 0 : n - 4;
    *bits = value < 16 ? 0 : (value >> 4) - (1 << *nbits);
  }
  
  #ifdef FASTLL_ENABLE_AVX2_INTRINSICS
  #include <immintrin.h>
  void EncodeChunk(const uint16_t* residuals, const PrefixCode& prefix_code,
                   BitWriter& output) {
    static_assert(kChunkSize == 16, "Chunk size must be 16");
    auto value = _mm256_load_si256((__m256i*)residuals);
  
    // we know that residuals[i] has at most 12 bits, so we just need 3 nibbles
    // and don't need to mask the third. However we do need to set the high
    // byte to 0xFF, which will make table lookups return 0.
    auto lo_nibble =
        _mm256_or_si256(_mm256_and_si256(value, _mm256_set1_epi16(0xF)),
                        _mm256_set1_epi16(0xFF00));
    auto mi_nibble = _mm256_or_si256(
        _mm256_and_si256(_mm256_srli_epi16(value, 4), _mm256_set1_epi16(0xF)),
        _mm256_set1_epi16(0xFF00));
    auto hi_nibble =
        _mm256_or_si256(_mm256_srli_epi16(value, 8), _mm256_set1_epi16(0xFF00));
  
    auto lo_lut = _mm256_broadcastsi128_si256(
        _mm_setr_epi8(0, 1, 2, 2, 3, 3, 3, 3, 4, 4, 4, 4, 4, 4, 4, 4));
    auto mi_lut = _mm256_broadcastsi128_si256(
        _mm_setr_epi8(0, 5, 6, 6, 7, 7, 7, 7, 8, 8, 8, 8, 8, 8, 8, 8));
    auto hi_lut = _mm256_broadcastsi128_si256(_mm_setr_epi8(
        0, 9, 10, 10, 11, 11, 11, 11, 12, 12, 12, 12, 12, 12, 12, 12));
  
    auto lo_token = _mm256_shuffle_epi8(lo_lut, lo_nibble);
    auto mi_token = _mm256_shuffle_epi8(mi_lut, mi_nibble);
    auto hi_token = _mm256_shuffle_epi8(hi_lut, hi_nibble);
  
    auto token = _mm256_max_epi16(lo_token, _mm256_max_epi16(mi_token, hi_token));
    auto nbits = _mm256_subs_epu16(token, _mm256_set1_epi16(1));
  
    // Compute 1<<nbits.
    auto pow2_lo_lut = _mm256_broadcastsi128_si256(
        _mm_setr_epi8(1 << 0, 1 << 1, 1 << 2, 1 << 3, 1 << 4, 1 << 5, 1 << 6,
                      1u << 7, 0, 0, 0, 0, 0, 0, 0, 0));
    auto pow2_hi_lut = _mm256_broadcastsi128_si256(
        _mm_setr_epi8(0, 0, 0, 0, 0, 0, 0, 0, 1 << 0, 1 << 1, 1 << 2, 1 << 3,
                      1 << 4, 1 << 5, 1 << 6, 1u << 7));
  
    auto nbits_masked = _mm256_or_si256(nbits, _mm256_set1_epi16(0xFF00));
  
    auto nbits_pow2_lo = _mm256_shuffle_epi8(pow2_lo_lut, nbits_masked);
    auto nbits_pow2_hi = _mm256_shuffle_epi8(pow2_hi_lut, nbits_masked);
  
    auto nbits_pow2 =
        _mm256_or_si256(_mm256_slli_epi16(nbits_pow2_hi, 8), nbits_pow2_lo);
  
    auto bits = _mm256_subs_epu16(value, nbits_pow2);
  
    auto token_masked = _mm256_or_si256(token, _mm256_set1_epi16(0xFF00));
  
    // huff_nbits <= 6.
    auto huff_nbits =
        _mm256_shuffle_epi8(_mm256_broadcastsi128_si256(
                                _mm_load_si128((__m128i*)prefix_code.raw_nbits)),
                            token_masked);
  
    auto huff_bits =
        _mm256_shuffle_epi8(_mm256_broadcastsi128_si256(
                                _mm_load_si128((__m128i*)prefix_code.raw_bits)),
                            token_masked);
  
    auto huff_nbits_masked =
        _mm256_or_si256(huff_nbits, _mm256_set1_epi16(0xFF00));
  
    auto bits_shifted = _mm256_mullo_epi16(
        bits, _mm256_shuffle_epi8(pow2_lo_lut, huff_nbits_masked));
  
    nbits = _mm256_add_epi16(nbits, huff_nbits);
    bits = _mm256_or_si256(bits_shifted, huff_bits);
  
    // Merge nbits and bits from 16-bit to 32-bit lanes.
    auto nbits_hi16 = _mm256_srli_epi32(nbits, 16);
    auto nbits_lo16 = _mm256_and_si256(nbits, _mm256_set1_epi32(0xFFFF));
    auto bits_hi16 = _mm256_srli_epi32(bits, 16);
    auto bits_lo16 = _mm256_and_si256(bits, _mm256_set1_epi32(0xFFFF));
  
    nbits = _mm256_add_epi32(nbits_hi16, nbits_lo16);
    bits = _mm256_or_si256(_mm256_sllv_epi32(bits_hi16, nbits_lo16), bits_lo16);
  
    // Merge 32 -> 64 bit lanes.
    auto nbits_hi32 = _mm256_srli_epi64(nbits, 32);
    auto nbits_lo32 = _mm256_and_si256(nbits, _mm256_set1_epi64x(0xFFFFFFFF));
    auto bits_hi32 = _mm256_srli_epi64(bits, 32);
    auto bits_lo32 = _mm256_and_si256(bits, _mm256_set1_epi64x(0xFFFFFFFF));
  
    nbits = _mm256_add_epi64(nbits_hi32, nbits_lo32);
    bits = _mm256_or_si256(_mm256_sllv_epi64(bits_hi32, nbits_lo32), bits_lo32);
  
    alignas(32) uint64_t nbits_simd[4] = {};
    alignas(32) uint64_t bits_simd[4] = {};
  
    _mm256_store_si256((__m256i*)nbits_simd, nbits);
    _mm256_store_si256((__m256i*)bits_simd, bits);
  
    // Manually merge the buffer bits with the SIMD bits.
    // Necessary because Write() is only guaranteed to work with <=56 bits.
    // Trying to SIMD-fy this code results in slower speed (and definitely less
    // clarity).
    {
      for (size_t i = 0; i < 4; i++) {
        output.buffer |= bits_simd[i] << output.bits_in_buffer;
        memcpy(output.data.get() + output.bytes_written, &output.buffer, 8);
        // If >> 64, next_buffer is unused.
        uint64_t next_buffer = bits_simd[i] >> (64 - output.bits_in_buffer);
        output.bits_in_buffer += nbits_simd[i];
        // This `if` seems to be faster than using ternaries.
        if (output.bits_in_buffer >= 64) {
          output.buffer = next_buffer;
          output.bits_in_buffer -= 64;
          output.bytes_written += 8;
        }
      }
      memcpy(output.data.get() + output.bytes_written, &output.buffer, 8);
      size_t bytes_in_buffer = output.bits_in_buffer / 8;
      output.bits_in_buffer -= bytes_in_buffer * 8;
      output.buffer >>= bytes_in_buffer * 8;
      output.bytes_written += bytes_in_buffer;
    }
  }
  #endif
  
  #ifdef FASTLL_ENABLE_NEON_INTRINSICS
  #include <arm_neon.h>
  
  void EncodeChunk(const uint16_t* residuals, const PrefixCode& code,
                   BitWriter& output) {
    uint16x8_t res = vld1q_u16(residuals);
    uint16x8_t token = vsubq_u16(vdupq_n_u16(16), vclzq_u16(res));
    uint16x8_t nbits = vqsubq_u16(token, vdupq_n_u16(1));
    uint16x8_t bits = vqsubq_u16(res, vshlq_s16(vdupq_n_s16(1), nbits));
    uint16x8_t huff_bits =
        vandq_u16(vdupq_n_u16(0xFF), vqtbl1q_u8(vld1q_u8(code.raw_bits), token));
    uint16x8_t huff_nbits =
        vandq_u16(vdupq_n_u16(0xFF), vqtbl1q_u8(vld1q_u8(code.raw_nbits), token));
    bits = vorrq_u16(vshlq_u16(bits, huff_nbits), huff_bits);
    nbits = vaddq_u16(nbits, huff_nbits);
  
    // Merge nbits and bits from 16-bit to 32-bit lanes.
    uint32x4_t nbits_lo16 = vandq_u32(nbits, vdupq_n_u32(0xFFFF));
    uint32x4_t bits_hi16 = vshlq_u32(vshrq_n_u32(bits, 16), nbits_lo16);
    uint32x4_t bits_lo16 = vandq_u32(bits, vdupq_n_u32(0xFFFF));
  
    uint32x4_t nbits32 = vsraq_n_u32(nbits_lo16, nbits, 16);
    uint32x4_t bits32 = vorrq_u32(bits_hi16, bits_lo16);
  
    // Merging up to 64 bits is not faster.
  
    // Manually merge the buffer bits with the SIMD bits.
    // A bit faster.
    for (size_t i = 0; i < 4; i++) {
      output.buffer |= bits32[i] << output.bits_in_buffer;
      memcpy(output.data.get() + output.bytes_written, &output.buffer, 8);
      output.bits_in_buffer += nbits32[i];
      size_t bytes_in_buffer = output.bits_in_buffer / 8;
      output.bits_in_buffer -= bytes_in_buffer * 8;
      output.buffer >>= bytes_in_buffer * 8;
      output.bytes_written += bytes_in_buffer;
    }
  }
  #endif
  
  template <size_t bytedepth>
  struct ChunkEncoder {
    static void EncodeRle(size_t count, const PrefixCode& code,
                          BitWriter& output) {
      if (count == 0) return;
      count -= kLZ77MinLength;
      unsigned token_div16, nbits, bits;
      EncodeHybridUint404_Mul16(count, &token_div16, &nbits, &bits);
      output.Write(
          code.lz77_nbits[token_div16] + nbits,
          (bits << code.lz77_nbits[token_div16]) | code.lz77_bits[token_div16]);
    }
  
    inline void Chunk(size_t run, uint16_t* residuals) {
      EncodeRle(run, *code, *output);
  #if defined(FASTLL_ENABLE_AVX2_INTRINSICS) && FASTLL_ENABLE_AVX2_INTRINSICS
      if (bytedepth == 1) {
        EncodeChunk(residuals, *code, *output);
        return;
      }
  #elif defined(FASTLL_ENABLE_NEON_INTRINSICS) && FASTLL_ENABLE_NEON_INTRINSICS
      if (bytedepth == 1) {
        EncodeChunk(residuals, *code, *output);
        if (kChunkSize > 8) {
          EncodeChunk(residuals + 8, *code, *output);
        }
        return;
      }
  #endif
      for (size_t ix = 0; ix < kChunkSize; ix++) {
        unsigned token, nbits, bits;
        EncodeHybridUint000(residuals[ix], &token, &nbits, &bits);
        output->Write(code->raw_nbits[token] + nbits,
                      code->raw_bits[token] | bits << code->raw_nbits[token]);
      }
    }
  
    inline void Finalize(size_t run) { EncodeRle(run, *code, *output); }
  
    const PrefixCode* code;
    BitWriter* output;
  };
  
  struct ChunkSampleCollector {
    void Rle(size_t count, uint64_t* lz77_counts) {
      if (count == 0) return;
      count -= kLZ77MinLength;
      unsigned token_div16, nbits, bits;
      EncodeHybridUint404_Mul16(count, &token_div16, &nbits, &bits);
      lz77_counts[token_div16]++;
    }
  
    inline void Chunk(size_t run, uint16_t* residuals) {
      // Run is broken. Encode the run and encode the individual vector.
      Rle(run, lz77_counts);
      run = 0;
      for (size_t ix = 0; ix < kChunkSize; ix++) {
        unsigned token, nbits, bits;
        EncodeHybridUint000(residuals[ix], &token, &nbits, &bits);
        raw_counts[token]++;
      }
    }
  
    // don't count final run since we don't know how long it really is
    void Finalize(size_t run) {}
  
    uint64_t* raw_counts;
    uint64_t* lz77_counts;
  };
  
  constexpr uint16_t PackSigned(int16_t value) {
    return (static_cast<uint16_t>(value) << 1) ^
           ((static_cast<uint16_t>(~value) >> 15) - 1);
  }
  
  template <typename T>
  struct ChannelRowProcessor {
    T* t;
    inline void ProcessChunk(const int16_t* row, const int16_t* row_left,
                             const int16_t* row_top, const int16_t* row_topleft) {
      bool continue_rle = true;
      alignas(32) uint16_t residuals[kChunkSize] = {};
      for (size_t ix = 0; ix < kChunkSize; ix++) {
        int16_t px = row[ix];
        int16_t left = row_left[ix];
        int16_t top = row_top[ix];
        int16_t topleft = row_topleft[ix];
        int16_t ac = left - topleft;
        int16_t ab = left - top;
        int16_t bc = top - topleft;
        int16_t grad = static_cast<int16_t>(static_cast<uint16_t>(ac) +
                                            static_cast<uint16_t>(top));
        int16_t d = ab ^ bc;
        int16_t clamp = d < 0 ? top : left;
        int16_t s = ac ^ bc;
        int16_t pred = s < 0 ? grad : clamp;
        residuals[ix] = PackSigned(px - pred);
        continue_rle &= residuals[ix] == last;
      }
      // Run continues, nothing to do.
      if (continue_rle) {
        run += kChunkSize;
      } else {
        // Run is broken. Encode the run and encode the individual vector.
        t->Chunk(run, residuals);
        run = 0;
      }
      last = residuals[kChunkSize - 1];
    }
    void ProcessRow(const int16_t* row, const int16_t* row_left,
                    const int16_t* row_top, const int16_t* row_topleft,
                    size_t xs) {
      for (size_t x = 0; x + kChunkSize <= xs; x += kChunkSize) {
        ProcessChunk(row + x, row_left + x, row_top + x, row_topleft + x);
      }
    }
  
    void Finalize() { t->Finalize(run); }
    size_t run = 0;
    uint16_t last = 0xFFFF;  // Can never appear
  };
  
  template <typename Processor, size_t nb_chans, size_t bytedepth>
  void ProcessImageArea(const unsigned char* rgba, size_t x0, size_t y0,
                        size_t oxs, size_t xs, size_t yskip, size_t ys,
                        size_t row_stride, Processor* processors) {
    constexpr size_t kPadding = 16;
  
    int16_t group_data[nb_chans][2][256 + kPadding * 2] = {};
    int16_t allzero[nb_chans] = {};
    int16_t allone[nb_chans];
    auto get_pixel = [&](size_t x, size_t y, size_t channel) {
      int16_t p = rgba[row_stride * (y0 + y) + (x0 + x) * nb_chans * bytedepth +
                       channel * bytedepth];
      if (bytedepth == 2) {
        p <<= 8;
        p |= rgba[row_stride * (y0 + y) + (x0 + x) * nb_chans * 2 + channel * 2 +
                  1];
      }
      return p;
    };
  
    for (size_t i = 0; i < nb_chans; i++) allone[i] = 0xffff;
    for (size_t y = 0; y < ys; y++) {
      // Pre-fill rows with YCoCg converted pixels.
      for (size_t x = 0; x < oxs; x++) {
        if (nb_chans < 3) {
          int16_t luma = get_pixel(x, y, 0);
          group_data[0][y & 1][x + kPadding] = luma;
          if (nb_chans == 2) {
            int16_t a = get_pixel(x, y, 1);
            group_data[1][y & 1][x + kPadding] = a;
          }
        } else {
          int16_t r = get_pixel(x, y, 0);
          int16_t g = get_pixel(x, y, 1);
          int16_t b = get_pixel(x, y, 2);
          if (nb_chans == 4) {
            int16_t a = get_pixel(x, y, 3);
            group_data[3][y & 1][x + kPadding] = a;
            group_data[1][y & 1][x + kPadding] = a ? r - b : 0;
            int16_t tmp = b + (group_data[1][y & 1][x + kPadding] >> 1);
            group_data[2][y & 1][x + kPadding] = a ? g - tmp : 0;
            group_data[0][y & 1][x + kPadding] =
                a ? tmp + (group_data[2][y & 1][x + kPadding] >> 1) : 0;
          } else {
            group_data[1][y & 1][x + kPadding] = r - b;
            int16_t tmp = b + (group_data[1][y & 1][x + kPadding] >> 1);
            group_data[2][y & 1][x + kPadding] = g - tmp;
            group_data[0][y & 1][x + kPadding] =
                tmp + (group_data[2][y & 1][x + kPadding] >> 1);
          }
        }
        for (size_t c = 0; c < nb_chans; c++) {
          allzero[c] |= group_data[c][y & 1][x + kPadding];
          allone[c] &= group_data[c][y & 1][x + kPadding];
        }
      }
      // Deal with x == 0.
      for (size_t c = 0; c < nb_chans; c++) {
        group_data[c][y & 1][kPadding - 1] =
            y > 0 ? group_data[c][(y - 1) & 1][kPadding] : 0;
        // Fix topleft.
        group_data[c][(y - 1) & 1][kPadding - 1] =
            y > 0 ? group_data[c][(y - 1) & 1][kPadding] : 0;
      }
      // Fill in padding.
      for (size_t c = 0; c < nb_chans; c++) {
        for (size_t x = oxs; x < xs; x++) {
          group_data[c][y & 1][kPadding + x] =
              group_data[c][y & 1][kPadding + oxs - 1];
        }
      }
      if (y < yskip) continue;
      for (size_t c = 0; c < nb_chans; c++) {
        if (y > 0 && (allzero[c] == 0 || (allone[c] == 0xff && bytedepth == 1))) {
          processors[c].run += xs;
          continue;
        }
  
        // Get pointers to px/left/top/topleft data to speedup loop.
        const int16_t* row = &group_data[c][y & 1][kPadding];
        const int16_t* row_left = &group_data[c][y & 1][kPadding - 1];
        const int16_t* row_top =
            y == 0 ? row_left : &group_data[c][(y - 1) & 1][kPadding];
        const int16_t* row_topleft =
            y == 0 ? row_left : &group_data[c][(y - 1) & 1][kPadding - 1];
  
        processors[c].ProcessRow(row, row_left, row_top, row_topleft, xs);
      }
    }
    for (size_t c = 0; c < nb_chans; c++) {
      processors[c].Finalize();
    }
  }
  
  template <size_t nb_chans, size_t bytedepth>
  void WriteACSection(const unsigned char* rgba, size_t x0, size_t y0, size_t oxs,
                      size_t ys, size_t row_stride, bool is_single_group,
                      const PrefixCode& code, std::array<BitWriter, 4>& output) {
    size_t xs = (oxs + kChunkSize - 1) / kChunkSize * kChunkSize;
    for (size_t i = 0; i < nb_chans; i++) {
      if (is_single_group && i == 0) continue;
      output[i].Allocate(16 * xs * ys * bytedepth + 4);
    }
    if (!is_single_group) {
      // Group header for modular image.
      // When the image is single-group, the global modular image is the one that
      // contains the pixel data, and there is no group header.
      output[0].Write(1, 1);     // Global tree
      output[0].Write(1, 1);     // All default wp
      output[0].Write(2, 0b00);  // 0 transforms
    }
  
    ChunkEncoder<bytedepth> encoders[nb_chans];
    ChannelRowProcessor<ChunkEncoder<bytedepth>> row_encoders[nb_chans];
    for (size_t c = 0; c < nb_chans; c++) {
      row_encoders[c].t = &encoders[c];
      encoders[c].output = &output[c];
      encoders[c].code = &code;
    }
    ProcessImageArea<ChannelRowProcessor<ChunkEncoder<bytedepth>>, nb_chans,
                     bytedepth>(rgba, x0, y0, oxs, xs, 0, ys, row_stride,
                                row_encoders);
  }
  
  constexpr int kHashExp = 16;
  constexpr uint32_t kHashSize = 1 << kHashExp;
  constexpr uint32_t kHashMultiplier = 2654435761;
  constexpr int kMaxColors = 512;
  
  // can be any function that returns a value in 0 .. kHashSize-1
  // has to map 0 to 0
  inline uint32_t pixel_hash(uint32_t p) {
    return (p * kHashMultiplier) >> (32 - kHashExp);
  }
  
  template <typename Processor, size_t nb_chans>
  void ProcessImageAreaPalette(const unsigned char* rgba, size_t x0, size_t y0,
                               size_t oxs, size_t xs, size_t yskip, size_t ys,
                               size_t row_stride, const int16_t* lookup,
                               Processor* processors) {
    constexpr size_t kPadding = 16;
  
    int16_t group_data[2][256 + kPadding * 2] = {};
    Processor& row_encoder = processors[0];
  
    for (size_t y = 0; y < ys; y++) {
      // Pre-fill rows with palette converted pixels.
      const unsigned char* inrow = rgba + row_stride * (y0 + y) + x0 * nb_chans;
      for (size_t x = 0; x < oxs; x++) {
        uint32_t p = 0;
        memcpy(&p, inrow + x * nb_chans, nb_chans);
        group_data[y & 1][x + kPadding] = lookup[pixel_hash(p)];
      }
      // Deal with x == 0.
      group_data[y & 1][kPadding - 1] =
          y > 0 ? group_data[(y - 1) & 1][kPadding] : 0;
      // Fix topleft.
      group_data[(y - 1) & 1][kPadding - 1] =
          y > 0 ? group_data[(y - 1) & 1][kPadding] : 0;
      // Fill in padding.
      for (size_t x = oxs; x < xs; x++) {
        group_data[y & 1][kPadding + x] = group_data[y & 1][kPadding + oxs - 1];
      }
      // Get pointers to px/left/top/topleft data to speedup loop.
      const int16_t* row = &group_data[y & 1][kPadding];
      const int16_t* row_left = &group_data[y & 1][kPadding - 1];
      const int16_t* row_top =
          y == 0 ? row_left : &group_data[(y - 1) & 1][kPadding];
      const int16_t* row_topleft =
          y == 0 ? row_left : &group_data[(y - 1) & 1][kPadding - 1];
  
      row_encoder.ProcessRow(row, row_left, row_top, row_topleft, xs);
    }
    row_encoder.Finalize();
  }
  
  template <size_t nb_chans>
  void WriteACSectionPalette(const unsigned char* rgba, size_t x0, size_t y0,
                             size_t oxs, size_t ys, size_t row_stride,
                             bool is_single_group, const PrefixCode& code,
                             const int16_t* lookup, BitWriter& output) {
    size_t xs = (oxs + kChunkSize - 1) / kChunkSize * kChunkSize;
  
    if (!is_single_group) {
      output.Allocate(16 * xs * ys + 4);
      // Group header for modular image.
      // When the image is single-group, the global modular image is the one that
      // contains the pixel data, and there is no group header.
      output.Write(1, 1);     // Global tree
      output.Write(1, 1);     // All default wp
      output.Write(2, 0b00);  // 0 transforms
    }
  
    ChunkEncoder<1> encoder;
    ChannelRowProcessor<ChunkEncoder<1>> row_encoder;
  
    row_encoder.t = &encoder;
    encoder.output = &output;
    encoder.code = &code;
    ProcessImageAreaPalette<ChannelRowProcessor<ChunkEncoder<1>>, nb_chans>(
        rgba, x0, y0, oxs, xs, 0, ys, row_stride, lookup, &row_encoder);
  }
  
  template <size_t nb_chans, size_t bytedepth>
  void CollectSamples(const unsigned char* rgba, size_t x0, size_t y0, size_t xs,
                      size_t row_stride, size_t row_count, uint64_t* raw_counts,
                      uint64_t* lz77_counts, bool palette,
                      const int16_t* lookup) {
    ChunkSampleCollector sample_collectors[nb_chans];
    ChannelRowProcessor<ChunkSampleCollector> row_sample_collectors[nb_chans];
    for (size_t c = 0; c < nb_chans; c++) {
      row_sample_collectors[c].t = &sample_collectors[c];
      sample_collectors[c].raw_counts = raw_counts;
      sample_collectors[c].lz77_counts = lz77_counts;
    }
    if (palette) {
      assert(bytedepth == 1);
      ProcessImageAreaPalette<ChannelRowProcessor<ChunkSampleCollector>,
                              nb_chans>(rgba, x0, y0, xs, xs, 1, 1 + row_count,
                                        row_stride, lookup,
                                        row_sample_collectors);
    } else {
      ProcessImageArea<ChannelRowProcessor<ChunkSampleCollector>, nb_chans,
                       bytedepth>(rgba, x0, y0, xs, xs, 1, 1 + row_count,
                                  row_stride, row_sample_collectors);
    }
  }
  
  void PrepareDCGlobalPalette(bool is_single_group, size_t width, size_t height,
                              const PrefixCode& code,
                              const std::vector<uint32_t>& palette,
                              size_t pcolors_real, BitWriter* output) {
    PrepareDCGlobalCommon(is_single_group, width, height, code, output);
    output->Write(2, 0b01);     // 1 transform
    output->Write(2, 0b01);     // Palette
    output->Write(5, 0b00000);  // Starting from ch 0
    output->Write(2, 0b10);     // 4-channel palette (RGBA)
    size_t pcolors = (pcolors_real + kChunkSize - 1) / kChunkSize * kChunkSize;
    // pcolors <= kMaxColors + kChunkSize - 1
    static_assert(kMaxColors + kChunkSize < 1281,
                  "add code to signal larger palette sizes");
    if (pcolors < 256) {
      output->Write(2, 0b00);
      output->Write(8, pcolors);
    } else {
      output->Write(2, 0b01);
      output->Write(10, pcolors - 256);
    }
  
    output->Write(2, 0b00);  // nb_deltas == 0
    output->Write(4, 0);     // Zero predictor for delta palette
    // Encode palette
    ChunkEncoder<1> encoder;
    ChannelRowProcessor<ChunkEncoder<1>> row_encoder;
    row_encoder.t = &encoder;
    encoder.output = output;
    encoder.code = &code;
    int16_t p[4][32 + 1024] = {};
    uint8_t prgba[4];
    int i = 0;
    int have_zero = 0;
    if (palette[pcolors_real - 1] == 0) have_zero = 1;
    for (; i < pcolors; i++) {
      if (i < pcolors_real) {
        memcpy(prgba, &palette[i], 4);
      }
      p[0][16 + i + have_zero] = prgba[0];
      p[1][16 + i + have_zero] = prgba[1];
      p[2][16 + i + have_zero] = prgba[2];
      p[3][16 + i + have_zero] = prgba[3];
    }
    p[0][15] = 0;
    row_encoder.ProcessRow(p[0] + 16, p[0] + 15, p[0] + 15, p[0] + 15, pcolors);
    p[1][15] = p[0][16];
    p[0][15] = p[0][16];
    row_encoder.ProcessRow(p[1] + 16, p[1] + 15, p[0] + 16, p[0] + 15, pcolors);
    p[2][15] = p[1][16];
    p[1][15] = p[1][16];
    row_encoder.ProcessRow(p[2] + 16, p[2] + 15, p[1] + 16, p[1] + 15, pcolors);
    p[3][15] = p[2][16];
    p[2][15] = p[2][16];
    row_encoder.ProcessRow(p[3] + 16, p[3] + 15, p[2] + 16, p[2] + 15, pcolors);
    row_encoder.Finalize();
  
    if (!is_single_group) {
      output->ZeroPadToByte();
    }
  }
  
  template <size_t nb_chans, size_t bytedepth>
  size_t LLEnc(const unsigned char* rgba, size_t width, size_t stride,
               size_t height, size_t bitdepth, int effort,
               unsigned char** output) {
    size_t bytes_per_sample = (bitdepth > 8 ? 2 : 1);
    assert(bytedepth == bytes_per_sample);
    assert(width != 0);
    assert(height != 0);
    assert(stride >= nb_chans * bytes_per_sample * width);
  
    // Count colors to try palette
    std::vector<uint32_t> palette(kHashSize);
    palette[0] = 1;
    int16_t lookup[kHashSize];
    lookup[0] = 0;
    int pcolors = 0;
    bool collided =
        effort < 2 || bitdepth != 8 || nb_chans < 4;  // todo: also do rgb palette
    for (size_t y = 0; y < height && !collided; y++) {
      const unsigned char* r = rgba + stride * y;
      size_t x = 0;
      if (nb_chans == 4) {
        // this is just an unrolling of the next loop
        for (; x + 7 < width; x += 8) {
          uint32_t p[8], index[8];
          memcpy(p, r + x * 4, 32);
          for (int i = 0; i < 8; i++) index[i] = pixel_hash(p[i]);
          for (int i = 0; i < 8; i++) {
            uint32_t init_entry = index[i] ? 0 : 1;
            if (init_entry != palette[index[i]] && p[i] != palette[index[i]]) {
              collided = true;
            }
          }
          for (int i = 0; i < 8; i++) palette[index[i]] = p[i];
        }
        for (; x < width; x++) {
          uint32_t p;
          memcpy(&p, r + x * 4, 4);
          uint32_t index = pixel_hash(p);
          uint32_t init_entry = index ? 0 : 1;
          if (init_entry != palette[index] && p != palette[index]) {
            collided = true;
          }
          palette[index] = p;
        }
      } else {
        for (; x < width; x++) {
          uint32_t p = 0;
          memcpy(&p, r + x * nb_chans, nb_chans);
          uint32_t index = pixel_hash(p);
          uint32_t init_entry = index ? 0 : 1;
          if (init_entry != palette[index] && p != palette[index]) {
            collided = true;
          }
          palette[index] = p;
        }
      }
    }
  
    int nb_entries = 0;
    if (!collided) {
      if (palette[0] == 0) pcolors = 1;
      if (palette[0] == 1) palette[0] = 0;
      bool have_color = false;
      uint8_t minG = 255, maxG = 0;
      for (int k = 0; k < kHashSize; k++) {
        if (palette[k] == 0) continue;
        uint8_t p[4];
        memcpy(p, &palette[k], 4);
        // move entries to front so sort has less work
        palette[nb_entries] = palette[k];
        if (p[0] != p[1] || p[0] != p[2]) have_color = true;
        if (p[1] < minG) minG = p[1];
        if (p[1] > maxG) maxG = p[1];
        nb_entries++;
        // don't do palette if too many colors are needed
        if (nb_entries + pcolors > kMaxColors) {
          collided = true;
          break;
        }
      }
      if (!have_color) {
        // don't do palette if it's just grayscale without many holes
        if (maxG - minG < nb_entries * 1.4f) collided = true;
      }
    }
    if (!collided) {
      std::sort(
          palette.begin(), palette.begin() + nb_entries,
          [](uint32_t ap, uint32_t bp) {
            if (ap == 0) return false;
            if (bp == 0) return true;
            uint8_t a[4], b[4];
            memcpy(a, &ap, 4);
            memcpy(b, &bp, 4);
            float ay, by;
            ay = (0.299f * a[0] + 0.587f * a[1] + 0.114f * a[2] + 0.01f) * a[3];
            by = (0.299f * b[0] + 0.587f * b[1] + 0.114f * b[2] + 0.01f) * b[3];
            return ay < by;  // sort on alpha*luma
          });
      for (int k = 0; k < nb_entries; k++) {
        if (palette[k] == 0) break;
        lookup[pixel_hash(palette[k])] = pcolors++;
      }
    }
  
    // Width gets padded to kChunkSize, but this computation doesn't change
    // because of that.
    size_t num_groups_x = (width + 255) / 256;
    size_t num_groups_y = (height + 255) / 256;
    size_t num_dc_groups_x = (width + 2047) / 2048;
    size_t num_dc_groups_y = (height + 2047) / 2048;
  
    uint64_t raw_counts[16] = {};
    uint64_t lz77_counts[17] = {};
  
    // sample the middle (effort * 2) rows of every group
    for (size_t g = 0; g < num_groups_y * num_groups_x; g++) {
      size_t xg = g % num_groups_x;
      size_t yg = g / num_groups_x;
      int y_offset = yg * 256;
      int y_max = std::min<size_t>(height - yg * 256, 256);
      int y_begin = y_offset + std::max<int>(0, y_max - 2 * effort) / 2;
      int y_count =
          std::min<int>(2 * effort * y_max / 256, y_offset + y_max - y_begin - 1);
      int x_max =
          std::min<size_t>(width - xg * 256, 256) / kChunkSize * kChunkSize;
      CollectSamples<nb_chans, bytedepth>(rgba, xg * 256, y_begin, x_max, stride,
                                          y_count, raw_counts, lz77_counts,
                                          !collided, lookup);
    }
  
    uint64_t base_raw_counts[16] = {3843, 852, 1270, 1214, 1014, 727, 481, 300,
                                    159,  51,  5,    1,    1,    1,   1,   1};
  
    bool doing_ycocg = nb_chans > 2 && collided;
    for (size_t i = bitdepth + 2 + (doing_ycocg ? 1 : 0); i < 16; i++) {
      base_raw_counts[i] = 0;
    }
    uint64_t base_lz77_counts[17] = {
        // short runs will be sampled, but long ones won't.
        // near full-group run is quite common (e.g. all-opaque alpha)
        18, 12, 9, 11, 15, 2, 2, 1, 1, 1, 1, 2, 300, 0, 0, 0, 0};
  
    for (size_t i = 0; i < 16; i++) {
      raw_counts[i] = (raw_counts[i] << 8) + base_raw_counts[i];
    }
    if (!collided) {
      unsigned token, nbits, bits;
      EncodeHybridUint000(PackSigned(pcolors - 1), &token, &nbits, &bits);
      // ensure all palette indices can actually be encoded
      for (size_t i = 0; i < token + 1; i++)
        raw_counts[i] = std::max<uint64_t>(raw_counts[i], 1);
      // these tokens are only used for the palette itself so they can get a bad
      // code
      for (size_t i = token + 1; i < 10; i++) raw_counts[i] = 1;
    }
    for (size_t i = 0; i < 17; i++) {
      lz77_counts[i] = (lz77_counts[i] << 8) + base_lz77_counts[i];
    }
    alignas(32) PrefixCode hcode(raw_counts, lz77_counts);
  
    BitWriter writer;
  
    bool onegroup = num_groups_x == 1 && num_groups_y == 1;
  
    size_t num_groups = onegroup ? 1
                                 : (2 + num_dc_groups_x * num_dc_groups_y +
                                    num_groups_x * num_groups_y);
  
    std::vector<std::array<BitWriter, 4>> group_data(num_groups);
    if (collided) {
      PrepareDCGlobal(onegroup, width, height, nb_chans, bitdepth, hcode,
                      &group_data[0][0]);
    } else {
      PrepareDCGlobalPalette(onegroup, width, height, hcode, palette, pcolors,
                             &group_data[0][0]);
    }
  #pragma omp parallel for
    for (size_t g = 0; g < num_groups_y * num_groups_x; g++) {
      size_t xg = g % num_groups_x;
      size_t yg = g / num_groups_x;
      size_t group_id =
          onegroup ? 0 : (2 + num_dc_groups_x * num_dc_groups_y + g);
      size_t xs = std::min<size_t>(width - xg * 256, 256);
      size_t ys = std::min<size_t>(height - yg * 256, 256);
      size_t x0 = xg * 256;
      size_t y0 = yg * 256;
      auto& gd = group_data[group_id];
      if (collided) {
        WriteACSection<nb_chans, bytedepth>(rgba, x0, y0, xs, ys, stride,
                                            onegroup, hcode, gd);
  
      } else {
        WriteACSectionPalette<nb_chans>(rgba, x0, y0, xs, ys, stride, onegroup,
                                        hcode, lookup, gd[0]);
      }
    }
  
    AssembleFrame(width, height, nb_chans, bitdepth, group_data, &writer);
  
    *output = writer.data.release();
    return writer.bytes_written;
  }
  
  size_t FastLosslessEncode(const unsigned char* rgba, size_t width,
                            size_t stride, size_t height, size_t c,
                            size_t bitdepth, int effort, unsigned char** output) {
    assert(bitdepth <= 12);
    assert(bitdepth > 0);
    assert(c <= 4);
    assert(c != 0);
    if (bitdepth <= 8) {
      if (c == 1)
        return LLEnc<1, 1>(rgba, width, stride, height, bitdepth, effort, output);
      if (c == 2)
        return LLEnc<2, 1>(rgba, width, stride, height, bitdepth, effort, output);
      if (c == 3)
        return LLEnc<3, 1>(rgba, width, stride, height, bitdepth, effort, output);
      if (c == 4)
        return LLEnc<4, 1>(rgba, width, stride, height, bitdepth, effort, output);
    } else {
      if (c == 1)
        return LLEnc<1, 2>(rgba, width, stride, height, bitdepth, effort, output);
      if (c == 2)
        return LLEnc<2, 2>(rgba, width, stride, height, bitdepth, effort, output);
      if (c == 3)
        return LLEnc<3, 2>(rgba, width, stride, height, bitdepth, effort, output);
      if (c == 4)
        return LLEnc<4, 2>(rgba, width, stride, height, bitdepth, effort, output);
    }
    return 0;
  }
  