lib/jpegli/bitstream.cc

   1 // Copyright (c) the JPEG XL Project Authors. All rights reserved.
   2 //
   3 // Use of this source code is governed by a BSD-style
   4 // license that can be found in the LICENSE file.
   5
   6 #include "lib/jpegli/bitstream.h"
   7
   8 #include <cmath>
   9
  10 #include "lib/jpegli/bit_writer.h"
  11 #include "lib/jpegli/error.h"
  12 #include "lib/jpegli/memory_manager.h"
  13
  14 namespace jpegli {
  15
  16 void WriteOutput(j_compress_ptr cinfo, const uint8_t* buf, size_t bufsize) {
  17   size_t pos = 0;
  18   while (pos < bufsize) {
  19     if (cinfo->dest->free_in_buffer == 0 &&
  20         !(*cinfo->dest->empty_output_buffer)(cinfo)) {
  21       JPEGLI_ERROR("Destination suspension is not supported in markers.");
  22     }
  23     size_t len = std::min<size_t>(cinfo->dest->free_in_buffer, bufsize - pos);
  24     memcpy(cinfo->dest->next_output_byte, buf + pos, len);
  25     pos += len;
  26     cinfo->dest->free_in_buffer -= len;
  27     cinfo->dest->next_output_byte += len;
  28   }
  29 }
  30
  31 void WriteOutput(j_compress_ptr cinfo, const std::vector<uint8_t>& bytes) {
  32   WriteOutput(cinfo, bytes.data(), bytes.size());
  33 }
  34
  35 void WriteOutput(j_compress_ptr cinfo, std::initializer_list<uint8_t> bytes) {
  36   WriteOutput(cinfo, bytes.begin(), bytes.size());
  37 }
  38
  39 void EncodeAPP0(j_compress_ptr cinfo) {
  40   WriteOutput(cinfo,
  41               {0xff, 0xe0, 0, 16, 'J', 'F', 'I', 'F', '\0',
  42                cinfo->JFIF_major_version, cinfo->JFIF_minor_version,
  43                cinfo->density_unit, static_cast<uint8_t>(cinfo->X_density >> 8),
  44                static_cast<uint8_t>(cinfo->X_density & 0xff),
  45                static_cast<uint8_t>(cinfo->Y_density >> 8),
  46                static_cast<uint8_t>(cinfo->Y_density & 0xff), 0, 0});
  47 }
  48
  49 void EncodeAPP14(j_compress_ptr cinfo) {
  50   uint8_t color_transform = cinfo->jpeg_color_space == JCS_YCbCr  ? 1
  51                             : cinfo->jpeg_color_space == JCS_YCCK ? 2
  52                                                                   : 0;
  53   WriteOutput(cinfo, {0xff, 0xee, 0, 14, 'A', 'd', 'o', 'b', 'e', 0, 100, 0, 0,
  54                       0, 0, color_transform});
  55 }
  56
  57 void WriteFileHeader(j_compress_ptr cinfo) {
  58   WriteOutput(cinfo, {0xFF, 0xD8});  // SOI
  59   if (cinfo->write_JFIF_header) {
  60     EncodeAPP0(cinfo);
  61   }
  62   if (cinfo->write_Adobe_marker) {
  63     EncodeAPP14(cinfo);
  64   }
  65 }
  66
  67 bool EncodeDQT(j_compress_ptr cinfo, bool write_all_tables) {
  68   uint8_t data[4 + NUM_QUANT_TBLS * (1 + 2 * DCTSIZE2)];  // 520 bytes
  69   size_t pos = 0;
  70   data[pos++] = 0xFF;
  71   data[pos++] = 0xDB;
  72   pos += 2;  // Length will be filled in later.
  73
  74   int send_table[NUM_QUANT_TBLS] = {};
  75   if (write_all_tables) {
  76     for (int i = 0; i < NUM_QUANT_TBLS; ++i) {
  77       if (cinfo->quant_tbl_ptrs[i]) send_table[i] = 1;
  78     }
  79   } else {
  80     for (int c = 0; c < cinfo->num_components; ++c) {
  81       send_table[cinfo->comp_info[c].quant_tbl_no] = 1;
  82     }
  83   }
  84
  85   bool is_baseline = true;
  86   for (int i = 0; i < NUM_QUANT_TBLS; ++i) {
  87     if (!send_table[i]) continue;
  88     JQUANT_TBL* quant_table = cinfo->quant_tbl_ptrs[i];
  89     if (quant_table == nullptr) {
  90       JPEGLI_ERROR("Missing quant table %d", i);
  91     }
  92     int precision = 0;
  93     for (size_t k = 0; k < DCTSIZE2; ++k) {
  94       if (quant_table->quantval[k] > 255) {
  95         precision = 1;
  96         is_baseline = false;
  97       }
  98     }
  99     if (quant_table->sent_table) {
 100       continue;
 101     }
 102     data[pos++] = (precision << 4) + i;
 103     for (size_t j = 0; j < DCTSIZE2; ++j) {
 104       int val_idx = kJPEGNaturalOrder[j];
 105       int val = quant_table->quantval[val_idx];
 106       if (val == 0) {
 107         JPEGLI_ERROR("Invalid quantval 0.");
 108       }
 109       if (precision) {
 110         data[pos++] = val >> 8;
 111       }
 112       data[pos++] = val & 0xFFu;
 113     }
 114     quant_table->sent_table = TRUE;
 115   }
 116   if (pos > 4) {
 117     data[2] = (pos - 2) >> 8u;
 118     data[3] = (pos - 2) & 0xFFu;
 119     WriteOutput(cinfo, data, pos);
 120   }
 121   return is_baseline;
 122 }
 123
 124 void EncodeSOF(j_compress_ptr cinfo, bool is_baseline) {
 125   if (cinfo->data_precision != kJpegPrecision) {
 126     is_baseline = false;
 127     JPEGLI_ERROR("Unsupported data precision %d", cinfo->data_precision);
 128   }
 129   const uint8_t marker = cinfo->progressive_mode ? 0xc2
 130                          : is_baseline           ? 0xc0
 131                                                  : 0xc1;
 132   const size_t n_comps = cinfo->num_components;
 133   const size_t marker_len = 8 + 3 * n_comps;
 134   std::vector<uint8_t> data(marker_len + 2);
 135   size_t pos = 0;
 136   data[pos++] = 0xFF;
 137   data[pos++] = marker;
 138   data[pos++] = marker_len >> 8u;
 139   data[pos++] = marker_len & 0xFFu;
 140   data[pos++] = kJpegPrecision;
 141   data[pos++] = cinfo->image_height >> 8u;
 142   data[pos++] = cinfo->image_height & 0xFFu;
 143   data[pos++] = cinfo->image_width >> 8u;
 144   data[pos++] = cinfo->image_width & 0xFFu;
 145   data[pos++] = n_comps;
 146   for (size_t i = 0; i < n_comps; ++i) {
 147     jpeg_component_info* comp = &cinfo->comp_info[i];
 148     data[pos++] = comp->component_id;
 149     data[pos++] = ((comp->h_samp_factor << 4u) | (comp->v_samp_factor));
 150     const uint32_t quant_idx = comp->quant_tbl_no;
 151     if (cinfo->quant_tbl_ptrs[quant_idx] == nullptr) {
 152       JPEGLI_ERROR("Invalid component quant table index %u.", quant_idx);
 153     }
 154     data[pos++] = quant_idx;
 155   }
 156   WriteOutput(cinfo, data);
 157 }
 158
 159 void WriteFrameHeader(j_compress_ptr cinfo) {
 160   jpeg_comp_master* m = cinfo->master;
 161   bool is_baseline = EncodeDQT(cinfo, /*write_all_tables=*/false);
 162   if (cinfo->progressive_mode || cinfo->arith_code ||
 163       cinfo->data_precision != 8) {
 164     is_baseline = false;
 165   }
 166   for (size_t i = 0; i < m->num_huffman_tables; ++i) {
 167     int slot_id = m->slot_id_map[i];
 168     if (slot_id > 0x11 || (slot_id > 0x01 && slot_id < 0x10)) {
 169       is_baseline = false;
 170     }
 171   }
 172   EncodeSOF(cinfo, is_baseline);
 173 }
 174
 175 void EncodeDRI(j_compress_ptr cinfo) {
 176   WriteOutput(cinfo, {0xFF, 0xDD, 0, 4,
 177                       static_cast<uint8_t>(cinfo->restart_interval >> 8),
 178                       static_cast<uint8_t>(cinfo->restart_interval & 0xFF)});
 179 }
 180
 181 void EncodeDHT(j_compress_ptr cinfo, size_t offset, size_t num) {
 182   jpeg_comp_master* m = cinfo->master;
 183   size_t marker_len = 2;
 184   for (size_t i = 0; i < num; ++i) {
 185     const JHUFF_TBL& table = m->huffman_tables[offset + i];
 186     if (table.sent_table) continue;
 187     marker_len += kJpegHuffmanMaxBitLength + 1;
 188     for (size_t j = 0; j <= kJpegHuffmanMaxBitLength; ++j) {
 189       marker_len += table.bits[j];
 190     }
 191   }
 192   std::vector<uint8_t> data(marker_len + 2);
 193   size_t pos = 0;
 194   data[pos++] = 0xFF;
 195   data[pos++] = 0xC4;
 196   data[pos++] = marker_len >> 8u;
 197   data[pos++] = marker_len & 0xFFu;
 198   for (size_t i = 0; i < num; ++i) {
 199     const JHUFF_TBL& table = m->huffman_tables[offset + i];
 200     if (table.sent_table) continue;
 201     size_t total_count = 0;
 202     for (size_t i = 0; i <= kJpegHuffmanMaxBitLength; ++i) {
 203       total_count += table.bits[i];
 204     }
 205     data[pos++] = m->slot_id_map[offset + i];
 206     for (size_t i = 1; i <= kJpegHuffmanMaxBitLength; ++i) {
 207       data[pos++] = table.bits[i];
 208     }
 209     for (size_t i = 0; i < total_count; ++i) {
 210       data[pos++] = table.huffval[i];
 211     }
 212   }
 213   if (marker_len > 2) {
 214     WriteOutput(cinfo, data);
 215   }
 216 }
 217
 218 void EncodeSOS(j_compress_ptr cinfo, int scan_index) {
 219   jpeg_comp_master* m = cinfo->master;
 220   const jpeg_scan_info* scan_info = &cinfo->scan_info[scan_index];
 221   const size_t marker_len = 6 + 2 * scan_info->comps_in_scan;
 222   std::vector<uint8_t> data(marker_len + 2);
 223   size_t pos = 0;
 224   data[pos++] = 0xFF;
 225   data[pos++] = 0xDA;
 226   data[pos++] = marker_len >> 8u;
 227   data[pos++] = marker_len & 0xFFu;
 228   data[pos++] = scan_info->comps_in_scan;
 229   for (int i = 0; i < scan_info->comps_in_scan; ++i) {
 230     int comp_idx = scan_info->component_index[i];
 231     data[pos++] = cinfo->comp_info[comp_idx].component_id;
 232     int dc_slot_id = m->slot_id_map[m->context_map[comp_idx]];
 233     int ac_context = m->ac_ctx_offset[scan_index] + i;
 234     int ac_slot_id = m->slot_id_map[m->context_map[ac_context]];
 235     data[pos++] = (dc_slot_id << 4u) + (ac_slot_id - 16);
 236   }
 237   data[pos++] = scan_info->Ss;
 238   data[pos++] = scan_info->Se;
 239   data[pos++] = ((scan_info->Ah << 4u) | (scan_info->Al));
 240   WriteOutput(cinfo, data);
 241 }
 242
 243 void WriteScanHeader(j_compress_ptr cinfo, int scan_index) {
 244   jpeg_comp_master* m = cinfo->master;
 245   const jpeg_scan_info* scan_info = &cinfo->scan_info[scan_index];
 246   cinfo->restart_interval = m->scan_token_info[scan_index].restart_interval;
 247   if (cinfo->restart_interval != m->last_restart_interval) {
 248     EncodeDRI(cinfo);
 249     m->last_restart_interval = cinfo->restart_interval;
 250   }
 251   size_t num_dht = 0;
 252   if (scan_index == 0) {
 253     // For the first scan we emit all DC and at most 4 AC Huffman codes.
 254     for (size_t i = 0, num_ac = 0; i < m->num_huffman_tables; ++i) {
 255       if (m->slot_id_map[i] >= 16 && num_ac++ >= 4) break;
 256       ++num_dht;
 257     }
 258   } else if (scan_info->Ss > 0) {
 259     // For multi-scan sequential and progressive DC scans we have already
 260     // emitted all Huffman codes that we need before the first scan. For
 261     // progressive AC scans we only need at most one new Huffman code.
 262     if (m->context_map[m->ac_ctx_offset[scan_index]] == m->next_dht_index) {
 263       num_dht = 1;
 264     }
 265   }
 266   if (num_dht > 0) {
 267     EncodeDHT(cinfo, m->next_dht_index, num_dht);
 268     m->next_dht_index += num_dht;
 269   }
 270   EncodeSOS(cinfo, scan_index);
 271 }
 272
 273 void WriteBlock(const int32_t* JXL_RESTRICT symbols,
 274                 const int32_t* JXL_RESTRICT extra_bits, const int num_nonzeros,
 275                 const bool emit_eob,
 276                 const HuffmanCodeTable* JXL_RESTRICT dc_code,
 277                 const HuffmanCodeTable* JXL_RESTRICT ac_code,
 278                 JpegBitWriter* JXL_RESTRICT bw) {
 279   int symbol = symbols[0];
 280   WriteBits(bw, dc_code->depth[symbol], dc_code->code[symbol] | extra_bits[0]);
 281   for (int i = 1; i < num_nonzeros; ++i) {
 282     symbol = symbols[i];
 283     if (symbol > 255) {
 284       WriteBits(bw, ac_code->depth[0xf0], ac_code->code[0xf0]);
 285       symbol -= 256;
 286       if (symbol > 255) {
 287         WriteBits(bw, ac_code->depth[0xf0], ac_code->code[0xf0]);
 288         symbol -= 256;
 289         if (symbol > 255) {
 290           WriteBits(bw, ac_code->depth[0xf0], ac_code->code[0xf0]);
 291           symbol -= 256;
 292         }
 293       }
 294     }
 295     WriteBits(bw, ac_code->depth[symbol],
 296               ac_code->code[symbol] | extra_bits[i]);
 297   }
 298   if (emit_eob) {
 299     WriteBits(bw, ac_code->depth[0], ac_code->code[0]);
 300   }
 301 }
 302
 303 namespace {
 304
 305 static JXL_INLINE void EmitMarker(JpegBitWriter* bw, int marker) {
 306   bw->data[bw->pos++] = 0xFF;
 307   bw->data[bw->pos++] = marker;
 308 }
 309
 310 void WriteTokens(j_compress_ptr cinfo, int scan_index, JpegBitWriter* bw) {
 311   jpeg_comp_master* m = cinfo->master;
 312   HuffmanCodeTable* coding_tables = &m->coding_tables[0];
 313   int next_restart_marker = 0;
 314   const ScanTokenInfo& sti = m->scan_token_info[scan_index];
 315   size_t num_token_arrays = m->cur_token_array + 1;
 316   size_t total_tokens = 0;
 317   size_t restart_idx = 0;
 318   size_t next_restart = sti.restarts[restart_idx];
 319   uint8_t* context_map = m->context_map;
 320   for (size_t i = 0; i < num_token_arrays; ++i) {
 321     Token* tokens = m->token_arrays[i].tokens;
 322     size_t num_tokens = m->token_arrays[i].num_tokens;
 323     if (sti.token_offset < total_tokens + num_tokens &&
 324         total_tokens < sti.token_offset + sti.num_tokens) {
 325       size_t start_ix =
 326           total_tokens < sti.token_offset ? sti.token_offset - total_tokens : 0;
 327       size_t end_ix = std::min(sti.token_offset + sti.num_tokens - total_tokens,
 328                                num_tokens);
 329       size_t cycle_len = bw->len / 8;
 330       size_t next_cycle = cycle_len;
 331       for (size_t i = start_ix; i < end_ix; ++i) {
 332         if (total_tokens + i == next_restart) {
 333           JumpToByteBoundary(bw);
 334           EmitMarker(bw, 0xD0 + next_restart_marker);
 335           next_restart_marker += 1;
 336           next_restart_marker &= 0x7;
 337           next_restart = sti.restarts[++restart_idx];
 338         }
 339         Token t = tokens[i];
 340         const HuffmanCodeTable* code = &coding_tables[context_map[t.context]];
 341         WriteBits(bw, code->depth[t.symbol], code->code[t.symbol] | t.bits);
 342         if (--next_cycle == 0) {
 343           if (!EmptyBitWriterBuffer(bw)) {
 344             JPEGLI_ERROR(
 345                 "Output suspension is not supported in "
 346                 "finish_compress");
 347           }
 348           next_cycle = cycle_len;
 349         }
 350       }
 351     }
 352     total_tokens += num_tokens;
 353   }
 354 }
 355
 356 void WriteACRefinementTokens(j_compress_ptr cinfo, int scan_index,
 357                              JpegBitWriter* bw) {
 358   jpeg_comp_master* m = cinfo->master;
 359   const ScanTokenInfo& sti = m->scan_token_info[scan_index];
 360   const uint8_t context = m->ac_ctx_offset[scan_index];
 361   const HuffmanCodeTable* code = &m->coding_tables[m->context_map[context]];
 362   size_t cycle_len = bw->len / 64;
 363   size_t next_cycle = cycle_len;
 364   size_t refbit_idx = 0;
 365   size_t eobrun_idx = 0;
 366   size_t restart_idx = 0;
 367   size_t next_restart = sti.restarts[restart_idx];
 368   int next_restart_marker = 0;
 369   for (size_t i = 0; i < sti.num_tokens; ++i) {
 370     if (i == next_restart) {
 371       JumpToByteBoundary(bw);
 372       EmitMarker(bw, 0xD0 + next_restart_marker);
 373       next_restart_marker += 1;
 374       next_restart_marker &= 0x7;
 375       next_restart = sti.restarts[++restart_idx];
 376     }
 377     RefToken t = sti.tokens[i];
 378     int symbol = t.symbol & 253;
 379     uint16_t bits = 0;
 380     if ((symbol & 1) == 0) {
 381       int r = symbol >> 4;
 382       if (r > 0 && r < 15) {
 383         bits = sti.eobruns[eobrun_idx++];
 384       }
 385     } else {
 386       bits = (t.symbol >> 1) & 1;
 387     }
 388     WriteBits(bw, code->depth[symbol], code->code[symbol] | bits);
 389     for (int j = 0; j < t.refbits; ++j) {
 390       WriteBits(bw, 1, sti.refbits[refbit_idx++]);
 391     }
 392     if (--next_cycle == 0) {
 393       if (!EmptyBitWriterBuffer(bw)) {
 394         JPEGLI_ERROR("Output suspension is not supported in finish_compress");
 395       }
 396       next_cycle = cycle_len;
 397     }
 398   }
 399 }
 400
 401 void WriteDCRefinementBits(j_compress_ptr cinfo, int scan_index,
 402                            JpegBitWriter* bw) {
 403   jpeg_comp_master* m = cinfo->master;
 404   const ScanTokenInfo& sti = m->scan_token_info[scan_index];
 405   size_t restart_idx = 0;
 406   size_t next_restart = sti.restarts[restart_idx];
 407   int next_restart_marker = 0;
 408   size_t cycle_len = bw->len * 4;
 409   size_t next_cycle = cycle_len;
 410   size_t refbit_idx = 0;
 411   for (size_t i = 0; i < sti.num_tokens; ++i) {
 412     if (i == next_restart) {
 413       JumpToByteBoundary(bw);
 414       EmitMarker(bw, 0xD0 + next_restart_marker);
 415       next_restart_marker += 1;
 416       next_restart_marker &= 0x7;
 417       next_restart = sti.restarts[++restart_idx];
 418     }
 419     WriteBits(bw, 1, sti.refbits[refbit_idx++]);
 420     if (--next_cycle == 0) {
 421       if (!EmptyBitWriterBuffer(bw)) {
 422         JPEGLI_ERROR(
 423             "Output suspension is not supported in "
 424             "finish_compress");
 425       }
 426       next_cycle = cycle_len;
 427     }
 428   }
 429 }
 430
 431 }  // namespace
 432
 433 void WriteScanData(j_compress_ptr cinfo, int scan_index) {
 434   const jpeg_scan_info* scan_info = &cinfo->scan_info[scan_index];
 435   JpegBitWriter* bw = &cinfo->master->bw;
 436   if (scan_info->Ah == 0) {
 437     WriteTokens(cinfo, scan_index, bw);
 438   } else if (scan_info->Ss > 0) {
 439     WriteACRefinementTokens(cinfo, scan_index, bw);
 440   } else {
 441     WriteDCRefinementBits(cinfo, scan_index, bw);
 442   }
 443   if (!bw->healthy) {
 444     JPEGLI_ERROR("Unknown Huffman coded symbol found in scan %d", scan_index);
 445   }
 446   JumpToByteBoundary(bw);
 447   if (!EmptyBitWriterBuffer(bw)) {
 448     JPEGLI_ERROR("Output suspension is not supported in finish_compress");
 449   }
 450 }
 451
 452 }  // namespace jpegli