Fix AV1 encoder coefficient ownership

HEIF_IMPLEMENTATION_PLAN.md

@@ -823,6 +823,7 @@ Encoder verification contract:
 - [~] The tile writer now publishes one packed coefficient context per covered 4x4 edge unit and derives luma/chroma skip plus DC-sign contexts from the complete transform edges using current-libaom units. Complete tile traversal, initialized picture state, and verified CDF update behavior remain.
 - [~] Encoder mode information now uses a frame-owned reference grid over its contiguous allocation, matching current libaom's `mi_grid_base` and `mi_alloc` ownership without per-block tail copies. Signed relative neighbor lookup, 4x4-unit addressing, and mutable selected skip syntax have focused contracts; complete mode decision still remains.
 - [~] Per-block transform, palette-size, and prediction syntax now uses fixed inline storage matching libaom's embedded block state. All 16 transform entries are immediately usable, directional deltas retain their signed range, and traversal performs no managed allocation; picture-level block storage and complete decision state remain.
+- [~] Finalized transform coefficients now use raster-ordered, per-superblock plane segments matching current libaom's frame coefficient-pool geometry. ImageSharp may segment the frame owner between superblock rows, avoiding decoder borders and giant contiguous rentals while preserving disjoint luma/chroma ranges; the forward transform and mode-decision stages still need to populate this owner.
 - [ ] Implement legal deblocking, CDEF, restoration, super-resolution, and film-grain signaling decisions.
 - [~] The coefficient symbol encoder now reuses tile-lifetime level and context workspaces instead of allocating per transform. Every remaining encoder fragment must be audited before it becomes active.
 - [~] The planar conversion, forward transform, and forward quantizer use descending SIMD dispatch: Vector512, Vector256, Vector128, then scalar. Apply the same rule to every later hot-path family.

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

