diff --git a/src/ImageSharp/Formats/Jxl/Processing/Decoder/JxlContextMapDecoder.cs b/src/ImageSharp/Formats/Jxl/Processing/Decoder/JxlContextMapDecoder.cs index 78b989910a..9695384b48 100644 --- a/src/ImageSharp/Formats/Jxl/Processing/Decoder/JxlContextMapDecoder.cs +++ b/src/ImageSharp/Formats/Jxl/Processing/Decoder/JxlContextMapDecoder.cs @@ -1,8 +1,109 @@ -using System; +// Copyright (c) Six Labors. +// Licensed under the Six Labors Split License. + +using System.Numerics.Tensors; +using System.Runtime.InteropServices; +using SixLabors.ImageSharp.Formats.Jxl.Processing.Decoder.Ans; +using SixLabors.ImageSharp.Formats.Jxl.Processing.Primitives; namespace SixLabors.ImageSharp.Formats.Jxl.Processing.Decoder; -public class JxlContextMapDecoder +internal static class JxlContextMapDecoder { + private const int MaxClusters = 256; + + public static void VerifyContextMap(List contextMap, int numHTrees) + { + Span haveHTree = stackalloc bool[numHTrees]; + + int numFound = 0; + + foreach (byte htree in contextMap) + { + if (htree >= numHTrees) + { + throw new InvalidOperationException("Invalid histogram index in context map"); + } + + if (!haveHTree[htree]) + { + haveHTree[htree] = true; + numFound++; + } + } + + if (numFound != numHTrees) + { + throw new InvalidOperationException("Incomplete context map"); + } + } + + public static void DecodeContextMap(Configuration configuration, List contextMap, ref int numHTrees, JxlBitReader input) + { + // True - simple context map, without any ANS entropy coding, + // false - the context map uses entropy coding and LZ77. + bool isSimpleContextMap = input.ReadBoolean(); + + Span cmSpan = CollectionsMarshal.AsSpan(contextMap); + + if (isSimpleContextMap) + { + uint bitsPerEntry = input.ReadBits32(2); + + if (bitsPerEntry != 0) + { + for (int i = 0; i < cmSpan.Length; i++) + { + cmSpan[i] = (byte)input.ReadBits32(bitsPerEntry); + } + } + else + { + // No context map at all + cmSpan.Clear(); + } + } + else + { + bool useMtf = input.ReadBoolean(); // use Move to Front transform? + JxlAnsCode ansCode = new(); + List sinkContextMap = []; + + if (!JxlAnsReader.DecodeHistograms(configuration, input, 1, ansCode, sinkContextMap, contextMap.Count <= 2)) + { + throw new InvalidOperationException("Context map histogram couldn't be decoded"); + } + + JxlAnsSymbolReader ans = JxlAnsSymbolReader.Create(ansCode, input); + + int i = 0; + int maxSymbols = 0; + + while (i < cmSpan.Length) + { + uint sym = ans.ReadHybridUInt(true, 0, input, sinkContextMap); + maxSymbols = sym > maxSymbols ? (int)sym : maxSymbols; + cmSpan[i] = (byte)sym; + i++; + } + + if (maxSymbols >= MaxClusters) + { + throw new InvalidOperationException("Invalid cluster ID"); + } + + if (!ans.CheckAnsFinalState()) + { + throw new InvalidOperationException("Invalid context map"); + } + + if (useMtf) + { + JxlInverseMtf.InverseMoveToFrontTransform(cmSpan); + } + } + numHTrees = TensorPrimitives.Max((ReadOnlySpan)cmSpan) + 1; + VerifyContextMap(contextMap, numHTrees); + } }