Fix the JPEG tile decode issue in the case of JPEGs having restart
markers.

The fix stores the restart information on each index point and restores
the restart information when we do the tile decode.

Bug: 3312406
Change-Id: I5cd8717d4b60ffb5d2e7c97189ee044d67472fe2
diff --git a/jdhuff.c b/jdhuff.c
index 0d704a5..bc5d4fd 100644
--- a/jdhuff.c
+++ b/jdhuff.c
@@ -295,20 +295,10 @@
 		      int nbits)
 /* Load up the bit buffer to a depth of at least nbits */
 {
-  j_decompress_ptr cinfo = state->cinfo;
-  if (cinfo->tile_decode &&
-      cinfo->restart_interval == 0 &&
-      cinfo->unread_marker >= 0xd0 &&
-      cinfo->unread_marker <= 0xd7 &&
-      nbits > bits_left
-      ) {
-      // Skip the restart marker.
-    cinfo->marker->next_restart_num = cinfo->unread_marker - 0xd0;
-    process_restart(cinfo);
-  }
   /* Copy heavily used state fields into locals (hopefully registers) */
   register const JOCTET * next_input_byte = state->next_input_byte;
   register size_t bytes_in_buffer = state->bytes_in_buffer;
+  j_decompress_ptr cinfo = state->cinfo;
 
   /* Attempt to load at least MIN_GET_BITS bits into get_buffer. */
   /* (It is assumed that no request will be for more than that many bits.) */
@@ -509,24 +499,20 @@
 }
 
 /*
- * Configure the Huffman decoder reader position and bit buffer.
+ * Save the current Huffman deocde position and the DC coefficients
+ * for each component into bitstream_offset and dc_info[], respectively.
  */
-GLOBAL(void)
-jpeg_configure_huffman_decoder(j_decompress_ptr cinfo,
-        huffman_offset_data offset)
+METHODDEF(void)
+get_huffman_decoder_configuration(j_decompress_ptr cinfo,
+        huffman_offset_data *offset)
 {
-  unsigned int bitstream_offset = offset.bitstream_offset;
-  int blkn, i;
-
-  cinfo->restart_interval = 0;
-  cinfo->unread_marker = 0;
-
-  unsigned int byte_offset = bitstream_offset >> LOG_TWO_BIT_BUF_SIZE;
-  unsigned int bit_in_bit_buffer =
-      bitstream_offset & ((1 << LOG_TWO_BIT_BUF_SIZE) - 1);
-
-  jset_input_stream_position_bit(cinfo, byte_offset,
-          bit_in_bit_buffer, offset.get_buffer);
+  huff_entropy_ptr entropy = (huff_entropy_ptr) cinfo->entropy;
+  short int *dc_info = offset->prev_dc;
+  int i;
+  jpeg_get_huffman_decoder_configuration(cinfo, offset);
+  for (i = 0; i < cinfo->comps_in_scan; i++) {
+    dc_info[i] = entropy->saved.last_dc_val[i];
+  }
 }
 
 /*
@@ -546,6 +532,10 @@
 	return;
   }
 
+  // Save restarts_to_go and next_restart_num
+  offset->restarts_to_go = (unsigned short) entropy->restarts_to_go;
+  offset->next_restart_num = cinfo->marker->next_restart_num;
+
   offset->bitstream_offset =
       (jget_input_stream_position(cinfo) << LOG_TWO_BIT_BUF_SIZE)
       + entropy->bitstate.bits_left;
@@ -570,20 +560,28 @@
 }
 
 /*
- * Save the current Huffman deocde position and the DC coefficients
- * for each component into bitstream_offset and dc_info[], respectively.
+ * Configure the Huffman decoder reader position and bit buffer.
  */