@@ -0,0 +1,109 @@
+// Copyright (c) Six Labors.
+// Licensed under the Six Labors Split License.
+
+using SixLabors.ImageSharp.Formats.Heif.Av1.OpenBitstreamUnit;
+using SixLabors.ImageSharp.Memory;
+
+namespace SixLabors.ImageSharp.Formats.Heif.Av1.Tiling;
+
+/// <summary>
+/// Owns the raster-ordered transform coefficients retained for every superblock in one encoded frame.
+/// </summary>
+internal sealed class Av1EncoderCoefficientBuffer : IDisposable
+{
+    /// <summary>
+    /// Stores one complete superblock's luma and chroma coefficients in each row.
+    /// </summary>
+    private readonly Buffer2D<int> coefficients;
+
+    /// <summary>
+    /// Initializes a new instance of the <see cref="Av1EncoderCoefficientBuffer"/> class.
+    /// </summary>
+    /// <param name="configuration">The configuration providing the frame allocator.</param>
+    /// <param name="sequenceHeader">The sequence header defining superblock and chroma geometry.</param>
+    /// <param name="width">The coded luma width.</param>
+    /// <param name="height">The coded luma height.</param>
+    public Av1EncoderCoefficientBuffer(
+        Configuration configuration,
+        ObuSequenceHeader sequenceHeader,
+        int width,
+        int height)
+    {
+        int superblockSizeLog2 = sequenceHeader.SuperblockSizeLog2;
+        int superblockSize = 1 << superblockSizeLog2;
+        this.SuperblockColumnCount = Av1Math.DivideLog2Ceiling(width, superblockSizeLog2);
+        this.SuperblockRowCount = Av1Math.DivideLog2Ceiling(height, superblockSizeLog2);
+        this.SuperblockCount = this.SuperblockColumnCount * this.SuperblockRowCount;
+        this.LumaCoefficientCount = superblockSize * superblockSize;
+
+        ObuColorConfig colorConfig = sequenceHeader.ColorConfig;
+        int chromaSubsampling = (colorConfig.SubSamplingX ? 1 : 0) + (colorConfig.SubSamplingY ? 1 : 0);
+        this.ChromaCoefficientCount = colorConfig.IsMonochrome ? 0 : this.LumaCoefficientCount >> chromaSubsampling;
+        this.CoefficientsPerSuperblock = this.LumaCoefficientCount + (2 * this.ChromaCoefficientCount);
+
+        // libaom stores finalized coefficients by raster-ordered superblock. A two-dimensional owner preserves that
+        // layout while allowing ImageSharp's allocator to segment the frame instead of demanding one giant rental.
+        this.coefficients = configuration.MemoryAllocator.Allocate2D<int>(
+            this.CoefficientsPerSuperblock,
+            this.SuperblockCount);
+    }
+
+    /// <summary>
+    /// Gets the number of superblock columns covering the coded frame.
+    /// </summary>
+    public int SuperblockColumnCount { get; }
+
+    /// <summary>
+    /// Gets the number of superblock rows covering the coded frame.
+    /// </summary>
+    public int SuperblockRowCount { get; }
+
+    /// <summary>
+    /// Gets the number of superblocks covering the coded frame.
+    /// </summary>
+    public int SuperblockCount { get; }
+
+    /// <summary>
+    /// Gets the number of luma coefficient positions reserved for each superblock.
+    /// </summary>
+    public int LumaCoefficientCount { get; }
+
+    /// <summary>
+    /// Gets the number of coefficient positions reserved for each chroma plane in each superblock.
+    /// </summary>
+    public int ChromaCoefficientCount { get; }
+
+    /// <summary>
+    /// Gets the number of coefficient positions reserved for each complete superblock.
+    /// </summary>
+    public int CoefficientsPerSuperblock { get; }
+
+    /// <summary>
+    /// Gets the total number of coefficient positions retained for the frame.
+    /// </summary>
+    public long TotalCoefficientCount => (long)this.CoefficientsPerSuperblock * this.SuperblockCount;
+
+    /// <summary>
+    /// Gets one component plane's coefficient span for a raster-ordered superblock.
+    /// </summary>
+    /// <param name="superblockIndex">The raster-ordered superblock index.</param>
+    /// <param name="plane">The requested component plane.</param>
+    /// <returns>The complete coefficient span reserved for that plane and superblock.</returns>
+    public Span<int> GetPlaneSpan(int superblockIndex, Av1Plane plane)
+    {
+        Span<int> superblock = this.coefficients.DangerousGetRowSpan(superblockIndex);
+        return plane switch
+        {
+            Av1Plane.Y => superblock[..this.LumaCoefficientCount],
+            Av1Plane.U => superblock.Slice(this.LumaCoefficientCount, this.ChromaCoefficientCount),
+            _ => superblock.Slice(
+                this.LumaCoefficientCount + this.ChromaCoefficientCount,
+                this.ChromaCoefficientCount)
+        };
+    }
+
+    /// <summary>
+    /// Releases the frame coefficient storage.
+    /// </summary>
+    public void Dispose() => this.coefficients.Dispose();
+}

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

@@ -59,14 +59,14 @@ internal partial class Av1TileWriter
     /// <param name="ec_ctx">The entropy-coding position state for the superblock.</param>
     /// <param name="writer">The tile symbol encoder.</param>
     /// <param name="superblock">The encoder decisions for the superblock.</param>
