Test case for coefficient round trip

src/ImageSharp/Formats/Heif/Av1/Entropy/Av1SymbolDecoder.cs

@@ -206,6 +206,88 @@ internal ref struct Av1SymbolDecoder
         return r.ReadSymbol(this.chromeForLumaAlpha[context]);
     }
 
+    /// <summary>
+    /// 5.11.39. Coefficients syntax.
+    /// </summary>
+    /// <remarks>
+    /// The implementation is taken from SVT-AV1 library, which deviates from the code flow in the specification.
+    /// </remarks>
+    public int ReadCoefficients(
+        Av1BlockModeInfo modeInfo,
+        Point blockPosition,
+        int[] aboveContexts,
+        int[] leftContexts,
+        int aboveOffset,
+        int leftOffset,
+        int plane,
+        int blocksWide,
+        int blocksHigh,
+        Av1TransformBlockContext transformBlockContext,
+        Av1TransformSize transformSize,
+        bool isLossless,
+        bool useReducedTransformSet,
+        Av1TransformInfo transformInfo,
+        int modeBlocksToRightEdge,
+        int modeBlocksToBottomEdge,
+        Span<int> coefficientBuffer)
+    {
+        int width = transformSize.GetWidth();
+        int height = transformSize.GetHeight();
+        Av1TransformSize transformSizeContext = (Av1TransformSize)(((int)transformSize.GetSquareSize() + ((int)transformSize.GetSquareUpSize() + 1)) >> 1);
+        Av1PlaneType planeType = (Av1PlaneType)Math.Min(plane, 1);
+        int culLevel = 0;
+
+        byte[] levelsBuffer = new byte[Av1Constants.TransformPad2d];
+        Span<byte> levels = levelsBuffer.AsSpan()[(Av1Constants.TransformPadTop * (width + Av1Constants.TransformPadHorizontal))..];
+
+        bool allZero = this.ReadTransformBlockSkip(transformSizeContext, transformBlockContext.SkipContext);
+        int bwl = transformSize.GetBlockWidthLog2();
+        int endOfBlock;
+        if (allZero)
+        {
+            if (plane == 0)
+            {
+                transformInfo.Type = Av1TransformType.DctDct;
+                transformInfo.CodeBlockFlag = false;
+            }
+
+            this.UpdateCoefficientContext(modeInfo, aboveContexts, leftContexts, blocksWide, blocksHigh, transformSize, blockPosition, aboveOffset, leftOffset, culLevel, modeBlocksToRightEdge, modeBlocksToBottomEdge);
+            return 0;
+        }
+
+        transformInfo.Type = ComputeTransformType(planeType, modeInfo, isLossless, transformSize, transformInfo, useReducedTransformSet);
+        Av1TransformClass transformClass = transformInfo.Type.ToClass();
+        Av1ScanOrder scanOrder = Av1ScanOrderConstants.GetScanOrder(transformSize, transformInfo.Type);
+        ReadOnlySpan<short> scan = scanOrder.Scan;
+
+        endOfBlock = this.ReadEndOfBlockPosition(transformSize, transformClass, transformSizeContext, planeType);
+        if (endOfBlock > 1)
+        {
+            Array.Fill(levelsBuffer, (byte)0, 0, ((width + Av1Constants.TransformPadHorizontal) * (height + Av1Constants.TransformPadVertical)) + Av1Constants.TransformPadEnd);
+        }
+
+        this.ReadCoefficientsEndOfBlock(transformClass, endOfBlock, height, scan, bwl, levels, transformSizeContext, planeType);
+        if (endOfBlock > 1)
+        {
+            if (transformClass == Av1TransformClass.Class2D)
+            {
+                this.ReadCoefficientsReverse2d(transformSize, 1, endOfBlock - 1 - 1, scan, bwl, levels, transformSizeContext, planeType);
+                this.ReadCoefficientsReverse(transformSize, transformClass, 0, 0, scan, bwl, levels, transformSizeContext, planeType);
+            }
+            else
+            {
+                this.ReadCoefficientsReverse(transformSize, transformClass, 0, endOfBlock - 1 - 1, scan, bwl, levels, transformSizeContext, planeType);
+            }
+        }
+
+        DebugGuard.MustBeGreaterThan(scan.Length, 0, nameof(scan));
+        culLevel = this.ReadCoefficientsDc(coefficientBuffer, endOfBlock, scan, bwl, levels, transformBlockContext.DcSignContext, planeType);
+        this.UpdateCoefficientContext(modeInfo, aboveContexts, leftContexts, blocksWide, blocksHigh, transformSize, blockPosition, aboveOffset, leftOffset, culLevel, modeBlocksToRightEdge, modeBlocksToBottomEdge);
+
+        transformInfo.CodeBlockFlag = true;
+        return endOfBlock;
+    }
+
     public int ReadEndOfBlockPosition(Av1TransformSize transformSize, Av1TransformClass transformClass, Av1TransformSize transformSizeContext, Av1PlaneType planeType)
     {
         int endOfBlockExtra = 0;
@@ -581,6 +663,106 @@ internal ref struct Av1SymbolDecoder
         }
     }
 
