diff --git a/src/ImageSharp/Formats/Jxl/IO/JxlAnsHelper.cs b/src/ImageSharp/Formats/Jxl/IO/JxlAnsHelper.cs index e8b3080ac3..35630eb103 100644 --- a/src/ImageSharp/Formats/Jxl/IO/JxlAnsHelper.cs +++ b/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 CreateFlatHistogram(Configuration configuration, int length, int totalCount) + public static IMemoryOwner 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 result = configuration.MemoryAllocator.Allocate(length); - Span resultSpan = result.Memory.Span; + IMemoryOwner result = configuration.MemoryAllocator.Allocate(length); + Span 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; diff --git a/src/ImageSharp/Formats/Jxl/IO/JxlAnsReader.cs b/src/ImageSharp/Formats/Jxl/IO/JxlAnsReader.cs new file mode 100644 index 0000000000..8fd8bd7c8c --- /dev/null +++ b/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? ReadHistogram(Configuration configuration, int precisionBits, JxlBitReader reader) + { + int range = 1 << precisionBits; + bool isSimpleCode = reader.ReadBoolean(); + + IMemoryOwner counts; + + if (isSimpleCode) + { + Span 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((int)maxSymbol + 1); + Span 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((int)length); + Span 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 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 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; + } +} diff --git a/src/ImageSharp/Formats/Jxl/IO/JxlBitReader.cs b/src/ImageSharp/Formats/Jxl/IO/JxlBitReader.cs index 55e3ed3efd..99b1599915 100644 --- a/src/ImageSharp/Formats/Jxl/IO/JxlBitReader.cs +++ b/src/ImageSharp/Formats/Jxl/IO/JxlBitReader.cs @@ -162,4 +162,6 @@ internal sealed class JxlBitReader(ReadOnlyMemory 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; }