Write AV1 selected transform sizes

HEIF_IMPLEMENTATION_PLAN.md

@@ -824,7 +824,7 @@ Encoder verification contract:
 - [~] 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.
 - [~] Final block decisions now use contiguous value storage with palette and prediction syntax inline. Macroblock edge and neighbor state is reused by the entropy-coding operation instead of allocating one managed object per final block; directional deltas retain their signed range, while complete mode decision still remains.
 - [~] Finalized transform coefficients and packed EOB/type state now use raster-ordered, per-superblock plane segments matching current libaom's coefficient-pool geometry. One ImageSharp allocator owner replaces libaom's separate coefficient, EOB, and entropy-context allocations while preserving the full 1024 luma and 256-per-chroma 4x4 state capacity of a 128x128 4:2:0 superblock; the forward transform and mode-decision stages still need to populate this owner.
-- [~] Tile partition writing now follows current libaom's recursive `write_modes_sb` preorder traversal and `update_ext_partition_context` edge updates directly. The obsolete SVT-derived global geometry catalog and its unimplemented lookup are removed; transform geometry is derived in libaom's bounded 64x64 residual order from the selected block transform size, frame-edge and segmentation syntax use mode-information units, and 128x128 CDEF units use libaom's 0-to-3 indexing and first-block strength ownership. Partition analysis and selected-transform syntax still need to populate and encode these retained decisions.
+- [~] Tile partition writing now follows current libaom's recursive `write_modes_sb` preorder traversal and `update_ext_partition_context` edge updates directly. The obsolete SVT-derived global geometry catalog and its unimplemented lookup are removed; transform geometry is derived in libaom's bounded 64x64 residual order, fixed intra transform-size symbols use the reference depth and neighbor contexts, frame-edge and segmentation syntax use mode-information units, and 128x128 CDEF units use libaom's 0-to-3 indexing and first-block strength ownership. Partition and mode analysis still need to populate these retained decisions; variable inter-transform syntax remains part of later inter-frame support.
 - [ ] 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/Entropy/Av1SymbolEncoder.cs

@@ -101,6 +101,11 @@ internal class Av1SymbolEncoder : IDisposable
     /// </summary>
     private readonly Av1Distribution[][][] intraExtendedTransform;
 
+    /// <summary>
+    /// The tile-adaptive fixed transform-size distributions.
+    /// </summary>
+    private readonly Av1Distribution[][] transformSize;
+
     /// <summary>
     /// The tile-adaptive spatial segment-identifier distributions.
     /// </summary>
@@ -175,6 +180,7 @@ internal class Av1SymbolEncoder : IDisposable
         this.filterIntraMode = Av1DefaultDistributions.FilterIntraMode;
         this.deltaQuantizerAbsolute = Av1DefaultDistributions.DeltaQuantizerAbsolute;
         this.intraExtendedTransform = Av1DefaultDistributions.IntraExtendedTransform;
+        this.transformSize = Av1DefaultDistributions.TransformSize;
         this.segmentId = Av1DefaultDistributions.SegmentId;
         this.angleDelta = Av1DefaultDistributions.AngleDelta;
         this.skip = Av1DefaultDistributions.Skip;
@@ -425,6 +431,36 @@ internal class Av1SymbolEncoder : IDisposable
         w.WriteSymbol(skip, this.transformBlockSkip[(int)transformSizeContext][skipContext]);
     }
 
+    /// <summary>
+    /// Writes the selected transform size as its subdivision depth from the block maximum.
+    /// </summary>
+    /// <param name="blockSize">The block size defining the maximum transform.</param>
+    /// <param name="transformSize">The selected transform size.</param>
+    /// <param name="context">The neighboring transform-size context.</param>
+    public void WriteTransformSize(Av1BlockSize blockSize, Av1TransformSize transformSize, int context)
+    {
+        Av1TransformSize maximumTransformSize = blockSize.GetMaximumTransformSize();
+        Av1TransformSize currentTransformSize = maximumTransformSize;
+        int categoryDepth = 0;
+        while (currentTransformSize != Av1TransformSize.Size4x4)
+        {
+            categoryDepth++;
+            currentTransformSize = currentTransformSize.GetSubSize();
+        }
+
+        int selectedDepth = 0;
+        currentTransformSize = maximumTransformSize;
+        while (currentTransformSize != transformSize && selectedDepth < Av1Constants.MaxVarTransform)
+        {
+            selectedDepth++;
+            currentTransformSize = currentTransformSize.GetSubSize();
+        }
+
+        DebugGuard.IsTrue(currentTransformSize == transformSize, nameof(transformSize));
+        ref Av1SymbolWriter w = ref this.writer;
+        w.WriteSymbol(selectedDepth, this.transformSize[categoryDepth - 1][context]);
+    }
+
     /// <summary>
     /// Finalizes the range-coded tile payload and transfers ownership of its memory.
     /// </summary>

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

