1 files changed, 19 insertions, 16 deletions
diff --git a/lib/jxl/enc_quant_weights.cc b/lib/jxl/enc_quant_weights.cc
index d8a9931..0d3f628 100644
--- a/lib/jxl/enc_quant_weights.cc
+++ b/lib/jxl/enc_quant_weights.cc
@@ -5,7 +5,6 @@
 
 #include "lib/jxl/enc_quant_weights.h"
 
-#include <stdio.h>
 #include <stdlib.h>
 
 #include <algorithm>
@@ -13,12 +12,11 @@
 #include <limits>
 #include <utility>
 
-#include "lib/jxl/aux_out.h"
-#include "lib/jxl/aux_out_fwd.h"
 #include "lib/jxl/base/bits.h"
+#include "lib/jxl/base/common.h"
 #include "lib/jxl/base/status.h"
-#include "lib/jxl/common.h"
 #include "lib/jxl/dct_scales.h"
+#include "lib/jxl/enc_aux_out.h"
 #include "lib/jxl/enc_bit_writer.h"
 #include "lib/jxl/enc_modular.h"
 #include "lib/jxl/fields.h"
@@ -28,6 +26,8 @@
 
 namespace jxl {
 
+struct AuxOut;
+
 namespace {
 
 Status EncodeDctParams(const DctQuantWeightParams& params, BitWriter* writer) {
@@ -138,7 +138,7 @@ Status DequantMatricesEncode(const DequantMatrices* matrices, BitWriter* writer,
           DequantMatrices::required_size_y[i], writer, modular_frame_encoder));
     }
   }
-  ReclaimAndCharge(writer, &allotment, layer, aux_out);
+  allotment.ReclaimAndCharge(writer, layer, aux_out);
   return true;
 }
 
@@ -159,7 +159,7 @@ Status DequantMatricesEncodeDC(const DequantMatrices* matrices,
       JXL_RETURN_IF_ERROR(F16Coder::Write(dc_quant[c] * 128.0f, writer));
     }
   }
-  ReclaimAndCharge(writer, &allotment, layer, aux_out);
+  allotment.ReclaimAndCharge(writer, layer, aux_out);
   return true;
 }
 
@@ -183,6 +183,18 @@ void DequantMatricesScaleDC(DequantMatrices* matrices, const float scale) {
   DequantMatricesSetCustomDC(matrices, dc);
 }
 
+void DequantMatricesRoundtrip(DequantMatrices* matrices) {
+  // Do not pass modular en/decoder, as they only change entropy and not
+  // values.
+  BitWriter writer;
+  JXL_CHECK(DequantMatricesEncode(matrices, &writer, 0, nullptr));
+  writer.ZeroPadToByte();
+  BitReader br(writer.GetSpan());
+  // Called only in the encoder: should fail only for programmer errors.
+  JXL_CHECK(matrices->Decode(&br));
+  JXL_CHECK(br.Close());
+}
+
 void DequantMatricesSetCustom(DequantMatrices* matrices,
                               const std::vector<QuantEncoding>& encodings,
                               ModularFrameEncoder* encoder) {
@@ -195,16 +207,7 @@ void DequantMatricesSetCustom(DequantMatrices* matrices,
                              encodings[i], i);
     }
   }
-  // Roundtrip encode/decode the matrices to ensure same values as decoder.
-  // Do not pass modular en/decoder, as they only change entropy and not
-  // values.
-  BitWriter writer;
-  JXL_CHECK(DequantMatricesEncode(matrices, &writer, 0, nullptr));
-  writer.ZeroPadToByte();
-  BitReader br(writer.GetSpan());
-  // Called only in the encoder: should fail only for programmer errors.
-  JXL_CHECK(matrices->Decode(&br));
-  JXL_CHECK(br.Close());
+  DequantMatricesRoundtrip(matrices);
 }
 
 }  // namespace jxl