+
+    private void UpdateCoefficientContext(
+        Av1BlockModeInfo modeInfo,
+        int[] aboveContexts,
+        int[] leftContexts,
+        int blocksWide,
+        int blocksHigh,
+        Av1TransformSize transformSize,
+        Point blockPosition,
+        int aboveOffset,
+        int leftOffset,
+        int culLevel,
+        int modeBlockToRightEdge,
+        int modeBlockToBottomEdge)
+    {
+        int transformSizeWide = transformSize.Get4x4WideCount();
+        int transformSizeHigh = transformSize.Get4x4HighCount();
+
+        if (modeBlockToRightEdge < 0)
+        {
+            int aboveContextCount = Math.Min(transformSizeWide, blocksWide - aboveOffset);
+            Array.Fill(aboveContexts, culLevel, 0, aboveContextCount);
+            Array.Fill(aboveContexts, 0, aboveContextCount, transformSizeWide - aboveContextCount);
+        }
+        else
+        {
+            Array.Fill(aboveContexts, culLevel, 0, transformSizeWide);
+        }
+
+        if (modeBlockToBottomEdge < 0)
+        {
+            int leftContextCount = Math.Min(transformSizeHigh, blocksHigh - leftOffset);
+            Array.Fill(leftContexts, culLevel, 0, leftContextCount);
+            Array.Fill(leftContexts, 0, leftContextCount, transformSizeWide - leftContextCount);
+        }
+        else
+        {
+            Array.Fill(leftContexts, culLevel, 0, transformSizeHigh);
+        }
+    }
+
+    private static Av1TransformType ComputeTransformType(Av1PlaneType planeType, Av1BlockModeInfo modeInfo, bool isLossless, Av1TransformSize transformSize, Av1TransformInfo transformInfo, bool useReducedTransformSet)
+    {
+        Av1TransformType transformType = Av1TransformType.DctDct;
+        if (isLossless || transformSize.GetSquareUpSize() > Av1TransformSize.Size32x32)
+        {
+            transformType = Av1TransformType.DctDct;
+        }
+        else
+        {
+            if (planeType == Av1PlaneType.Y)
+            {
+                transformType = transformInfo.Type;
+            }
+            else
+            {
+                // In intra mode, uv planes don't share the same prediction mode as y
+                // plane, so the tx_type should not be shared
+                transformType = ConvertIntraModeToTransformType(modeInfo, Av1PlaneType.Uv);
+            }
+        }
+
+        Av1TransformSetType transformSetType = GetExtendedTransformSetType(transformSize, useReducedTransformSet);
+        if (!transformType.IsExtendedSetUsed(transformSetType))
+        {
+            transformType = Av1TransformType.DctDct;
+        }
+
+        return transformType;
+    }
+
+    private static Av1TransformSetType GetExtendedTransformSetType(Av1TransformSize transformSize, bool useReducedSet)
+    {
+        Av1TransformSize squareUpSize = transformSize.GetSquareUpSize();
+
+        if (squareUpSize >= Av1TransformSize.Size32x32)
+        {
+            return Av1TransformSetType.DctOnly;
+        }
+
+        if (useReducedSet)
+        {
+            return Av1TransformSetType.Dtt4Identity;
+        }
+
+        Av1TransformSize squareSize = transformSize.GetSquareSize();
+        return squareSize == Av1TransformSize.Size16x16 ? Av1TransformSetType.Dtt4Identity : Av1TransformSetType.Dtt4Identity1dDct;
+    }
+
+    private static Av1TransformType ConvertIntraModeToTransformType(Av1BlockModeInfo modeInfo, Av1PlaneType planeType)
+    {
+        Av1PredictionMode mode = (planeType == Av1PlaneType.Y) ? modeInfo.YMode : modeInfo.UvMode;
+        if (mode == Av1PredictionMode.UvChromaFromLuma)
+        {
+            mode = Av1PredictionMode.DC;
+        }
+
+        return mode.ToTransformType();
+    }
+
     internal static Av1Distribution GetSplitOrHorizontalDistribution(Av1Distribution[] inputs, Av1BlockSize blockSize, int context)
     {
         Av1Distribution input = inputs[context];

src/ImageSharp/Formats/Heif/Av1/Entropy/Av1SymbolEncoder.cs

@@ -415,13 +415,13 @@ internal class Av1SymbolEncoder : IDisposable
         ref byte ls = ref levels[0];
 
         Unsafe.InitBlock(ref levels[-Av1Constants.TransformPadTop * stride], 0, (uint)(Av1Constants.TransformPadTop * stride));
-        Unsafe.InitBlock(ref levels[stride * height], 0, (uint)(Av1Constants.TransformPadBottom * stride + Av1Constants.TransformPadEnd));
+        Unsafe.InitBlock(ref levels[stride * height], 0, (uint)((Av1Constants.TransformPadBottom * stride) + Av1Constants.TransformPadEnd));
 
-        for (int i = 0; i < height; i++)
+        for (int y = 0; y < height; y++)
         {
-            for (int j = 0; j < width; j++)
+            for (int x = 0; x < width; x++)
             {
-                ls = (byte)Av1Math.Clamp(Math.Abs(coefficientBuffer[i * width + j]), 0, byte.MaxValue);
+                ls = (byte)Av1Math.Clamp(Math.Abs(coefficientBuffer[(y * width) + x]), 0, byte.MaxValue);
                 ls = ref Unsafe.Add(ref ls, 1);
             }
 
@@ -433,7 +433,10 @@ internal class Av1SymbolEncoder : IDisposable
     /// SVT: set_levels from EbCommonUtils.h
     /// </summary>
     private static Span<byte> SetLevels(Span<byte> levelsBuffer, int width)
-        => levelsBuffer.Slice(Av1Constants.TransformPadTop * (width + Av1Constants.TransformPadHorizontal));
+    {
+        int stride = width + Av1Constants.TransformPadHorizontal;
+        return levelsBuffer[(Av1Constants.TransformPadTop * stride)..];
+    }
 
     private void WriteSkip(bool hasEndOfBlock, int context)
     {

src/ImageSharp/Formats/Heif/Av1/Tiling/Av1BlockModeInfoEncoder.cs

@@ -1,30 +1,27 @@
-// Copyright (c) Six Labors.
+// Copyright (c) Six Labors.
 // Licensed under the Six Labors Split License.
 
 using SixLabors.ImageSharp.Formats.Heif.Av1.Prediction;
 
 namespace SixLabors.ImageSharp.Formats.Heif.Av1.Tiling;
 
-internal partial class Av1TileWriter
+internal class Av1BlockModeInfoEncoder
 {
-    internal class Av1BlockModeInfoEncoder
-    {
-        public Av1BlockSize BlockSize { get; }
+    public Av1BlockSize BlockSize { get; }
 
-        public Av1PredictionMode PredictionMode { get; }
+    public Av1PredictionMode PredictionMode { get; }
 
-        public Av1PartitionType PartitionType { get; }
+    public Av1PartitionType PartitionType { get; }
 
-        public Av1PredictionMode UvPredictionMode { get; }
+    public Av1PredictionMode UvPredictionMode { get; }
 
-        public bool Skip { get; } = true;
+    public bool Skip { get; } = true;
 
-        public bool SkipMode { get; } = true;
+    public bool SkipMode { get; } = true;
 
-        public bool UseIntraBlockCopy { get; } = true;
+    public bool UseIntraBlockCopy { get; } = true;
 
-        public int SegmentId { get; }
+    public int SegmentId { get; }
 
-        public int TransformDepth { get; internal set; }
-    }
+    public int TransformDepth { get; internal set; }
 }

src/ImageSharp/Formats/Heif/Av1/Tiling/Av1Common.cs

@@ -5,18 +5,15 @@ using SixLabors.ImageSharp.Formats.Heif.Av1.OpenBitstreamUnit;
 
 namespace SixLabors.ImageSharp.Formats.Heif.Av1.Tiling;
 
-internal partial class Av1TileWriter
+internal class Av1Common
 {
-    internal class Av1Common
-    {
-        public int ModeInfoRowCount { get; internal set; }
+    public int ModeInfoRowCount { get; internal set; }
 
-        public int ModeInfoColumnCount { get; internal set; }
+    public int ModeInfoColumnCount { get; internal set; }
 
-        public int ModeInfoStride { get; internal set; }
+    public int ModeInfoStride { get; internal set; }
 
-        public required ObuFrameSize FrameSize { get; internal set; }
+    public required ObuFrameSize FrameSize { get; internal set; }
 
-        public required ObuTileGroupHeader TilesInfo { get; internal set; }
-    }
+    public required ObuTileGroupHeader TilesInfo { get; internal set; }
 }

src/ImageSharp/Formats/Heif/Av1/Tiling/Av1MacroBlockModeInfo.cs

@@ -3,14 +3,11 @@
 
 namespace SixLabors.ImageSharp.Formats.Heif.Av1.Tiling;
 
-internal partial class Av1TileWriter
+internal class Av1MacroBlockModeInfo
 {
-    internal class Av1MacroBlockModeInfo
-    {
-        public required Av1BlockModeInfoEncoder Block { get; internal set; }
+    public required Av1BlockModeInfoEncoder Block { get; internal set; }
 
-        public required Av1PaletteLumaModeInfo Palette { get; internal set; }
+    public required Av1PaletteLumaModeInfo Palette { get; internal set; }
 
-        public int CdefStrength { get; internal set; }
-    }
+    public int CdefStrength { get; internal set; }
 }

src/ImageSharp/Formats/Heif/Av1/Tiling/Av1NeighborArrayUnit.cs

@@ -11,14 +11,6 @@ internal class Av1NeighborArrayUnit<T>
 {
     public static readonly T InvalidNeighborData = T.MaxValue;
 
-    [Flags]
-    public enum UnitMask
-    {
-        Left = 1,
-        Top = 2,
-        TopLeft = 4,
-    }
-
     private readonly T[] left;
     private readonly T[] top;
     private readonly T[] topLeft;
@@ -30,6 +22,14 @@ internal class Av1NeighborArrayUnit<T>
         this.topLeft = new T[topLeftSize];
     }
 
+    [Flags]
+    public enum UnitMask
+    {
+        Left = 1,
+        Top = 2,
+        TopLeft = 4,
+    }
+
     public Span<T> Left => this.left;
 
     public Span<T> Top => this.top;

src/ImageSharp/Formats/Heif/Av1/Tiling/Av1PaletteLumaModeInfo.cs

@@ -1,11 +1,8 @@
-// Copyright (c) Six Labors.
+// Copyright (c) Six Labors.
 // Licensed under the Six Labors Split License.
 
 namespace SixLabors.ImageSharp.Formats.Heif.Av1.Tiling;
 
-internal partial class Av1TileWriter
+internal class Av1PaletteLumaModeInfo
 {
-    internal class Av1PaletteLumaModeInfo
-    {
-    }
 }

src/ImageSharp/Formats/Heif/Av1/Tiling/Av1PictureControlSet.cs

@@ -3,20 +3,17 @@
 
 namespace SixLabors.ImageSharp.Formats.Heif.Av1.Tiling;
 
-internal partial class Av1TileWriter
+internal class Av1PictureControlSet
 {
-    internal class Av1PictureControlSet
-    {
-        public required Av1NeighborArrayUnit[] luma_dc_sign_level_coeff_na { get; internal set; }
+    public required Av1NeighborArrayUnit<byte>[] luma_dc_sign_level_coeff_na { get; internal set; }
 
-        public required Av1NeighborArrayUnit[] cr_dc_sign_level_coeff_na { get; internal set; }
+    public required Av1NeighborArrayUnit<byte>[] cr_dc_sign_level_coeff_na { get; internal set; }
 
-        public required Av1NeighborArrayUnit[] cb_dc_sign_level_coeff_na { get; internal set; }
+    public required Av1NeighborArrayUnit<byte>[] cb_dc_sign_level_coeff_na { get; internal set; }
 
-        public required Av1NeighborArrayUnit[] txfm_context_array { get; internal set; }
+    public required Av1NeighborArrayUnit<byte>[] txfm_context_array { get; internal set; }
 
-        public required Av1SequenceControlSet Sequence { get; internal set; }
+    public required Av1SequenceControlSet Sequence { get; internal set; }
 
-        public required Av1PictureParentControlSet Parent { get; internal set; }
-    }
+    public required Av1PictureParentControlSet Parent { get; internal set; }
 }

src/ImageSharp/Formats/Heif/Av1/Tiling/Av1PictureParentControlSet.cs

@@ -5,18 +5,17 @@ using SixLabors.ImageSharp.Formats.Heif.Av1.OpenBitstreamUnit;
 
 namespace SixLabors.ImageSharp.Formats.Heif.Av1.Tiling;
 
-internal partial class Av1TileWriter
+internal class Av1PictureParentControlSet
 {
-    internal class Av1PictureParentControlSet
-    {
-        public required Av1Common Common { get; internal set; }
+    public required Av1Common Common { get; internal set; }
 
-        public required ObuFrameHeader FrameHeader { get; internal set; }
+    public required ObuFrameHeader FrameHeader { get; internal set; }
 
-        public required int[] PreviousQIndex { get; internal set; }
+    public required int[] PreviousQIndex { get; internal set; }
 
-        public int PaletteLevel { get; internal set; }
-        public int AlignedWidth { get; internal set; }
-        public int AlignedHeight { get; internal set; }
-    }
+    public int PaletteLevel { get; internal set; }
+
+    public int AlignedWidth { get; internal set; }
+
+    public int AlignedHeight { get; internal set; }
 }

src/ImageSharp/Formats/Heif/Av1/Tiling/Av1SequenceControlSet.cs

@@ -1,14 +1,11 @@
-// Copyright (c) Six Labors.
+// Copyright (c) Six Labors.
 // Licensed under the Six Labors Split License.
 
 using SixLabors.ImageSharp.Formats.Heif.Av1.OpenBitstreamUnit;
 
 namespace SixLabors.ImageSharp.Formats.Heif.Av1.Tiling;
 
-internal partial class Av1TileWriter
+internal class Av1SequenceControlSet
 {
-    internal class Av1SequenceControlSet
-    {
-        public required ObuSequenceHeader SequenceHeader { get; internal set; }
-    }
+    public required ObuSequenceHeader SequenceHeader { get; internal set; }
 }

src/ImageSharp/Formats/Heif/Av1/Tiling/Av1TileReader.cs

@@ -491,22 +491,6 @@ internal class Av1TileReader : IAv1TileReader
         return hasChroma;
     }
 
-    private Av1TransformSize GetSize(int plane, object transformSize) => throw new NotImplementedException();
-
-    /// <summary>
-    /// 5.11.38. Get plane residual size function.
-    /// The GetPlaneResidualSize returns the size of a residual block for the specified plane. (The residual block will always
-    /// have width and height at least equal to 4.)
-    /// </summary>
-    private Av1BlockSize GetPlaneResidualSize(Av1BlockSize sizeChunk, int plane)
-    {
-        bool subsamplingX = this.SequenceHeader.ColorConfig.SubSamplingX;
-        bool subsamplingY = this.SequenceHeader.ColorConfig.SubSamplingY;
-        bool subX = plane > 0 && subsamplingX;
-        bool subY = plane > 0 && subsamplingY;
-        return sizeChunk.GetSubsampled(subX, subY);
-    }
-
     /// <summary>
     /// 5.11.35. Transform block syntax.
     /// </summary>
@@ -563,152 +547,15 @@ internal class Av1TileReader : IAv1TileReader
         int height = transformSize.GetHeight();
         Av1TransformSize transformSizeContext = (Av1TransformSize)(((int)transformSize.GetSquareSize() + ((int)transformSize.GetSquareUpSize() + 1)) >> 1);
         Av1PlaneType planeType = (Av1PlaneType)Math.Min(plane, 1);
-        int culLevel = 0;
-
-        byte[] levelsBuffer = new byte[Av1Constants.TransformPad2d];
-        Span<byte> levels = levelsBuffer.AsSpan()[(Av1Constants.TransformPadTop * (width + Av1Constants.TransformPadHorizontal))..];
-
-        bool allZero = reader.ReadTransformBlockSkip(transformSizeContext, transformBlockContext.SkipContext);
-        int bwl = transformSize.GetBlockWidthLog2();
-        int endOfBlock;
-        if (allZero)
-        {
-            if (plane == 0)
-            {
-                transformInfo.Type = Av1TransformType.DctDct;
-                transformInfo.CodeBlockFlag = false;
-            }
-
-            this.UpdateCoefficientContext(plane, partitionInfo, transformSize, blockRow, blockColumn, aboveOffset, leftOffset, culLevel);
-            return 0;
-        }
-
-        transformInfo.Type = this.ComputeTransformType(planeType, partitionInfo, transformSize, transformInfo);
-        Av1TransformClass transformClass = transformInfo.Type.ToClass();
-        Av1ScanOrder scanOrder = Av1ScanOrderConstants.GetScanOrder(transformSize, transformInfo.Type);
-        ReadOnlySpan<short> scan = scanOrder.Scan;
-
-        endOfBlock = reader.ReadEndOfBlockPosition(transformSize, transformClass, transformSizeContext, planeType);
-        if (endOfBlock > 1)
-        {
-            Array.Fill(levelsBuffer, (byte)0, 0, ((width + Av1Constants.TransformPadHorizontal) * (height + Av1Constants.TransformPadVertical)) + Av1Constants.TransformPadEnd);
-        }
-
-        reader.ReadCoefficientsEndOfBlock(transformClass, endOfBlock, height, scan, bwl, levels, transformSizeContext, planeType);
-        if (endOfBlock > 1)
-        {
-            if (transformClass == Av1TransformClass.Class2D)
-            {
-                reader.ReadCoefficientsReverse2d(transformSize, 1, endOfBlock - 1 - 1, scan, bwl, levels, transformSizeContext, planeType);
-                reader.ReadCoefficientsReverse(transformSize, transformClass, 0, 0, scan, bwl, levels, transformSizeContext, planeType);
-            }
-            else
-            {
-                reader.ReadCoefficientsReverse(transformSize, transformClass, 0, endOfBlock - 1 - 1, scan, bwl, levels, transformSizeContext, planeType);
-            }
-        }
-
-        DebugGuard.MustBeGreaterThan(scan.Length, 0, nameof(scan));
-        culLevel = reader.ReadCoefficientsDc(coefficientBuffer, endOfBlock, scan, bwl, levels, transformBlockContext.DcSignContext, planeType);
-        this.UpdateCoefficientContext(plane, partitionInfo, transformSize, blockRow, blockColumn, aboveOffset, leftOffset, culLevel);
-
-        transformInfo.CodeBlockFlag = true;
-        return endOfBlock;
-    }
-
-    private void UpdateCoefficientContext(int plane, Av1PartitionInfo partitionInfo, Av1TransformSize transformSize, int blockRow, int blockColumn, int aboveOffset, int leftOffset, int culLevel)
-    {
+        Point blockPosition = new(blockColumn, blockRow);
+        bool isLossless = this.FrameHeader.LosslessArray[partitionInfo.ModeInfo.SegmentId];
         bool subX = this.SequenceHeader.ColorConfig.SubSamplingX;
         bool subY = this.SequenceHeader.ColorConfig.SubSamplingY;
-        int[] aboveContexts = this.aboveNeighborContext.GetContext(plane);
-        int[] leftContexts = this.leftNeighborContext.GetContext(plane);
-        int transformSizeWide = transformSize.Get4x4WideCount();
-        int transformSizeHigh = transformSize.Get4x4HighCount();
-
-        if (partitionInfo.ModeBlockToRightEdge < 0)
-        {
-            Av1BlockSize planeBlockSize = partitionInfo.ModeInfo.BlockSize.GetSubsampled(subX, subY);
-            int blocksWide = partitionInfo.GetMaxBlockWide(planeBlockSize, subX);
-            int aboveContextCount = Math.Min(transformSizeWide, blocksWide - aboveOffset);
-            Array.Fill(aboveContexts, culLevel, 0, aboveContextCount);
-            Array.Fill(aboveContexts, 0, aboveContextCount, transformSizeWide - aboveContextCount);
-        }
-        else
-        {
-            Array.Fill(aboveContexts, culLevel, 0, transformSizeWide);
-        }
-
-        if (partitionInfo.ModeBlockToBottomEdge < 0)
-        {
-            Av1BlockSize planeBlockSize = partitionInfo.ModeInfo.BlockSize.GetSubsampled(subX, subY);
-            int blocksHigh = partitionInfo.GetMaxBlockHigh(planeBlockSize, subY);
-            int leftContextCount = Math.Min(transformSizeHigh, blocksHigh - leftOffset);
-            Array.Fill(leftContexts, culLevel, 0, leftContextCount);
-            Array.Fill(leftContexts, 0, leftContextCount, transformSizeWide - leftContextCount);
-        }
-        else
-        {
-            Array.Fill(leftContexts, culLevel, 0, transformSizeHigh);
-        }
-    }
-
-    private Av1TransformType ComputeTransformType(Av1PlaneType planeType, Av1PartitionInfo partitionInfo, Av1TransformSize transformSize, Av1TransformInfo transformInfo)
-    {
-        Av1TransformType transformType = Av1TransformType.DctDct;
-        if (this.FrameHeader.LosslessArray[partitionInfo.ModeInfo.SegmentId] || transformSize.GetSquareUpSize() > Av1TransformSize.Size32x32)
-        {
-            transformType = Av1TransformType.DctDct;
-        }
-        else
-        {
-            if (planeType == Av1PlaneType.Y)
-            {
-                transformType = transformInfo.Type;
-            }
-            else
-            {
-                // In intra mode, uv planes don't share the same prediction mode as y
-                // plane, so the tx_type should not be shared
-                transformType = ConvertIntraModeToTransformType(partitionInfo.ModeInfo, Av1PlaneType.Uv);
-            }
-        }
-
-        Av1TransformSetType transformSetType = GetExtendedTransformSetType(transformSize, this.FrameHeader.UseReducedTransformSet);
-        if (!transformType.IsExtendedSetUsed(transformSetType))
-        {
-            transformType = Av1TransformType.DctDct;
-        }
-
-        return transformType;
-    }
-
-    private static Av1TransformSetType GetExtendedTransformSetType(Av1TransformSize transformSize, bool useReducedSet)
-    {
-        Av1TransformSize squareUpSize = transformSize.GetSquareUpSize();
-
-        if (squareUpSize >= Av1TransformSize.Size32x32)
-        {
-            return Av1TransformSetType.DctOnly;
-        }
-
-        if (useReducedSet)
-        {
-            return Av1TransformSetType.Dtt4Identity;
-        }
-
-        Av1TransformSize squareSize = transformSize.GetSquareSize();
-        return squareSize == Av1TransformSize.Size16x16 ? Av1TransformSetType.Dtt4Identity : Av1TransformSetType.Dtt4Identity1dDct;
-    }
-
-    private static Av1TransformType ConvertIntraModeToTransformType(Av1BlockModeInfo modeInfo, Av1PlaneType planeType)
-    {
-        Av1PredictionMode mode = (planeType == Av1PlaneType.Y) ? modeInfo.YMode : modeInfo.UvMode;
-        if (mode == Av1PredictionMode.UvChromaFromLuma)
-        {
-            mode = Av1PredictionMode.DC;
-        }
+        Av1BlockSize planeBlockSize = partitionInfo.ModeInfo.BlockSize.GetSubsampled(subX, subY);
+        int blocksWide = partitionInfo.GetMaxBlockWide(planeBlockSize, subX);
+        int blocksHigh = partitionInfo.GetMaxBlockHigh(planeBlockSize, subY);
 
-        return mode.ToTransformType();
+        return reader.ReadCoefficients(partitionInfo.ModeInfo, blockPosition, this.aboveNeighborContext.GetContext(plane), this.leftNeighborContext.GetContext(plane), aboveOffset, leftOffset, plane, blocksWide, blocksHigh, transformBlockContext, transformSize, isLossless, this.FrameHeader.UseReducedTransformSet, transformInfo, partitionInfo.ModeBlockToRightEdge, partitionInfo.ModeBlockToBottomEdge, coefficientBuffer);
     }
 
     private Av1TransformBlockContext GetTransformBlockContext(Av1TransformSize transformSize, int plane, Av1BlockSize planeBlockSize, int transformBlockUnitHighCount, int transformBlockUnitWideCount, int startY, int startX)

tests/ImageSharp.Tests/Formats/Heif/Av1/Av1CoefficientsEntropyTests.cs

@@ -0,0 +1,57 @@
+// Copyright (c) Six Labors.
+// Licensed under the Six Labors Split License.
+
+using System.Buffers;
+using Microsoft.VisualBasic;
+using SixLabors.ImageSharp.Formats.Heif.Av1;
+using SixLabors.ImageSharp.Formats.Heif.Av1.Entropy;
+using SixLabors.ImageSharp.Formats.Heif.Av1.Prediction;
+using SixLabors.ImageSharp.Formats.Heif.Av1.Tiling;
+using SixLabors.ImageSharp.Formats.Heif.Av1.Transform;
+using SixLabors.ImageSharp.Memory;
+
+namespace SixLabors.ImageSharp.Tests.Formats.Heif.Av1;
+
+[Trait("Format", "Avif")]
+public class Av1CoefficientsEntropyTests
+{
+    private const int BaseQIndex = 23;
+
+    [Fact]
+    public void RoundTripZeroEndOfBlock()
+    {
+        // Assign
+        const short transformBlockSkipContext = 0;
+        const short dcSignContext = 0;
+        const int txbIndex = 0;
+        Av1BlockSize blockSize = Av1BlockSize.Block4x4;
+        Av1TransformSize transformSize = Av1TransformSize.Size4x4;
+        Av1TransformType transformType = Av1TransformType.Identity;
+        Av1PredictionMode intraDirection = Av1PredictionMode.DC;
+        Av1ComponentType componentType = Av1ComponentType.Luminance;
+        Av1FilterIntraMode filterIntraMode = Av1FilterIntraMode.DC;
+        ushort endOfBlock = 16;
+        Av1BlockModeInfo modeInfo = new(Av1Constants.MaxPlanes, blockSize, new Point(0, 0));
+        Av1TransformInfo transformInfo = new(transformSize, 0, 0);
+        int[] aboveContexts = new int[1];
+        int[] leftContexts = new int[1];
+        Av1TransformBlockContext transformBlockContext = new();
+        Configuration configuration = Configuration.Default;
+        Av1SymbolEncoder encoder = new(configuration, 100 / 8, BaseQIndex);
+        Span<int> coefficientsBuffer = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16];
+        Span<int> actuals = new int[16];
+
+        // Act
+        encoder.WriteCoefficients(transformSize, transformType, txbIndex, intraDirection, coefficientsBuffer, componentType, transformBlockSkipContext, dcSignContext, endOfBlock, true, BaseQIndex, filterIntraMode);
+
+        using IMemoryOwner<byte> encoded = encoder.Exit();
+
+        Av1SymbolDecoder decoder = new(encoded.GetSpan(), 0);
+        Av1SymbolReader reader = new(encoded.GetSpan());
+        decoder.ReadCoefficients(modeInfo, new Point(0, 0), aboveContexts, leftContexts, 0, 0, 0, 1, 1, transformBlockContext, transformSize, false, true, transformInfo, 0, 0, actuals);
+
+        // Assert
+        Assert.Equal(coefficientsBuffer, actuals);
+    }
+
+}