predict idat size correctly for interlaced images

3 files changed, 99 insertions, 7 deletions

diff --git a/example_png_info.cpp b/example_png_info.cpp
index e7f78a2..c6df5c1 100644
--- a/example_png_info.cpp
+++ b/example_png_info.cpp
@@ -280,12 +280,18 @@ Main
 */
 int main(int argc, char *argv[]) /*list the chunks*/
 {
-  if(argc < 2)
+  bool ignore_checksums = false;
+  std::string filename = "";
+  for (int i = 1; i < argc; i++)
+  {
+    if(std::string(argv[i]) == "--ignore_checksums") ignore_checksums = true;
+    else filename = argv[i];
+  }
+  if(filename == "")
   {
     std::cout << "Please provide a filename to preview" << std::endl;
     return 0;
   }
-  const char* filename = argv[1];
 
   std::vector<unsigned char> buffer;
   std::vector<unsigned char> image;
@@ -294,6 +300,12 @@ int main(int argc, char *argv[]) /*list the chunks*/
   lodepng::load_file(buffer, filename); //load the image file with given filename
 
   lodepng::State state;
+  if(ignore_checksums)
+  {
+    state.decoder.ignore_crc = 1;
+    state.decoder.zlibsettings.ignore_adler32 = 1;
+  }
+
   unsigned error = lodepng::decode(image, w, h, state, buffer);
 
   if(error)
diff --git a/lodepng.cpp b/lodepng.cpp
index 81044eb..ee9f168 100644
--- a/lodepng.cpp
+++ b/lodepng.cpp
@@ -4564,22 +4564,40 @@ static void decodeGeneric(unsigned char** out, unsigned* w, unsigned* h,
 
   ucvector_init(&scanlines);
   /*predict output size, to allocate exact size for output buffer to avoid more dynamic allocation.
-  The prediction is currently not correct for interlaced PNG images.*/
-  predict = lodepng_get_raw_size_idat(*w, *h, &state->info_png.color) + *h;
+  If the decompressed size does not match the prediction, the image must be corrupt.*/
+  if(state->info_png.interlace_method == 0)
+  {
+    /*The extra *h is added because this are the filter bytes every scanline starts with*/
+    predict = lodepng_get_raw_size_idat(*w, *h, &state->info_png.color) + *h;
+  }
+  else
+  {
+    /*Adam-7 interlaced: predicted size is the sum of the 7 sub-images sizes*/
+    const LodePNGColorMode* color = &state->info_png.color;
+    predict = 0;
+    predict += lodepng_get_raw_size_idat((*w + 7) / 8, (*h + 7) / 8, color) + (*h + 7) / 8;
+    if(*w > 4) predict += lodepng_get_raw_size_idat((*w + 3) / 8, (*h + 7) / 8, color) + (*h + 7) / 8;
+    predict += lodepng_get_raw_size_idat((*w + 3) / 4, (*h + 3) / 8, color) + (*h + 3) / 8;
+    if(*w > 2) predict += lodepng_get_raw_size_idat((*w + 1) / 4, (*h + 3) / 4, color) + (*h + 3) / 4;
+    predict += lodepng_get_raw_size_idat((*w + 1) / 2, (*h + 1) / 4, color) + (*h + 1) / 4;
+    if(*w > 1) predict += lodepng_get_raw_size_idat((*w + 0) / 2, (*h + 1) / 2, color) + (*h + 1) / 2;
+    predict += lodepng_get_raw_size_idat((*w + 0) / 1, (*h + 0) / 2, color) + (*h + 0) / 2;
+  }
   if(!state->error && !ucvector_reserve(&scanlines, predict)) state->error = 83; /*alloc fail*/
   if(!state->error)
   {
     state->error = zlib_decompress(&scanlines.data, &scanlines.size, idat.data,
                                    idat.size, &state->decoder.zlibsettings);
+    if(!state->error && scanlines.size != predict) state->error = 91; /*decompressed size doesn't match prediction*/
   }
   ucvector_cleanup(&idat);
 
   if(!state->error)
   {
+    size_t outsize = lodepng_get_raw_size(*w, *h, &state->info_png.color);
     ucvector outv;
     ucvector_init(&outv);
-    if(!ucvector_resizev(&outv,
-        lodepng_get_raw_size(*w, *h, &state->info_png.color), 0)) state->error = 83; /*alloc fail*/
+    if(!ucvector_resizev(&outv, outsize, 0)) state->error = 83; /*alloc fail*/
     if(!state->error) state->error = postProcessScanlines(outv.data, scanlines.data, *w, *h, &state->info_png);
     *out = outv.data;
   }
@@ -5861,6 +5879,7 @@ const char* lodepng_error_text(unsigned code)
     case 89: return "text chunk keyword too short or long: must have size 1-79";
     /*the windowsize in the LodePNGCompressSettings. Requiring POT(==> & instead of %) makes encoding 12% faster.*/
     case 90: return "windowsize must be a power of two";
+    case 91: return "invalid decompressed idat size";
   }
   return "unknown error code";
 }