-    /// <param name="frameBuffer">The transformed coefficients for the frame.</param>
+    /// <param name="coefficientBuffer">The transformed coefficients retained by raster-ordered superblock.</param>
     /// <param name="tileIndex">The zero-based tile index.</param>
     public static void WriteSuperblock(
         Av1PictureControlSet pcs,
         Av1EntropyCodingContext ec_ctx,
         ref Av1SymbolEncoder writer,
         Av1Superblock superblock,
-        Av1FrameBuffer<int> frameBuffer,
+        Av1EncoderCoefficientBuffer coefficientBuffer,
         ushort tileIndex)
     {
         Av1SequenceControlSet scs = pcs.Sequence;
@@ -124,79 +124,79 @@ internal partial class Av1TileWriter
                 switch (superblock.CodingUnitPartitionTypes[blockIndex])
                 {
                     case Av1PartitionType.None:
-                        WriteModesBlock(pcs, ec_ctx, ref writer, superblock, blk_ptr, tileIndex, frameBuffer);
+                        WriteModesBlock(pcs, ec_ctx, ref writer, superblock, blk_ptr, tileIndex, coefficientBuffer);
                         break;
 
                     case Av1PartitionType.Horizontal:
-                        WriteModesBlock(pcs, ec_ctx, ref writer, superblock, blk_ptr, tileIndex, frameBuffer);
+                        WriteModesBlock(pcs, ec_ctx, ref writer, superblock, blk_ptr, tileIndex, coefficientBuffer);
 
                         if (mi_row + hbs < cm.ModeInfoRowCount)
                         {
                             finalBlockIndex++;
                             blk_ptr = superblock.FinalBlocks[finalBlockIndex];
-                            WriteModesBlock(pcs, ec_ctx, ref writer, superblock, blk_ptr, tileIndex, frameBuffer);
+                            WriteModesBlock(pcs, ec_ctx, ref writer, superblock, blk_ptr, tileIndex, coefficientBuffer);
                         }
 
                         break;
 
                     case Av1PartitionType.Vertical:
-                        WriteModesBlock(pcs, ec_ctx, ref writer, superblock, blk_ptr, tileIndex, frameBuffer);
+                        WriteModesBlock(pcs, ec_ctx, ref writer, superblock, blk_ptr, tileIndex, coefficientBuffer);
                         if (mi_col + hbs < cm.ModeInfoColumnCount)
                         {
                             finalBlockIndex++;
                             blk_ptr = superblock.FinalBlocks[finalBlockIndex];
-                            WriteModesBlock(pcs, ec_ctx, ref writer, superblock, blk_ptr, tileIndex, frameBuffer);
+                            WriteModesBlock(pcs, ec_ctx, ref writer, superblock, blk_ptr, tileIndex, coefficientBuffer);
                         }
 
                         break;
                     case Av1PartitionType.Split:
                         break;
                     case Av1PartitionType.HorizontalA:
-                        WriteModesBlock(pcs, ec_ctx, ref writer, superblock, blk_ptr, tileIndex, frameBuffer);
+                        WriteModesBlock(pcs, ec_ctx, ref writer, superblock, blk_ptr, tileIndex, coefficientBuffer);
 
                         finalBlockIndex++;
                         blk_ptr = superblock.FinalBlocks[finalBlockIndex];
-                        WriteModesBlock(pcs, ec_ctx, ref writer, superblock, blk_ptr, tileIndex, frameBuffer);
+                        WriteModesBlock(pcs, ec_ctx, ref writer, superblock, blk_ptr, tileIndex, coefficientBuffer);
 
                         finalBlockIndex++;
                         blk_ptr = superblock.FinalBlocks[finalBlockIndex];
-                        WriteModesBlock(pcs, ec_ctx, ref writer, superblock, blk_ptr, tileIndex, frameBuffer);
+                        WriteModesBlock(pcs, ec_ctx, ref writer, superblock, blk_ptr, tileIndex, coefficientBuffer);
 
                         break;
                     case Av1PartitionType.HorizontalB:
-                        WriteModesBlock(pcs, ec_ctx, ref writer, superblock, blk_ptr, tileIndex, frameBuffer);
+                        WriteModesBlock(pcs, ec_ctx, ref writer, superblock, blk_ptr, tileIndex, coefficientBuffer);
 
                         finalBlockIndex++;
                         blk_ptr = superblock.FinalBlocks[finalBlockIndex];
-                        WriteModesBlock(pcs, ec_ctx, ref writer, superblock, blk_ptr, tileIndex, frameBuffer);
+                        WriteModesBlock(pcs, ec_ctx, ref writer, superblock, blk_ptr, tileIndex, coefficientBuffer);
 
                         finalBlockIndex++;
                         blk_ptr = superblock.FinalBlocks[finalBlockIndex];
-                        WriteModesBlock(pcs, ec_ctx, ref writer, superblock, blk_ptr, tileIndex, frameBuffer);
+                        WriteModesBlock(pcs, ec_ctx, ref writer, superblock, blk_ptr, tileIndex, coefficientBuffer);
 
                         break;
                     case Av1PartitionType.VerticalA:
-                        WriteModesBlock(pcs, ec_ctx, ref writer, superblock, blk_ptr, tileIndex, frameBuffer);
+                        WriteModesBlock(pcs, ec_ctx, ref writer, superblock, blk_ptr, tileIndex, coefficientBuffer);
 
                         finalBlockIndex++;
                         blk_ptr = superblock.FinalBlocks[finalBlockIndex];
-                        WriteModesBlock(pcs, ec_ctx, ref writer, superblock, blk_ptr, tileIndex, frameBuffer);
+                        WriteModesBlock(pcs, ec_ctx, ref writer, superblock, blk_ptr, tileIndex, coefficientBuffer);
 
                         finalBlockIndex++;
                         blk_ptr = superblock.FinalBlocks[finalBlockIndex];
-                        WriteModesBlock(pcs, ec_ctx, ref writer, superblock, blk_ptr, tileIndex, frameBuffer);
+                        WriteModesBlock(pcs, ec_ctx, ref writer, superblock, blk_ptr, tileIndex, coefficientBuffer);
 
                         break;
                     case Av1PartitionType.VerticalB:
-                        WriteModesBlock(pcs, ec_ctx, ref writer, superblock, blk_ptr, tileIndex, frameBuffer);
+                        WriteModesBlock(pcs, ec_ctx, ref writer, superblock, blk_ptr, tileIndex, coefficientBuffer);
 
                         finalBlockIndex++;
                         blk_ptr = superblock.FinalBlocks[finalBlockIndex];
-                        WriteModesBlock(pcs, ec_ctx, ref writer, superblock, blk_ptr, tileIndex, frameBuffer);
+                        WriteModesBlock(pcs, ec_ctx, ref writer, superblock, blk_ptr, tileIndex, coefficientBuffer);
 
                         finalBlockIndex++;
                         blk_ptr = superblock.FinalBlocks[finalBlockIndex];
-                        WriteModesBlock(pcs, ec_ctx, ref writer, superblock, blk_ptr, tileIndex, frameBuffer);
+                        WriteModesBlock(pcs, ec_ctx, ref writer, superblock, blk_ptr, tileIndex, coefficientBuffer);
 
                         break;
                     case Av1PartitionType.Horizontal4:
@@ -217,7 +217,7 @@ internal partial class Av1TileWriter
                                 blk_ptr = superblock.FinalBlocks[finalBlockIndex];
                             }
 
-                            WriteModesBlock(pcs, ec_ctx, ref writer, superblock, blk_ptr, tileIndex, frameBuffer);
+                            WriteModesBlock(pcs, ec_ctx, ref writer, superblock, blk_ptr, tileIndex, coefficientBuffer);
                         }
 
                         break;