-METHODDEF(void)
-get_huffman_decoder_configuration(j_decompress_ptr cinfo,
-        huffman_offset_data *offset)
+GLOBAL(void)
+jpeg_configure_huffman_decoder(j_decompress_ptr cinfo,
+        huffman_offset_data offset)
 {
   huff_entropy_ptr entropy = (huff_entropy_ptr) cinfo->entropy;
-  short int *dc_info = offset->prev_dc;
-  int i;
-  jpeg_get_huffman_decoder_configuration(cinfo, offset);
-  for (i = 0; i < cinfo->comps_in_scan; i++) {
-    dc_info[i] = entropy->saved.last_dc_val[i];
-  }
+
+  // Restore restarts_to_go and next_restart_num
+  cinfo->unread_marker = 0;
+  entropy->restarts_to_go = offset.restarts_to_go;
+  cinfo->marker->next_restart_num = offset.next_restart_num;
+
+  unsigned int bitstream_offset = offset.bitstream_offset;
+  int blkn, i;
+
+  unsigned int byte_offset = bitstream_offset >> LOG_TWO_BIT_BUF_SIZE;
+  unsigned int bit_in_bit_buffer =
+      bitstream_offset & ((1 << LOG_TWO_BIT_BUF_SIZE) - 1);
+
+  jset_input_stream_position_bit(cinfo, byte_offset,
+          bit_in_bit_buffer, offset.get_buffer);
 }
 
 /*
diff --git a/jdphuff.c b/jdphuff.c
index a58cdd2..922017e 100644
--- a/jdphuff.c
+++ b/jdphuff.c
@@ -632,21 +632,6 @@
 }
 
 /*
- * Configure the Huffman decoder to decode the image
- * starting from (iMCU_row_offset, iMCU_col_offset).
- */
-METHODDEF(void)
-configure_huffman_decoder(j_decompress_ptr cinfo, huffman_offset_data offset)
-{
-  int i;
-  phuff_entropy_ptr entropy = (phuff_entropy_ptr) cinfo->entropy;
-  jpeg_configure_huffman_decoder(cinfo, offset);
-  entropy->saved.EOBRUN = offset.EOBRUN;
-  for (i = 0; i < cinfo->comps_in_scan; i++)
-    entropy->saved.last_dc_val[i] = offset.prev_dc[i];
-}
-
-/*
  * Save the current Huffman deocde position and the DC coefficients
  * for each component into bitstream_offset and dc_info[], respectively.
  */
@@ -656,16 +641,103 @@
 {
   int i;
   phuff_entropy_ptr entropy = (phuff_entropy_ptr) cinfo->entropy;
-  jpeg_get_huffman_decoder_configuration(cinfo, offset);
+  jpeg_get_huffman_decoder_configuration_progressive(cinfo, offset);
   offset->EOBRUN = entropy->saved.EOBRUN;
   for (i = 0; i < cinfo->comps_in_scan; i++)
     offset->prev_dc[i] = entropy->saved.last_dc_val[i];
 }
 
