Implement ANS entropy helpers

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

@@ -8,24 +8,9 @@ namespace SixLabors.ImageSharp.Formats.Jxl.IO;
 [StructLayout(LayoutKind.Sequential)]
 internal struct JxlAnsEntry
 {
-    // Although the Entry struct looks like this:
-    //     uint8_t cutoff;
-    //     uint8_t right_value;
-    //     uint16_t freq0;
-    //     uint16_t offsets1;
-    //     uint16_t freq1_xor_freq0;
-    // and clearly uses smaller types (e.g. byte or ushort),
-    // prefer using int here as otherwise we have to
-    // introduce many casts: some to assign values, others to
-    // convert from unsigned to signed kinds.
-    //
-    // This struct is 20 bytes which is more than the recommended
-    // maximum of 16 bytes, but I believe it justifies more due to
-    // reduced heap allocations that would be introduced if this
-    // struct would be turned into a class.
-    public int Entry;
-    public int RightValue;
-    public int Frequency0;
-    public int Offsets1;
-    public int Frequency1XorFrequency0;
+    public byte Cutoff;
+    public byte RightValue;
+    public ushort Frequency0;
+    public ushort Offsets1;
+    public ushort Frequency1XorFrequency0;
 }

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

@@ -46,4 +46,206 @@ internal static class JxlAnsHelper
 
         return result;
     }
+
+    public static JxlAnsSymbol Lookup(ReadOnlySpan<JxlAnsEntry> table, int value, int logEntrySize, int entrySizeMinus1)
+    {
+        int i = value >> logEntrySize;
+        int pos = value & entrySizeMinus1;
+
+        JxlAnsEntry entry = table[i];
+
+        int cutoff = entry.Cutoff;
+        int rightValue = entry.RightValue;
+        int freq0 = entry.Frequency0;
+
+        bool greater = pos >= cutoff;
+
+        int offsets1or0 = greater ? entry.Offsets1 : 0;
+        int freq1xorfreq0or0 = greater ? entry.Frequency1XorFrequency0 : 0;
+
+        JxlAnsSymbol symbol = new()
+        {
+            Value = greater ? rightValue : i,
+            Offset = offsets1or0 + pos,
+            Frequency = freq0 ^ freq1xorfreq0or0
+        };
+
+        return symbol;
+    }
+
+    public static bool InitAliasTable(Span<int> preDistribution, uint logRange, int logAlphaSize, Span<JxlAnsEntry> entries)
+    {
+        int range = 1 << (int)logRange;
+        int tableSize = 1 << logAlphaSize;
+
+        Debug.Assert(tableSize <= range, "table_size must be <= range");
+
+        int distributionPointer = preDistribution.Length - 1;
+
+        while (distributionPointer >= 0 && preDistribution[distributionPointer] == 0)
+        {
+            distributionPointer--;
+        }
+
+        if (distributionPointer < 0)
+        {
+            preDistribution[0] = range;
+            distributionPointer = 0;
+        }
+
+        Span<int> distribution = preDistribution[..(distributionPointer + 1)];
+
+        if (distribution.Length > tableSize)
+        {
+            Debug.Fail("Too many items in the distribution");
+
+            return false;
+        }
+
+        int entrySize = range >> logAlphaSize;
+        int singleSymbol = -1;
+        int sum = 0;
+
+        for (int sym = 0; sym < distribution.Length; sym++)
+        {
+            int value = distribution[sym];
+            sum += value;
+
+            if (value == AnsTableSize)
+            {
+                if (singleSymbol != -1)
+                {
+                    return false;
+                }
+
+                singleSymbol = sym;
+            }
+        }
+
+        if (sum != range)
+        {
+            return false;
+        }
+
+        if (singleSymbol != -1)
+        {
+            byte sym = (byte)singleSymbol;
+            if (singleSymbol != sym)
+            {
+                return false;
+            }
+
+            for (int i = 0; i < tableSize; i++)
+            {
+                ref JxlAnsEntry jxlEntry = ref entries[i];
+
+                jxlEntry.RightValue = sym;
+                jxlEntry.Cutoff = 0;
+                jxlEntry.Offsets1 = (ushort)(entrySize * i);
+                jxlEntry.Frequency0 = 0;
+                jxlEntry.Frequency1XorFrequency0 = AnsTableSize;
+            }
+
+            return true;
+        }
+
+        Span<uint> underfullPosn = stackalloc uint[distribution.Length];
+        Span<uint> overfullPosn = stackalloc uint[distribution.Length];
+        Span<uint> cutoffs = stackalloc uint[1 << logAlphaSize];
+
+        int underfullPointer = 0;
+        int overfullPointer = 0;
+
+        for (int i = 0; i < distribution.Length; i++)
+        {
+            uint currentCutoff = (uint)distribution[i];
+
+            cutoffs[i] = currentCutoff;
+
+            if (currentCutoff > entrySize)
+            {
+                overfullPosn[overfullPointer] = (uint)i;
+                overfullPointer++;
+            }
+            else if (currentCutoff < entrySize)
+            {
+                underfullPosn[underfullPointer] = (uint)i;
+                underfullPointer++;
+            }
+        }
+
+        for (int i = distribution.Length; i < tableSize; i++)
+        {
+            cutoffs[i] = 0;
+            underfullPosn[underfullPointer] = (uint)i;
+            underfullPointer++;
+        }
+
+        uint unsignedEntrySize = (uint)entrySize;
+
+        while (overfullPointer >= 0)
+        {
+            uint overfullIndex = overfullPosn[overfullPointer];
+            overfullPointer--;
+
+            if (underfullPointer <= -1)
+            {
+                return false;
+            }
+
+            uint underfullIndex = underfullPosn[underfullPointer];
+            underfullPointer--;
+
+            int signedOverfullIndex = (int)overfullIndex;
+            int signedUnderfullIndex = (int)underfullIndex;
+
+            uint underfullBy = unsignedEntrySize - cutoffs[signedUnderfullIndex];
+            cutoffs[signedOverfullIndex] -= underfullBy;
+
+            ref JxlAnsEntry currentEntry = ref entries[signedUnderfullIndex];
+
+            currentEntry.RightValue = unchecked((byte)overfullIndex);
+            currentEntry.Offsets1 = unchecked((ushort)cutoffs[signedOverfullIndex]);
+
+            uint currentCutoff = cutoffs[signedOverfullIndex];
+
+            if (currentCutoff < entrySize)
+            {
+                underfullPosn[underfullPointer] = overfullIndex;
+                underfullPointer++;
+            }
+            else if (currentCutoff > entrySize)
+            {
+                overfullPosn[overfullPointer] = overfullIndex;
+                overfullPointer++;
+            }
+        }
+
+        for (uint i = 0; i < tableSize; i++)
+        {
+            uint currentCutoff = cutoffs[(int)i];
+            ref JxlAnsEntry entry = ref entries[(int)i];
+
+            if (currentCutoff == entrySize)
+            {
+                entry.RightValue = (byte)i;
+                entry.Offsets1 = 0;
+                entry.Cutoff = 0;
+            }
+            else
+            {
+                entry.Offsets1 -= (ushort)currentCutoff;
+                entry.Cutoff = (byte)currentCutoff;
+            }
+
+            int freq0 = i < distribution.Length ? distribution[(int)i] : 0;
+            int i1 = entry.RightValue;
+            int freq1 = i1 < distribution.Length ? distribution[i1] : 0;
+
+            entry.Frequency0 = (ushort)freq0;
+            entry.Frequency1XorFrequency0 = (ushort)(freq1 ^ freq0);
+        }
+
+        return true;
+    }
 }