@@ -239,7 +239,7 @@ internal partial class Av1TileWriter
                                 blk_ptr = superblock.FinalBlocks[finalBlockIndex];
                             }
 
-                            WriteModesBlock(pcs, ec_ctx, ref writer, superblock, blk_ptr, tileIndex, frameBuffer);
+                            WriteModesBlock(pcs, ec_ctx, ref writer, superblock, blk_ptr, tileIndex, coefficientBuffer);
                         }
 
                         break;
@@ -345,7 +345,7 @@ internal partial class Av1TileWriter
     /// <param name="tb_ptr">The containing superblock.</param>
     /// <param name="blk_ptr">The final encoder decisions for the block.</param>
     /// <param name="tile_idx">The zero-based tile index.</param>
-    /// <param name="coeff_ptr">The transformed coefficients for the frame.</param>
+    /// <param name="coefficientBuffer">The transformed coefficients retained by raster-ordered superblock.</param>
     private static void WriteModesBlock(
         Av1PictureControlSet pcs,
         Av1EntropyCodingContext entropyCodingContext,
@@ -353,7 +353,7 @@ internal partial class Av1TileWriter
         Av1Superblock tb_ptr,
         Av1EncoderBlockStruct blk_ptr,
         ushort tile_idx,
-        Av1FrameBuffer<int> coeff_ptr)
+        Av1EncoderCoefficientBuffer coefficientBuffer)
     {
         Av1SequenceControlSet scs = pcs.Sequence;
         ObuFrameHeader frm_hdr = pcs.Parent.FrameHeader;
@@ -530,7 +530,8 @@ internal partial class Av1TileWriter
                     blockOrigin,
                     intra_luma_mode,
                     blockSize,
-                    coeff_ptr,
+                    coefficientBuffer,
+                    tb_ptr.Index,
                     luma_dc_sign_level_coeff_na,
                     cr_dc_sign_level_coeff_na,
                     cb_dc_sign_level_coeff_na);