diff --git a/lodepng_unittest.cpp b/lodepng_unittest.cpp
index d3b3498..a633a97 100644
--- a/lodepng_unittest.cpp
+++ b/lodepng_unittest.cpp
@@ -355,11 +355,44 @@ void doCodecTestCPP(Image& image)
   assertPixels(image, &decoded[0], "Pixels C++");
 }
 
+//Test LodePNG encoding and decoding the encoded result, using the C++ interface, with interlace
+void doCodecTestInterlaced(Image& image)
+{
+  std::vector<unsigned char> encoded;
+  std::vector<unsigned char> decoded;
+  unsigned decoded_w;
+  unsigned decoded_h;
+
+  lodepng::State state;
+  state.info_png.interlace_method = 1;
+  state.info_raw.colortype = image.colorType;
+  state.info_raw.bitdepth = image.bitDepth;
+
+  unsigned error_enc = lodepng::encode(encoded, image.data, image.width, image.height, state);
+
+  assertNoPNGError(error_enc, "encoder error C++");
+
+  //if the image is large enough, compressing it should result in smaller size
+  if(image.data.size() > 512) assertTrue(encoded.size() < image.data.size(), "compressed size");
+
+  state.info_raw.colortype = image.colorType;
+  state.info_raw.bitdepth = image.bitDepth;
+  unsigned error_dec = lodepng::decode(decoded, decoded_w, decoded_h, state, encoded);
+
+  assertNoPNGError(error_dec, "decoder error C++");
+
+  ASSERT_EQUALS(image.width, decoded_w);
+  ASSERT_EQUALS(image.height, decoded_h);
+  ASSERT_EQUALS(image.data.size(), decoded.size());
+  assertPixels(image, &decoded[0], "Pixels C++");
+}
+
 //Test LodePNG encoding and decoding the encoded result
 void doCodecTest(Image& image)
 {
   doCodecTestC(image);
   doCodecTestCPP(image);
+  doCodecTestInterlaced(image);
 }
 
 
@@ -538,6 +571,27 @@ void testColor(int r, int g, int b, int a)
   testSinglePixel(r, g, b, a);
 }
 
+void testSize(int w, int h)
+{
+  std::cout << "codec test size " << w << " " << h << std::endl;
+  Image image;
+  image.width = w;
+  image.height = h;
+  image.colorType = LCT_RGBA;
+  image.bitDepth = 8;
+  image.data.resize(w * h * 4);
+  for(int y = 0; y < h; y++)
+  for(int x = 0; x < w; x++)
+  {
+    image.data[w * 4 * y + 4 * x + 0] = x % 256;
+    image.data[w * 4 * y + 4 * x + 0] = y % 256;
+    image.data[w * 4 * y + 4 * x + 0] = 255;
+    image.data[w * 4 * y + 4 * x + 0] = 255;
+  }
+
+  doCodecTest(image);
+}
+
 void testPNGCodec()
 {
   codecTest(1, 1);
@@ -566,11 +620,18 @@ void testPNGCodec()
   testColor(127, 127, 127, 128);
   testColor(128, 128, 128, 128);
   //transparent single pixels
+  testColor(0, 0, 0, 0);
   testColor(255, 0, 0, 0);
   testColor(1, 2, 3, 0);
   testColor(255, 255, 255, 0);
   testColor(254, 254, 254, 0);
-  testColor(0, 0, 0, 0);
+
+  // This is mainly to test the Adam7 interlacing
+  for(int h = 1; h < 12; h++)
+  for(int w = 1; w < 12; w++)
+  {
+    testSize(w, h);
+  }
 }
 
 //Tests some specific color conversions with specific color bit combinations