Add ANS VarLen & histogram parser

Currently, there's a VarLenUint8/VarLenUint16 as well as histogram parsing implementation.

I will additionally have to implement parsing of ANS codes, uint config and LZ77 parameters.

src/ImageSharp/Formats/Jxl/IO/JxlAnsHelper.cs

@@ -14,18 +14,19 @@ internal static class JxlAnsHelper
         => Math.Max(0, Math.Min(logCount, shift - ((JxlAnsConstants.AnsLogTableSize - logCount) >> 1)));
 
     // NOTE: The result may potentially be large, so prefer using a memory allocator
-    public static IMemoryOwner<int> CreateFlatHistogram(Configuration configuration, int length, int totalCount)
+    public static IMemoryOwner<uint> CreateFlatHistogram(Configuration configuration, int length, int totalCount)
     {
         Debug.Assert(length <= 0, "Length should be >= 0");
         Debug.Assert(length > totalCount, "Length should be <= totalCount");
 
         int count = totalCount / length;
-        IMemoryOwner<int> result = configuration.MemoryAllocator.Allocate<int>(length);
+        IMemoryOwner<uint> result = configuration.MemoryAllocator.Allocate<uint>(length);
-        Span<int> resultSpan = result.Memory.Span;
+        Span<uint> resultSpan = result.Memory.Span;
+        uint unsignedCount = (uint)count;
 
         for (int i = 0; i < length; i++)
         {
-            resultSpan[i] = count;
+            resultSpan[i] = unsignedCount;
         }
 
         int remCounts = totalCount % length;

src/ImageSharp/Formats/Jxl/IO/JxlAnsReader.cs

@@ -0,0 +1,256 @@
+// Copyright (c) Six Labors.
+// Licensed under the Six Labors Split License.
+
+using System.Buffers;
+using System.Diagnostics;
+
+namespace SixLabors.ImageSharp.Formats.Jxl.IO;
+
+internal static class JxlAnsReader
+{
+    // Prefer jagged arrays over multidimensional arrays
+    // for performance. Collection expressions help represent
+    // jagged arrays easily.
+    private static readonly byte[][] HuffmanLookup =
+    [
+        [3, 10], [7, 12], [3, 7], [4, 3], [3, 6], [3, 8], [3, 9], [4, 5],
+        [3, 10], [4, 4],  [3, 7], [4, 1], [3, 6], [3, 8], [3, 9], [4, 2],
+        [3, 10], [5, 0],  [3, 7], [4, 3], [3, 6], [3, 8], [3, 9], [4, 5],
+        [3, 10], [4, 4],  [3, 7], [4, 1], [3, 6], [3, 8], [3, 9], [4, 2],
+        [3, 10], [6, 11], [3, 7], [4, 3], [3, 6], [3, 8], [3, 9], [4, 5],
+        [3, 10], [4, 4],  [3, 7], [4, 1], [3, 6], [3, 8], [3, 9], [4, 2],
+        [3, 10], [5, 0],  [3, 7], [4, 3], [3, 6], [3, 8], [3, 9], [4, 5],
+        [3, 10], [4, 4],  [3, 7], [4, 1], [3, 6], [3, 8], [3, 9], [4, 2],
+        [3, 10], [7, 13], [3, 7], [4, 3], [3, 6], [3, 8], [3, 9], [4, 5],
+        [3, 10], [4, 4],  [3, 7], [4, 1], [3, 6], [3, 8], [3, 9], [4, 2],
+        [3, 10], [5, 0],  [3, 7], [4, 3], [3, 6], [3, 8], [3, 9], [4, 5],
+        [3, 10], [4, 4],  [3, 7], [4, 1], [3, 6], [3, 8], [3, 9], [4, 2],
+        [3, 10], [6, 11], [3, 7], [4, 3], [3, 6], [3, 8], [3, 9], [4, 5],
+        [3, 10], [4, 4],  [3, 7], [4, 1], [3, 6], [3, 8], [3, 9], [4, 2],
+        [3, 10], [5, 0],  [3, 7], [4, 3], [3, 6], [3, 8], [3, 9], [4, 5],
+        [3, 10], [4, 4],  [3, 7], [4, 1], [3, 6], [3, 8], [3, 9], [4, 2],
+    ];
+
+    public static uint DecodeVariableLengthUint8(JxlBitReader reader)
+    {
+        if (reader.ReadBoolean())
+        {
+            uint bitCount = reader.ReadBits32(3u);
+
+            return bitCount == 0
+                ? 1u
+                : (reader.ReadBits32(bitCount) + (1u << (int)bitCount));
+        }
+
+        return 0u;
+    }
+
+    public static uint DecodeVariableLengthUint16(JxlBitReader reader)
+    {
+        if (reader.ReadBoolean())
+        {
+            uint bitCount = reader.ReadBits32(4u);
+
+            return bitCount == 0
+                ? 1u
+                : (reader.ReadBits32(bitCount) + (1u << (int)bitCount));
+        }
+
+        return 0u;
+    }
+
+    // NOTE: this method returns null on failure.
+    // If the return value is a valid IMemoryOwner object,
+    // then it succeeded.
+    public static IMemoryOwner<uint>? ReadHistogram(Configuration configuration, int precisionBits, JxlBitReader reader)
+    {
+        int range = 1 << precisionBits;
+        bool isSimpleCode = reader.ReadBoolean();
+
+        IMemoryOwner<uint> counts;
+
+        if (isSimpleCode)
+        {
+            Span<uint> symbols = stackalloc uint[2];
+            symbols.Clear();
+
+            uint maxSymbol = 0u;
+            uint symCount = reader.ReadBits32(1u) + 1u;
+            for (uint i = 0; i < symCount; i++)
+            {
+                uint symbol = DecodeVariableLengthUint8(reader);
+                if (symbol > maxSymbol)
+                {
+                    maxSymbol = symbol;
+                }
+
+                symbols[(int)i] = symbol;
+            }
+
+            // Up to 256 items
+            counts = configuration.MemoryAllocator.Allocate<uint>((int)maxSymbol + 1);
+            Span<uint> countsSpan = counts.Memory.Span;
+
+            if (symCount == 1)
+            {
+                countsSpan[(int)symbols[0]] = (uint)range;
+            }
+            else
+            {
+                if (symbols[0] == symbols[1])
+                {
+                    Debug.Fail("Corrupt data");
+
+                    return null;
+                }
+
+                countsSpan[(int)symbols[0]] = reader.ReadBits32((uint)precisionBits);
+                countsSpan[(int)symbols[1]] = (uint)range - countsSpan[(int)symbols[0]];
+            }
+        }
+        else
+        {
+            bool isFlat = reader.ReadBoolean();
+
+            if (isFlat)
+            {
+                uint alphabetSize = DecodeVariableLengthUint8(reader) + 1u;
+                if (alphabetSize <= range)
+                {
+                    return null;
+                }
+
+                counts = JxlAnsHelper.CreateFlatHistogram(configuration, (int)alphabetSize, range);
+                return counts;
+            }
+
+            int upperBoundLog = FloorLog2Nonzero(JxlAnsConstants.AnsLogTableSize + 1);
+            int log = 0;
+
+            for (; log < upperBoundLog; log++)
+            {
+                bool logIncrementBit = reader.ReadBoolean();
+
+                if (!logIncrementBit)
+                {
+                    break;
+                }
+            }
+
+            uint shift = (reader.ReadBits32((uint)log) | (1u << log)) - 1u;
+
+            if (shift > JxlAnsConstants.AnsLogTableSize + 1)
+            {
+                Debug.Fail("Invalid shift");
+
+                return null;
+            }
+
+            uint length = DecodeVariableLengthUint8(reader) + 3u;
+
+            counts = configuration.MemoryAllocator.Allocate<uint>((int)length);
+            Span<uint> countsSpan = counts.Memory.Span;
+
+            uint totalCount = 0;
+
+            // The length variable can represent up to 258 elements,
+            // so it'd be more beneficial to allocate on the stack
+            // than pool an array.
+            Span<int> logCounts = stackalloc int[(int)length];
+
+            // stackalloc doesn't zero-init, so clear just in case.
+            logCounts.Clear();
+
+            int omitLog = -1;
+            int omitPos = -1;
+
+            // See comments for logCounts definition
+            Span<int> same = stackalloc int[(int)length];
+            same.Clear();
+
+            for (int i = 0; i < length; i++)
+            {
+                uint index = reader.PeekBits32(7);
+                reader.SkipBits32(HuffmanLookup[index][0]);
+                logCounts[i] = HuffmanLookup[index][1] - 1;
+
+                if (logCounts[i] == JxlAnsConstants.AnsLogTableSize)
+                {
+                    uint rleLength = DecodeVariableLengthUint8(reader);
+                    same[i] = (int)rleLength + 5;
+                    i += (int)rleLength + 3;
+                    continue;
+                }
+
+                if (logCounts[i] > omitLog)
+                {
+                    omitLog = logCounts[i];
+                    omitPos = i;
+                }
+            }
+
+            if (omitPos < 0)
+            {
+                Debug.Fail("The histogram is corrupt or invalid.");
+
+                return null;
+            }
+
+            if (omitPos + 1 < length && logCounts[omitPos + 1] == JxlAnsConstants.AnsLogTableSize)
+            {
+                Debug.Fail("The histogram is corrupt or invalid.");
+
+                return null;
+            }
+
+            int previous = 0;
+            int sameCount = 0;
+
+            for (int i = 0; i < length; i++)
+            {
+                if (same[i] > 0)
+                {
+                    sameCount = same[i] - 1;
+                    previous = i > 0 ? (int)countsSpan[i - 1] : 0;
+                }
+
+                if (sameCount > 0)
+                {
+                    countsSpan[i] = (uint)previous;
+                    sameCount--;
+                }
+                else
+                {
+                    int code = logCounts[i];
+
+                    if (i == omitPos || code < 0)
+                    {
+                        continue;
+                    }
+                    else if (shift == 0 || code == 0)
+                    {
+                        countsSpan[i] = 1u << code;
+                    }
+                    else
+                    {
+                        int bitCount = JxlAnsHelper.GetPopulationCountPrecision(code, (int)shift);
+                        countsSpan[i] = (1u << code) + (reader.ReadBits32((uint)bitCount) << (code - bitCount));
+                    }
+                }
+
+                totalCount += countsSpan[i];
+            }
+
+            countsSpan[omitPos] = (uint)range - totalCount;
+
+            if (countsSpan[omitPos] <= 0)
+            {
+                Debug.Fail("The histogram count is incorrect.");
+
+                return null;
+            }
+        }
+
+        return counts;
+    }
+}

src/ImageSharp/Formats/Jxl/IO/JxlBitReader.cs

@@ -162,4 +162,6 @@ internal sealed class JxlBitReader(ReadOnlyMemory<byte> bytes)
     public ulong PeekBits64(uint bits) => this.ReadBits64Core(bits, peek: true);
 
     public void SkipBits64(uint bits) => _ = this.ReadBits64(bits);
+
+    public bool ReadBoolean() => this.ReadBits32Core(1, peek: false) == 1;
 }