@@ -951,11 +952,11 @@ internal partial class Av1TileWriter
     /// <param name="blockOrigin">The block origin in samples.</param>
     /// <param name="intraLumaDir">The luma prediction direction.</param>
     /// <param name="planeBlockSize">The luma block size.</param>
-    /// <param name="coeff_ptr">The transformed coefficients for the frame.</param>
+    /// <param name="coefficientBuffer">The transformed coefficients retained by raster-ordered superblock.</param>
+    /// <param name="superblockIndex">The raster-ordered index of the containing superblock.</param>
     /// <param name="luma_dc_sign_level_coeff_na">The luma coefficient neighbor contexts.</param>
     /// <param name="cr_dc_sign_level_coeff_na">The red-difference chroma coefficient neighbor contexts.</param>
     /// <param name="cb_dc_sign_level_coeff_na">The blue-difference chroma coefficient neighbor contexts.</param>
-    /// <exception cref="NotImplementedException">The transform-depth path required by the block is not implemented.</exception>
     private static void EncodeCoefficients1d(
         Av1PictureControlSet pcs,
         Av1EntropyCodingContext ec_ctx,
@@ -965,42 +966,38 @@ internal partial class Av1TileWriter
         Point blockOrigin,
         Av1PredictionMode intraLumaDir,
         Av1BlockSize planeBlockSize,
-        Av1FrameBuffer<int> coeff_ptr,
+        Av1EncoderCoefficientBuffer coefficientBuffer,
+        int superblockIndex,
         Av1NeighborArrayUnit<byte> luma_dc_sign_level_coeff_na,
         Av1NeighborArrayUnit<byte> cr_dc_sign_level_coeff_na,
         Av1NeighborArrayUnit<byte> cb_dc_sign_level_coeff_na)
     {
-        if (mbmi.Block.TransformDepth != 0)
+        EncodeTransformCoefficientsY(
-        {
+            pcs,
-            EncodeTransformCoefficientsY(
+            ec_ctx,
-                pcs,
+            ref writer,
-                ec_ctx,
+            mbmi,
-                ref writer,
+            blk_ptr,
-                mbmi,
+            blockOrigin,
-                blk_ptr,
+            intraLumaDir,
-                blockOrigin,
+            planeBlockSize,
-                intraLumaDir,
+            coefficientBuffer,
-                planeBlockSize,
+            superblockIndex,
-                coeff_ptr,
+            luma_dc_sign_level_coeff_na);
-                luma_dc_sign_level_coeff_na);
 
-            EncodeTransformCoefficientsUv(
+        EncodeTransformCoefficientsUv(
-                pcs,
+            pcs,
-                ec_ctx,
+            ec_ctx,
-                ref writer,
+            ref writer,
-                mbmi,
+            mbmi,
-                blk_ptr,
+            blk_ptr,
-                blockOrigin,
+            blockOrigin,
-                intraLumaDir,
+            intraLumaDir,
-                planeBlockSize,
+            planeBlockSize,
-                coeff_ptr,
+            coefficientBuffer,
-                cr_dc_sign_level_coeff_na,
+            superblockIndex,
-                cb_dc_sign_level_coeff_na);
+            cr_dc_sign_level_coeff_na,
-        }
+            cb_dc_sign_level_coeff_na);
-        else
-        {
-            throw new NotImplementedException("Only capable to encode Largest transform mode.");
-        }
     }
 
     /// <summary>
