From a106793e970d6da38b9361303ac1baaeb61b16eb Mon Sep 17 00:00:00 2001
From: Luca Versari <veluca93@gmail.com>
Date: Sun, 3 Mar 2024 21:22:34 +0100
Subject: [PATCH] Add options to add custom chunks.

---
 fpnge.cc | 25 +++++++++++++++++++++++--
 fpnge.h  | 50 +++++++++++++++++++++++++++++++-------------------
 2 files changed, 54 insertions(+), 21 deletions(-)

diff --git a/fpnge.cc b/fpnge.cc
index be74386..200d0a6 100644
--- a/fpnge.cc
+++ b/fpnge.cc
@@ -373,6 +373,11 @@ struct BitWriter {
     }
   }
 
+  void WriteBytes(const char *data, size_t count) {
+    memcpy(this->data + bytes_written, data, count);
+    bytes_written += count;
+  }
+
   unsigned char *data;
   size_t bytes_written = 0;
   size_t bits_in_buffer = 0;
@@ -1407,6 +1412,8 @@ static void AppendBE32(size_t value, BitWriter *__restrict writer) {
 
 static void WriteHeader(size_t width, size_t height, size_t bytes_per_channel,
                         size_t num_channels, char cicp_colorspace,
+                        const FPNGEAdditionalChunk *additional_chunks,
+                        int num_additional_chunks,
                         BitWriter *__restrict writer) {
   constexpr uint8_t kPNGHeader[8] = {137, 80, 78, 71, 13, 10, 26, 10};
   for (size_t i = 0; i < 8; i++) {
@@ -1439,9 +1446,22 @@ static void WriteHeader(size_t width, size_t height, size_t bytes_per_channel,
     writer->Write(32, 0x01001009); // PQ, Rec2020
     writer->Write(32, 0xfe23234d); // CRC
   }
+  for (int i = 0; i < num_additional_chunks; i++) {
+    AppendBE32(additional_chunks[i].data_size, writer);
+    size_t crc_start = writer->bytes_written;
+    uint32_t name = 0;
+    memcpy(&name, additional_chunks[i].name, 4);
+    writer->Write(32, name);
+    writer->WriteBytes(additional_chunks[i].data,
+                       additional_chunks[i].data_size);
+    size_t crc_end = writer->bytes_written;
+    uint32_t crc =
+        Crc32().update_final(writer->data + crc_start, crc_end - crc_start);
+    AppendBE32(crc, writer);
+  }
 }
 
-}  // namespace
+} // namespace
 
 extern "C" size_t FPNGEEncode(size_t bytes_per_channel, size_t num_channels,
                               const void *data, size_t width, size_t row_stride,
@@ -1488,7 +1508,8 @@ extern "C" size_t FPNGEEncode(size_t bytes_per_channel, size_t num_channels,
   writer.data = static_cast<unsigned char *>(output);
 
   WriteHeader(width, height, bytes_per_channel, num_channels,
-              options->cicp_colorspace, &writer);
+              options->cicp_colorspace, options->additional_chunks,
+              options->num_additional_chunks, &writer);
 
   assert(writer.bits_in_buffer == 0);
   size_t chunk_length_pos = writer.bytes_written;
diff --git a/fpnge.h b/fpnge.h
index a48aaf1..fe94a4c 100644
--- a/fpnge.h
+++ b/fpnge.h
@@ -30,35 +30,47 @@ enum FPNGEOptionsPredictor {
   FPNGE_PREDICTOR_APPROX,
   FPNGE_PREDICTOR_BEST
 };
+
+struct FPNGEAdditionalChunk {
+  char name[4];
+  const char *data;
+  int data_size;
+};
+
 struct FPNGEOptions {
-  char predictor;        // FPNGEOptionsPredictor
-  char huffman_sample;   // 0-127: how much of the image to sample
-  char cicp_colorspace;  // FPNGECicpColorspace
+  char predictor;       // FPNGEOptionsPredictor
+  char huffman_sample;  // 0-127: how much of the image to sample
+  char cicp_colorspace; // FPNGECicpColorspace
+  int num_additional_chunks;
+  const struct FPNGEAdditionalChunk *additional_chunks;
 };
 
 #define FPNGE_COMPRESS_LEVEL_DEFAULT 4
 #define FPNGE_COMPRESS_LEVEL_BEST 5
 inline void FPNGEFillOptions(struct FPNGEOptions *options, int level,
                              int cicp_colorspace) {
-  if (level == 0) level = FPNGE_COMPRESS_LEVEL_DEFAULT;
+  if (level == 0)
+    level = FPNGE_COMPRESS_LEVEL_DEFAULT;
   options->cicp_colorspace = cicp_colorspace;
   options->huffman_sample = 1;
+  options->num_additional_chunks = 0;
+  options->additional_chunks = NULL;
   switch (level) {
-    case 1:
-      options->predictor = 2;
-      break;
-    case 2:
-      options->predictor = 4;
-      break;
-    case 3:
-      options->predictor = 5;
-      break;
-    case 5:
-      options->huffman_sample = 23;
-      // fall through
-    default:
-      options->predictor = 6;
-      break;
+  case 1:
+    options->predictor = 2;
+    break;
+  case 2:
+    options->predictor = 4;
+    break;
+  case 3:
+    options->predictor = 5;
+    break;
+  case 5:
+    options->huffman_sample = 23;
+    // fall through
+  default:
+    options->predictor = 6;
+    break;
   }
 }