@@ -748,6 +748,15 @@ internal partial class Av1TileWriter
                 }
             }
 
+            WriteTransformSize(
+                pcs,
+                writer,
+                macroBlockModeInfo,
+                macroBlock,
+                blockSize,
+                blockOrigin,
+                tile_idx);
+
             if (!skipWritingCoefficients)
             {
                 EncodeCoefficients1d(
@@ -770,6 +779,66 @@ internal partial class Av1TileWriter
         UpdateNeighbors(pcs, entropyCodingContext, blockOrigin, ref blk_ptr, tile_idx, blockSize);
     }
 
+    /// <summary>
+    /// Writes or derives the block transform size and publishes its edge contexts.
+    /// </summary>
+    /// <param name="pcs">The picture coding state.</param>
+    /// <param name="writer">The tile symbol encoder.</param>
+    /// <param name="macroBlockModeInfo">The selected block modes.</param>
+    /// <param name="macroBlock">The reusable macroblock edge and neighbor state.</param>
+    /// <param name="blockSize">The block size.</param>
+    /// <param name="blockOrigin">The block origin in samples.</param>
+    /// <param name="tileIndex">The zero-based tile index.</param>
+    internal static void WriteTransformSize(
+        Av1PictureControlSet pcs,
+        Av1SymbolEncoder writer,
+        Av1MacroBlockModeInfo macroBlockModeInfo,
+        Av1MacroBlockD macroBlock,
+        Av1BlockSize blockSize,
+        Point blockOrigin,
+        int tileIndex)
+    {
+        ObuFrameHeader frameHeader = pcs.Parent.FrameHeader;
+        bool isLossless = frameHeader.LosslessArray[macroBlockModeInfo.Block.SegmentId];
+        bool writesTransformSize = !isLossless &&
+            frameHeader.TransformMode == Av1TransformMode.Select &&
+            blockSize > Av1BlockSize.Block4x4;
+        Av1TransformSize transformSize = isLossless
+            ? Av1TransformSize.Size4x4
+            : writesTransformSize
+                ? macroBlockModeInfo.Block.TransformSize
+                : blockSize.GetMaximumTransformSize();
+
+        macroBlockModeInfo.Block.TransformSize = transformSize;
+        Av1NeighborArrayUnit<byte> transformContexts = pcs.TransformFunctionContexts[tileIndex];
+        if (writesTransformSize)
+        {
+            Av1TransformSize maximumTransformSize = blockSize.GetMaximumTransformSize();
+            int above = transformContexts.Top[transformContexts.GetTopIndex(blockOrigin)] >= maximumTransformSize.GetWidth() ? 1 : 0;
+            int left = transformContexts.Left[transformContexts.GetLeftIndex(blockOrigin)] >= maximumTransformSize.GetHeight() ? 1 : 0;
+            int context = macroBlock.IsUpAvailable
+                ? macroBlock.IsLeftAvailable ? above + left : above
+                : macroBlock.IsLeftAvailable ? left : 0;
+
+            writer.WriteTransformSize(blockSize, transformSize, context);
+        }
+
+        Size blockDimensions = new(blockSize.GetWidth(), blockSize.GetHeight());
+
+        // Above entries retain transform widths and left entries retain heights, including rectangular selections.
+        transformContexts.UnitModeWrite(
+            (byte)transformSize.GetWidth(),
+            blockOrigin,
+            blockDimensions,
+            Av1NeighborArrayUnit<byte>.UnitMask.Top);
+
+        transformContexts.UnitModeWrite(
+            (byte)transformSize.GetHeight(),
+            blockOrigin,
+            blockDimensions,
+            Av1NeighborArrayUnit<byte>.UnitMask.Left);
+    }
+
     /// <summary>
     /// Writes the chroma intra mode, chroma-from-luma alpha values, and directional angle adjustment for a block.
     /// </summary>

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

@@ -346,6 +346,70 @@ public class Av1CoefficientsEntropyTests
         }
     }
 