@@ -1014,7 +1011,8 @@ internal partial class Av1TileWriter
     /// <param name="blockOrigin">The block origin in samples.</param>
     /// <param name="intraLumaDir">The luma prediction direction.</param>
     /// <param name="plane_bsize">The luma block size.</param>
-    /// <param name="coeff_ptr">The transformed coefficients for the frame.</param>
+    /// <param name="coefficientBuffer">The transformed coefficients retained by raster-ordered superblock.</param>
+    /// <param name="superblockIndex">The raster-ordered index of the containing superblock.</param>
     /// <param name="luma_dc_sign_level_coeff_na">The luma coefficient neighbor contexts.</param>
     public static void EncodeTransformCoefficientsY(
         Av1PictureControlSet pcs,
@@ -1025,7 +1023,8 @@ internal partial class Av1TileWriter
         Point blockOrigin,
         Av1PredictionMode intraLumaDir,
         Av1BlockSize plane_bsize,
-        Av1FrameBuffer<int> coeff_ptr,
+        Av1EncoderCoefficientBuffer coefficientBuffer,
+        int superblockIndex,
         Av1NeighborArrayUnit<byte> luma_dc_sign_level_coeff_na)
     {
         // This writer currently emits intra frames, so coefficient contexts use only intra prediction state.
@@ -1033,6 +1032,7 @@ internal partial class Av1TileWriter
         int tx_depth = mbmi.Block.TransformDepth;
         int txb_count = blockGeometry.TransformBlockCount[mbmi.Block.TransformDepth];
         ObuFrameHeader frameHeader = pcs.Parent.FrameHeader;
+        Span<int> lumaCoefficients = coefficientBuffer.GetPlaneSpan(superblockIndex, Av1Plane.Y);
 
         for (int tx_index = 0; tx_index < txb_count; tx_index++)
         {
@@ -1041,7 +1041,7 @@ internal partial class Av1TileWriter
             Av1TransformSize tx_size = blockGeometry.TransformSize[tx_depth];
 
             int coeff1d_offset = entropyCodingContext.CodedAreaSuperblock;
-            Span<int> coeff_buffer = coeff_ptr.GetPlaneBuffer(Av1Plane.Y).DangerousGetSingleSpan()[coeff1d_offset..];
+            Span<int> coeff_buffer = lumaCoefficients[coeff1d_offset..];
 
             Point transformOrigin = blockGeometry.TransformOrigin[tx_depth][txb_itr];
             Av1TransformBlockContext blockContext = GetTransformBlockContexts(
@@ -1096,7 +1096,8 @@ internal partial class Av1TileWriter
     /// <param name="blockOrigin">The luma block origin in samples.</param>
     /// <param name="intraLumaDir">The luma prediction direction used by coefficient contexts.</param>
     /// <param name="plane_bsize">The luma block size.</param>
-    /// <param name="coeff_ptr">The transformed coefficients for the frame.</param>
+    /// <param name="coefficientBuffer">The transformed coefficients retained by raster-ordered superblock.</param>
+    /// <param name="superblockIndex">The raster-ordered index of the containing superblock.</param>
     /// <param name="cr_dc_sign_level_coeff_na">The red-difference chroma coefficient neighbor contexts.</param>
     /// <param name="cb_dc_sign_level_coeff_na">The blue-difference chroma coefficient neighbor contexts.</param>
     private static void EncodeTransformCoefficientsUv(
@@ -1108,7 +1109,8 @@ internal partial class Av1TileWriter
         Point blockOrigin,
         Av1PredictionMode intraLumaDir,
         Av1BlockSize plane_bsize,
-        Av1FrameBuffer<int> coeff_ptr,
+        Av1EncoderCoefficientBuffer coefficientBuffer,
+        int superblockIndex,
         Av1NeighborArrayUnit<byte> cr_dc_sign_level_coeff_na,
         Av1NeighborArrayUnit<byte> cb_dc_sign_level_coeff_na)
     {
@@ -1120,19 +1122,22 @@ internal partial class Av1TileWriter
         }
 
         int tx_depth = mbmi.Block.TransformDepth;
-        int transformBlockCount = 1;
         ObuFrameHeader frameHeader = pcs.Parent.FrameHeader;
+        Av1TransformSize chromaTransformSize = blockGeometry.TransformSizeUv[tx_depth];
+        int transformWidth = chromaTransformSize.GetWidth();
+        int transformHeight = chromaTransformSize.GetHeight();
+        int transformBlockCount = (blockGeometry.BlockWidthUv * blockGeometry.BlockHeightUv) /
+            (transformWidth * transformHeight);
+        Span<int> blueCoefficients = coefficientBuffer.GetPlaneSpan(superblockIndex, Av1Plane.U);
+        Span<int> redCoefficients = coefficientBuffer.GetPlaneSpan(superblockIndex, Av1Plane.V);
 
         for (int transformBlockIndex = 0; transformBlockIndex < transformBlockCount; ++transformBlockIndex)
         {
-            Av1TransformSize chromaTransformSize = blockGeometry.TransformSizeUv[tx_depth];
-            int transformWidth = chromaTransformSize.GetWidth();
-            int transformHeight = chromaTransformSize.GetHeight();
             Point transformOrigin = blockGeometry.TransformOrigin[tx_depth][transformBlockIndex];
             Point chromaOrigin = RoundUv(blockOrigin + (Size)transformOrigin - (Size)blockGeometry.Origin) >> 1;
 
             // Both chroma planes share transform geometry but retain independent coefficient contexts.
-            Span<int> coefficientBuffer = coeff_ptr.GetPlaneBuffer(Av1Plane.U).DangerousGetSingleSpan().Slice(entropyCodingContext.CodedAreaSuperblockUv);
+            Span<int> coefficients = blueCoefficients[entropyCodingContext.CodedAreaSuperblockUv..];
             Av1TransformBlockContext blockContext = GetTransformBlockContexts(
                 Av1ComponentType.Chroma,
                 cb_dc_sign_level_coeff_na,
@@ -1145,14 +1150,14 @@ internal partial class Av1TileWriter
                 chromaTransformSize,
                 chromaTransformType,
                 intraLumaDir,
-                coefficientBuffer,
+                coefficients,
                 Av1ComponentType.Chroma,
                 blockContext,
                 (ushort)endOfBlockCb,
                 frameHeader.UseReducedTransformSet,
                 blk_ptr.FilterIntraMode);
 
-            coefficientBuffer = coeff_ptr.GetPlaneBuffer(Av1Plane.V).DangerousGetSingleSpan().Slice(entropyCodingContext.CodedAreaSuperblockUv);
+            coefficients = redCoefficients[entropyCodingContext.CodedAreaSuperblockUv..];
             int endOfBlockCr = blk_ptr.TransformBlocks[transformBlockIndex].NzCoefficientCount[2];
 
             blockContext = GetTransformBlockContexts(
@@ -1166,7 +1171,7 @@ internal partial class Av1TileWriter
                 chromaTransformSize,
                 chromaTransformType,
                 intraLumaDir,
-                coefficientBuffer,
+                coefficients,
                 Av1ComponentType.Chroma,
                 blockContext,
                 (ushort)endOfBlockCr,

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

@@ -190,6 +190,63 @@ public class Av1CoefficientsEntropyTests
         Assert.Equal(0, allocated);
     }
 
+    [Theory]
+    [InlineData(false, 6, 4096, 1024, 6144, 36864L)]
+    [InlineData(true, 2, 16384, 4096, 24576, 49152L)]
+    public void EncoderCoefficientBufferMatchesLibaom420SuperblockLayout(
+        bool use128x128Superblock,
+        int expectedSuperblockCount,
+        int expectedLumaCount,
+        int expectedChromaCount,
+        int expectedCoefficientsPerSuperblock,
+        long expectedTotalCoefficientCount)
+    {
+        ObuSequenceHeader sequenceHeader = new() { Use128x128Superblock = use128x128Superblock };
+        sequenceHeader.ColorConfig.IsMonochrome = false;
+        sequenceHeader.ColorConfig.SubSamplingX = true;
+        sequenceHeader.ColorConfig.SubSamplingY = true;
+
+        using Av1EncoderCoefficientBuffer coefficients = new(
+            Configuration.Default,
+            sequenceHeader,
+            width: 129,
+            height: 65);
+
+        Assert.Equal(expectedSuperblockCount, coefficients.SuperblockCount);
+        Assert.Equal(expectedLumaCount, coefficients.LumaCoefficientCount);
+        Assert.Equal(expectedChromaCount, coefficients.ChromaCoefficientCount);
+        Assert.Equal(expectedCoefficientsPerSuperblock, coefficients.CoefficientsPerSuperblock);
+        Assert.Equal(expectedTotalCoefficientCount, coefficients.TotalCoefficientCount);
+        Assert.Equal(expectedLumaCount, coefficients.GetPlaneSpan(0, Av1Plane.Y).Length);
+        Assert.Equal(expectedChromaCount, coefficients.GetPlaneSpan(0, Av1Plane.U).Length);
+        Assert.Equal(expectedChromaCount, coefficients.GetPlaneSpan(0, Av1Plane.V).Length);
+    }
+
+    [Fact]
+    public void EncoderCoefficientBufferKeepsEveryPlaneAndSuperblockDisjoint()
+    {
+        ObuSequenceHeader sequenceHeader = new() { Use128x128Superblock = true };
+        sequenceHeader.ColorConfig.IsMonochrome = false;
+        sequenceHeader.ColorConfig.SubSamplingX = true;
+        sequenceHeader.ColorConfig.SubSamplingY = true;
+
+        using Av1EncoderCoefficientBuffer coefficients = new(
+            Configuration.Default,
+            sequenceHeader,
+            width: 129,
+            height: 65);
+
+        coefficients.GetPlaneSpan(0, Av1Plane.Y)[0] = 11;
+        coefficients.GetPlaneSpan(0, Av1Plane.U)[0] = 22;
+        coefficients.GetPlaneSpan(0, Av1Plane.V)[0] = 33;
+        coefficients.GetPlaneSpan(1, Av1Plane.Y)[0] = 44;
+
+        Assert.Equal(11, coefficients.GetPlaneSpan(0, Av1Plane.Y)[0]);
+        Assert.Equal(22, coefficients.GetPlaneSpan(0, Av1Plane.U)[0]);
+        Assert.Equal(33, coefficients.GetPlaneSpan(0, Av1Plane.V)[0]);
+        Assert.Equal(44, coefficients.GetPlaneSpan(1, Av1Plane.Y)[0]);
+    }
+
     [Fact]
     public void RoundTripZeroEndOfBlock()
     {