+
+/*
+ * Save the current Huffman decoder position and the bit buffer
+ * into bitstream_offset and get_buffer, respectively.
+ */
+GLOBAL(void)
+jpeg_get_huffman_decoder_configuration_progressive(j_decompress_ptr cinfo,
+        huffman_offset_data *offset)
+{
+  phuff_entropy_ptr entropy = (phuff_entropy_ptr) cinfo->entropy;
+
+  if (cinfo->restart_interval) {
+    // We are at the end of a data segment
+    if (entropy->restarts_to_go == 0)
+      if (! process_restart(cinfo))
+	return;
+  }
+
+  // Save restarts_to_go and next_restart_num.
+  offset->restarts_to_go = (unsigned short) entropy->restarts_to_go;
+  offset->next_restart_num = cinfo->marker->next_restart_num;
+
+  offset->bitstream_offset =
+      (jget_input_stream_position(cinfo) << LOG_TWO_BIT_BUF_SIZE)
+      + entropy->bitstate.bits_left;
+
+  offset->get_buffer = entropy->bitstate.get_buffer;
+}
+
+
+/*
+ * Configure the Huffman decoder to decode the image
+ * starting from (iMCU_row_offset, iMCU_col_offset).
+ */
+METHODDEF(void)
+configure_huffman_decoder(j_decompress_ptr cinfo, huffman_offset_data offset)
+{
+  int i;
+  phuff_entropy_ptr entropy = (phuff_entropy_ptr) cinfo->entropy;
+  jpeg_configure_huffman_decoder_progressive(cinfo, offset);
+  entropy->saved.EOBRUN = offset.EOBRUN;
+  for (i = 0; i < cinfo->comps_in_scan; i++)
+    entropy->saved.last_dc_val[i] = offset.prev_dc[i];
+}
+
+/*
+ * Configure the Huffman decoder reader position and bit buffer.
+ */
+GLOBAL(void)
+jpeg_configure_huffman_decoder_progressive(j_decompress_ptr cinfo,
+        huffman_offset_data offset)
+{
+	phuff_entropy_ptr entropy = (phuff_entropy_ptr) cinfo->entropy;
+
+  // Restore restarts_to_go and next_restart_num
+  cinfo->unread_marker = 0;
+  entropy->restarts_to_go = offset.restarts_to_go;
+  cinfo->marker->next_restart_num = offset.next_restart_num;
+
+  unsigned int bitstream_offset = offset.bitstream_offset;
+  int blkn, i;
+
+  unsigned int byte_offset = bitstream_offset >> LOG_TWO_BIT_BUF_SIZE;
+  unsigned int bit_in_bit_buffer =
+      bitstream_offset & ((1 << LOG_TWO_BIT_BUF_SIZE) - 1);
+
+  jset_input_stream_position_bit(cinfo, byte_offset,
+          bit_in_bit_buffer, offset.get_buffer);
+}
+
+GLOBAL(void)
+jpeg_configure_huffman_index_scan(j_decompress_ptr cinfo,
+        huffman_index *index, int scan_no, int offset)
+{
+  phuff_entropy_ptr entropy = (phuff_entropy_ptr) cinfo->entropy;
+  if (scan_no >= index->scan_count) {
+    index->scan = realloc(index->scan,
+                    (scan_no + 1) * sizeof(huffman_scan_header));
+    index->mem_used += (scan_no - index->scan_count + 1)
+      * (sizeof(huffman_scan_header) + cinfo->total_iMCU_rows
+      * sizeof(huffman_offset_data*));
+    index->scan_count = scan_no + 1;
+  }
+  index->scan[scan_no].offset = (huffman_offset_data**)malloc(
+          cinfo->total_iMCU_rows * sizeof(huffman_offset_data*));
+  index->scan[scan_no].bitstream_offset = offset;
+}
+
 /*
  * Module initialization routine for progressive Huffman entropy decoding.
  */
-
 GLOBAL(void)
 jinit_phuff_decoder (j_decompress_ptr cinfo)
 {
@@ -697,22 +769,4 @@
       *coef_bit_ptr++ = -1;
 }
 
-GLOBAL(void)
-jpeg_configure_huffman_index_scan(j_decompress_ptr cinfo,
-        huffman_index *index, int scan_no, int offset)
-{
-  phuff_entropy_ptr entropy = (phuff_entropy_ptr) cinfo->entropy;
-  if (scan_no >= index->scan_count) {
-    index->scan = realloc(index->scan,
-                    (scan_no + 1) * sizeof(huffman_scan_header));
-    index->mem_used += (scan_no - index->scan_count + 1)
-      * (sizeof(huffman_scan_header) + cinfo->total_iMCU_rows
-      * sizeof(huffman_offset_data*));
-    index->scan_count = scan_no + 1;
-  }
-  index->scan[scan_no].offset = (huffman_offset_data**)malloc(
-          cinfo->total_iMCU_rows * sizeof(huffman_offset_data*));
-  index->scan[scan_no].bitstream_offset = offset;
-}
-
 #endif /* D_PROGRESSIVE_SUPPORTED */
diff --git a/jpeglib.h b/jpeglib.h
index 83bed4a..07e6872 100644
--- a/jpeglib.h
+++ b/jpeglib.h
@@ -649,6 +649,10 @@
 
   // save the decoder current bit buffer, entropy->bitstate.get_buffer.
   INT32 get_buffer;
+
+  // save the restart info.
+  unsigned short restarts_to_go;
+  unsigned char next_restart_num;
 } huffman_offset_data;
 
 typedef struct {