+    [Fact]
+    public void SelectedTransformSizeRoundTripsAndPublishesRectangularEdgeContexts()
+    {
+        Av1PictureControlSet picture = CreateEncoderPicture(16, 16);
+        picture.Parent.FrameHeader.TransformMode = Av1TransformMode.Select;
+        Av1MacroBlockModeInfo modeInfo = picture.ModeInfoAllocation[0].MacroBlockModeInfo;
+        modeInfo.Block.BlockSize = Av1BlockSize.Block16x32;
+        modeInfo.Block.TransformSize = Av1TransformSize.Size8x8;
+        modeInfo.Block.SegmentId = 0;
+        Point blockOrigin = new(16, 16);
+        Av1MacroBlockD macroBlock = new()
+        {
+            Tile = new Av1TileInfo(0, 0, picture.Parent.FrameHeader),
+            IsUpAvailable = true,
+            IsLeftAvailable = true
+        };
+
+        using Av1NeighborArrayUnit<byte> transforms = new(
+            Configuration.Default,
+            leftSize: 64,
+            topSize: 64,
+            topLeftSize: 128)
+        {
+            GranularityNormalLog2 = Av1Constants.ModeInfoSizeLog2,
+            GranularityTopLeftLog2 = Av1Constants.ModeInfoSizeLog2
+        };
+
+        int topIndex = transforms.GetTopIndex(blockOrigin);
+        int leftIndex = transforms.GetLeftIndex(blockOrigin);
+        transforms.Top[topIndex] = 16;
+        transforms.Left[leftIndex] = 16;
+        picture.TransformFunctionContexts = [transforms];
+
+        Av1SymbolEncoder writer = new(Configuration.Default, 64, BaseQIndex);
+        Av1TileWriter.WriteTransformSize(
+            picture,
+            writer,
+            modeInfo,
+            macroBlock,
+            modeInfo.Block.BlockSize,
+            blockOrigin,
+            tileIndex: 0);
+
+        using IMemoryOwner<byte> encoded = writer.Exit();
+        writer.Dispose();
+
+        Av1SymbolDecoder reader = new(Configuration.Default, encoded.GetSpan(), BaseQIndex);
+        Assert.Equal(
+            Av1TransformSize.Size8x8,
+            reader.ReadTransformSize(Av1BlockSize.Block16x32, context: 1));
+
+        for (int index = 0; index < transforms.Top.Length; index++)
+        {
+            byte expected = index >= topIndex && index < topIndex + 4 ? (byte)8 : (byte)0;
+            Assert.Equal(expected, transforms.Top[index]);
+        }
+
+        for (int index = 0; index < transforms.Left.Length; index++)
+        {
+            byte expected = index >= leftIndex && index < leftIndex + 8 ? (byte)8 : (byte)0;
+            Assert.Equal(expected, transforms.Left[index]);
+        }
+    }
+
     [Theory]
     [InlineData((int)Av1PartitionType.None, 24, 24)]
     [InlineData((int)Av1PartitionType.Horizontal, 24, 28)]
@@ -510,10 +574,21 @@ public class Av1CoefficientsEntropyTests
             GranularityTopLeftLog2 = 2
         };
 
+        using Av1NeighborArrayUnit<byte> transforms = new(
+            Configuration.Default,
+            leftSize: 16,
+            topSize: 32,
+            topLeftSize: 48)
+        {
+            GranularityNormalLog2 = 2,
+            GranularityTopLeftLog2 = 2
+        };
+
         picture.PartitionContexts = [partitions];
         picture.LuminanceDcSignLevelCoefficientNeighbors = [luma];
         picture.CrDcSignLevelCoefficientNeighbors = [red];
         picture.CbDcSignLevelCoefficientNeighbors = [blue];
+        picture.TransformFunctionContexts = [transforms];
         Av1TileInfo tile = new(0, 0, picture.Parent.FrameHeader);
         Point[] blockPositions = [new(16, 0), new(24, 0), new(16, 8), new(24, 8)];
         Av1EncoderBlockStruct[] blocks = new Av1EncoderBlockStruct[blockPositions.Length];
@@ -579,6 +654,16 @@ public class Av1CoefficientsEntropyTests
         {
             Assert.Equal(24, partitions.Left[index].Left);
         }
+
+        for (int index = 0; index < transforms.Top.Length; index++)
+        {
+            Assert.Equal(index < 16 ? 0 : 32, transforms.Top[index]);
+        }
+
+        for (int index = 0; index < transforms.Left.Length; index++)
+        {
+            Assert.Equal(32, transforms.Left[index]);
+        }
     }
 
     [Fact]