Correct AV1 transform coefficient storage

HEIF_IMPLEMENTATION_PLAN.md

@@ -59,6 +59,7 @@ This snapshot pins or classifies the available references and failures; it does
 | `Av1YuvConverter.ConvertToRgb`, `ConvertFromRgb`, scalar row conversion, and chroma reconstruction | H.273 formulas 20-31 and the identity, YCgCo, and non-constant-luminance matrix formulas; AV1 section 6.4.2 chroma sample positions | libavif `src/reformat.c` and `src/colr.c` at `092276ce89098ead06db80975173191e5fee1826`; libaom `aom/aom_image.h` at `03087864cf4bea6abb0d28f95cf7843511413d8f` | Scalar behavioral oracle for 8-bit full/limited-range conversion. Decode covers monochrome, YUV 4:2:0, 4:2:2, and 4:4:4 with AV1 chroma sample positioning; encode remains YUV 4:4:4 at this snapshot. Later high-bit-depth and SIMD paths must match it. |
 | `Av1FrameBuffer` high-bit-depth sample layout and `Av1YuvConverter` 10/12-bit output conversion | AV1 section 6.4.1 bit depth and H.273 sample-range scaling | libaom `aom_scale/yv12config.h`, `av1/common/idct.c`, and `av1/common/reconintra.c` at `03087864cf4bea6abb0d28f95cf7843511413d8f`; libavif `src/avif.c` and `src/reformat.c` at `092276ce89098ead06db80975173191e5fee1826` | Establish two-byte native sample storage with sample-unit strides for 10/12-bit reconstruction and use the same scalar color model at every supported bit depth. |
 | `Av1PredictionDecoder`, `Av1HighBitDepthPredictor`, `Av1ChromaFromLumaContext`, `Av1PartitionInfo`, and the scalar DC, directional, Paeth, smooth, filter-intra, and chroma-from-luma predictors | AV1 sections 7.11.2 and 7.11.2.3 intra prediction | libaom `aom_dsp/intrapred.c`, `av1/common/reconintra.c`, `av1/common/av1_common_int.h`, `av1/common/cfl.c`, and `av1/common/cfl.h` at `03087864cf4bea6abb0d28f95cf7843511413d8f` | Behavioral oracle for luma/chroma mode-neighbor addressing, directional upsampling, Paeth selection, smooth normalization, filter-intra taps, high-bit-depth clipping, chroma-from-luma storage/subsampling, and chroma-from-luma row strides. Existing managed scalar tables and predictors remain the implementation base. The WIP rectangular byte-pipeline smooth digest expectations encode width/height-swapped weights and must be replaced only from an independently generated oracle, not regenerated from this implementation. |
+| `Av1FrameInfo`, `Av1TileReader`, and `Av1BlockDecoder` transform/coefficient storage | AV1 section 5.11.39 coefficient syntax and section 7.11.2 reconstruction | libaom `av1/decoder/decodetxb.c` and `av1/decoder/decoder.h` at `03087864cf4bea6abb0d28f95cf7843511413d8f` | Preserve separate luma and chroma transform coefficients at monotonically advancing per-plane offsets within each superblock so reconstruction consumes the same transform-block order produced by tile parsing. |
 | `Av1Inverse2dTransformer` and `Av1InverseTransformerFactory` | AV1 section 7.11.2 inverse transform and reconstruction | libaom `av1/common/av1_inv_txfm1d.c`, `av1/common/av1_inv_txfm2d.c`, and `av1/common/idct.c` at `03087864cf4bea6abb0d28f95cf7843511413d8f` | Scalar transform oracle for coefficient-row traversal, intermediate layout, stage ranges, clipping, and high-bit-depth sample addition. The managed 16-bit overload is also used as a parity oracle for the byte overload. |
 
 This table is intentionally incomplete. Add a row before each additional AV1 or HEVC algorithm is ported or materially reshaped.

src/ImageSharp/Formats/Heif/Av1/Pipeline/Av1FrameDecoder.cs

@@ -26,7 +26,7 @@ internal class Av1FrameDecoder : IAv1FrameDecoder
         this.frameBuffer = frameBuffer;
         this.inverseQuantizer = new(sequenceHeader, frameHeader);
         this.deQuants = new(sequenceHeader, frameHeader);
-        this.blockDecoder = new(this.sequenceHeader, this.frameHeader, this.frameInfo, this.frameBuffer);
+        this.blockDecoder = new(this.sequenceHeader, this.frameHeader, this.frameBuffer);
     }
 
     public void DecodeFrame()

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

@@ -60,12 +60,6 @@ internal partial class Av1FrameInfo
             {
                 Point point = new(x, y);
                 this.superblockInfos[i] = new(this, point);
-                for (int j = 0; j < this.modeInfoCountPerSuperblock; j++)
-                {
-                    this.transformInfosY[j] = new Av1TransformInfo();
-                    this.transformInfosUv[j] = new Av1TransformInfo();
-                }
-
                 i++;
             }
         }
@@ -76,9 +70,11 @@ internal partial class Av1FrameInfo
         // Factor: 444 => 0, 422 => 1, 420 => 2.
         this.subsamplingFactor = (subX && subY) ? 2 : (subX && !subY) ? 1 : (!subX && !subY) ? 0 : -1;
         Guard.IsFalse(this.subsamplingFactor == -1, nameof(this.subsamplingFactor), "Invalid combination of subsampling.");
-        this.coefficientsY = new int[superblockCount * this.modeInfoCountPerSuperblock * CoefficientCountPerModeInfo];
+        int lumaCoefficientCountPerSuperblock = this.modeInfoCountPerSuperblock * CoefficientCountPerModeInfo;
-        this.coefficientsU = new int[(this.modeInfoCountPerSuperblock * CoefficientCountPerModeInfo) >> this.subsamplingFactor];
+        int chromaCoefficientCountPerSuperblock = lumaCoefficientCountPerSuperblock >> this.subsamplingFactor;
-        this.coefficientsV = new int[(this.modeInfoCountPerSuperblock * CoefficientCountPerModeInfo) >> this.subsamplingFactor];
+        this.coefficientsY = new int[superblockCount * lumaCoefficientCountPerSuperblock];
+        this.coefficientsU = new int[superblockCount * chromaCoefficientCountPerSuperblock];
+        this.coefficientsV = new int[superblockCount * chromaCoefficientCountPerSuperblock];
         this.deltaQ = new int[superblockCount];
 
         // Superblock size: 128x128 has sizelog2 = 7, 64x64 = 6. Factor should be 128x128 => 4 and 64x64 => 1.
@@ -150,37 +146,31 @@ internal partial class Av1FrameInfo
     {
         Span<Av1TransformInfo> span = this.transformInfosUv;
         int offset = (((index.Y * this.superblockColumnCount) + index.X) * this.modeInfoCountPerSuperblock) << 1;
-        return span.Slice(offset, this.modeInfoCountPerSuperblock);
+        return span.Slice(offset, this.modeInfoCountPerSuperblock << 1);
     }
 
-    public Span<int> GetCoefficients(int plane) =>
-        plane switch
-        {
-            0 => (Span<int>)this.coefficientsY,
-            1 => (Span<int>)this.coefficientsY,
-            2 => (Span<int>)this.coefficientsY,
-            _ => null,
-        };
-
     public Span<int> GetCoefficientsY(Point index)
     {
         Span<int> span = this.coefficientsY;
-        int i = ((index.Y * this.modeInfoCountPerSuperblock) + index.X) * CoefficientCountPerModeInfo;
+        int count = this.modeInfoCountPerSuperblock * CoefficientCountPerModeInfo;
-        return span.Slice(i, CoefficientCountPerModeInfo);
+        int superblock = (index.Y * this.superblockColumnCount) + index.X;
+        return span.Slice(superblock * count, count);
     }
 
     public Span<int> GetCoefficientsU(Point index)
     {
         Span<int> span = this.coefficientsU;
-        int i = ((index.Y * this.modeInfoCountPerSuperblock) + index.X) * CoefficientCountPerModeInfo;
+        int count = (this.modeInfoCountPerSuperblock * CoefficientCountPerModeInfo) >> this.subsamplingFactor;
-        return span.Slice(i >> this.subsamplingFactor, CoefficientCountPerModeInfo);
+        int superblock = (index.Y * this.superblockColumnCount) + index.X;
+        return span.Slice(superblock * count, count);
     }
 
     public Span<int> GetCoefficientsV(Point index)
     {
         Span<int> span = this.coefficientsV;
-        int i = ((index.Y * this.modeInfoCountPerSuperblock) + index.X) * CoefficientCountPerModeInfo;
+        int count = (this.modeInfoCountPerSuperblock * CoefficientCountPerModeInfo) >> this.subsamplingFactor;
-        return span.Slice(i >> this.subsamplingFactor, CoefficientCountPerModeInfo);
+        int superblock = (index.Y * this.superblockColumnCount) + index.X;
+        return span.Slice(superblock * count, count);
     }
 
     public ref int GetDeltaQuantizationIndex(Point index)

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

@@ -120,16 +120,16 @@ internal class Av1PartitionInfo
         this.HeightInPixels[2] = Math.Max(1, bh4 >> subY) * modeInfoSize;
     }
 
-    public void PopulateModeInfoNeighbors(Av1FrameInfo frameInfo, ObuColorConfig colorConfig)
+    public void PopulateModeInfoNeighbors(ObuColorConfig colorConfig)
     {
         if (this.AvailableAbove)
         {
-            this.AboveModeInfo = frameInfo.GetModeInfoAt(new Point(this.ColumnIndex, this.RowIndex - 1));
+            this.AboveModeInfo = this.SuperblockInfo.GetModeInfoAt(new Point(this.ColumnIndex, this.RowIndex - 1));
         }
 
         if (this.AvailableLeft)
         {
-            this.LeftModeInfo = frameInfo.GetModeInfoAt(new Point(this.ColumnIndex - 1, this.RowIndex));
+            this.LeftModeInfo = this.SuperblockInfo.GetModeInfoAt(new Point(this.ColumnIndex - 1, this.RowIndex));
         }
 
         if (!this.IsChroma)
@@ -145,12 +145,12 @@ internal class Av1PartitionInfo
         // Chroma neighbors refer to the bottom-right luma mode covered by each adjacent chroma block.
         if (this.AvailableAboveForChroma)
         {
-            this.AboveModeInfoForChroma = frameInfo.GetModeInfoAt(new Point(chromaBaseColumn + subX, chromaBaseRow - 1));
+            this.AboveModeInfoForChroma = this.SuperblockInfo.GetModeInfoAt(new Point(chromaBaseColumn + subX, chromaBaseRow - 1));
         }
 
         if (this.AvailableLeftForChroma)
         {
-            this.LeftModeInfoForChroma = frameInfo.GetModeInfoAt(new Point(chromaBaseColumn - 1, chromaBaseRow + subY));
+            this.LeftModeInfoForChroma = this.SuperblockInfo.GetModeInfoAt(new Point(chromaBaseColumn - 1, chromaBaseRow + subY));
         }
     }
 

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

@@ -47,6 +47,8 @@ internal class Av1SuperblockInfo
 
     public Span<Av1TransformInfo> GetTransformInfoUv() => this.frameInfo.GetSuperblockTransformUv(this.Position);
 
+    public Span<Av1TransformInfo> GetTransformInfo(int plane) => this.frameInfo.GetSuperblockTransform(plane, this.Position);
+
     /// <summary>
     /// Gets the mode information records parsed for this superblock in bitstream order.
     /// </summary>
@@ -54,6 +56,8 @@ internal class Av1SuperblockInfo
 
     public Av1BlockModeInfo GetModeInfo(Point index) => this.frameInfo.GetModeInfo(this.Position, index);
 
+    public Av1BlockModeInfo GetModeInfoAt(Point index) => this.frameInfo.GetModeInfoAt(index);
+
     public Span<int> GetCoefficients(Av1Plane plane) => plane switch
     {
         Av1Plane.Y => this.CoefficientsY,

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

@@ -125,6 +125,7 @@ internal class Av1TileReader : IAv1TileReader
                 this.FrameInfo.ClearCdef(superblockPosition);
                 this.firstTransformOffset[0] = 0;
                 this.firstTransformOffset[1] = 0;
+                this.coefficientIndex.AsSpan().Clear();
                 this.ReadLoopRestoration(modeInfoPosition, superBlockSize);
                 this.ParsePartition(ref reader, modeInfoPosition, superBlockSize, superblockInfo, tileInfo);
 
@@ -315,7 +316,7 @@ internal class Av1TileReader : IAv1TileReader
             }
         }
 
-        partitionInfo.PopulateModeInfoNeighbors(this.FrameInfo, this.SequenceHeader.ColorConfig);
+        partitionInfo.PopulateModeInfoNeighbors(this.SequenceHeader.ColorConfig);
 
         this.ReadModeInfo(ref reader, partitionInfo);
         ReadPaletteTokens(ref reader, partitionInfo);
@@ -438,7 +439,8 @@ internal class Av1TileReader : IAv1TileReader
 
                         if (!partitionInfo.ModeInfo.Skip)
                         {
-                            endOfBlock = this.ParseTransformBlock(ref reader, partitionInfo, coefficientIndex, transformInfo, plane, blockColumn, blockRow, startX, startY, transformInfo.Size, subX != 0, subY != 0);
+                            Span<int> coefficientBuffer = superblockInfo.GetCoefficients((Av1Plane)plane)[coefficientIndex..];
+                            endOfBlock = this.ParseTransformBlock(ref reader, partitionInfo, coefficientBuffer, transformInfo, plane, blockColumn, blockRow, startX, startY, transformInfo.Size, subX != 0, subY != 0);
                         }
 
                         if (endOfBlock != 0)
@@ -480,7 +482,7 @@ internal class Av1TileReader : IAv1TileReader
     private int ParseTransformBlock(
         ref Av1SymbolDecoder reader,
         Av1PartitionInfo partitionInfo,
-        int coefficientIndex,
+        Span<int> coefficientBuffer,
         Av1TransformInfo transformInfo,
         int plane,
         int blockColumn,
@@ -509,7 +511,7 @@ internal class Av1TileReader : IAv1TileReader
         }
 
         Av1TransformBlockContext transformBlockContext = this.GetTransformBlockContext(transformSize, plane, planeBlockSize, transformBlockUnitHighCount, transformBlockUnitWideCount, startY, startX);
-        endOfBlock = this.ParseCoefficients(ref reader, partitionInfo, startY, startX, blockRow, blockColumn, plane, transformBlockContext, transformSize, coefficientIndex, transformInfo);
+        endOfBlock = this.ParseCoefficients(ref reader, partitionInfo, startY, startX, blockRow, blockColumn, plane, transformBlockContext, transformSize, transformInfo, coefficientBuffer);
 
         return endOfBlock;
     }
@@ -520,9 +522,8 @@ internal class Av1TileReader : IAv1TileReader
     /// <remarks>
     /// The implementation is taken from SVT-AV1 library, which deviates from the code flow in the specification.
     /// </remarks>
-    private int ParseCoefficients(ref Av1SymbolDecoder reader, Av1PartitionInfo partitionInfo, int blockRow, int blockColumn, int aboveOffset, int leftOffset, int plane, Av1TransformBlockContext transformBlockContext, Av1TransformSize transformSize, int coefficientIndex, Av1TransformInfo transformInfo)
+    private int ParseCoefficients(ref Av1SymbolDecoder reader, Av1PartitionInfo partitionInfo, int blockRow, int blockColumn, int aboveOffset, int leftOffset, int plane, Av1TransformBlockContext transformBlockContext, Av1TransformSize transformSize, Av1TransformInfo transformInfo, Span<int> coefficientBuffer)
     {
-        Span<int> coefficientBuffer = this.FrameInfo.GetCoefficients(plane);
         int width = transformSize.GetWidth();
         int height = transformSize.GetHeight();
         Av1TransformSize transformSizeContext = Av1SymbolContextHelper.GetTransformSizeContext(transformSize);
@@ -895,7 +896,8 @@ internal class Av1TileReader : IAv1TileReader
             ref Av1TransformInfo infoV = ref chromaTransformInfo[transformInfoUvIndex];
             for (int i = 0; i < totalChromaTransformUnitCount; i++)
             {
-                infoV = originalInfo;
+                // U and V share transform geometry, but their entropy state and coefficients remain independent.
+                infoV = new Av1TransformInfo(originalInfo);
                 originalInfo = ref Unsafe.Add(ref originalInfo, 1);
                 infoV = ref Unsafe.Add(ref infoV, 1);
             }

src/ImageSharp/Formats/Heif/Av1/Transform/Av1BlockDecoder.cs

@@ -16,8 +16,6 @@ internal class Av1BlockDecoder
 
     private readonly ObuFrameHeader frameHeader;
 
-    private readonly Av1FrameInfo frameInfo;
-
     private readonly Av1FrameBuffer<byte> frameBuffer;
 
     private readonly bool isLoopFilterEnabled;
@@ -26,11 +24,10 @@ internal class Av1BlockDecoder
 
     private readonly Av1ChromaFromLumaContext chromaFromLumaContext;
 
-    public Av1BlockDecoder(ObuSequenceHeader sequenceHeader, ObuFrameHeader frameHeader, Av1FrameInfo frameInfo, Av1FrameBuffer<byte> frameBuffer)
+    public Av1BlockDecoder(ObuSequenceHeader sequenceHeader, ObuFrameHeader frameHeader, Av1FrameBuffer<byte> frameBuffer)
     {
         this.sequenceHeader = sequenceHeader;
         this.frameHeader = frameHeader;
-        this.frameInfo = frameInfo;
         this.frameBuffer = frameBuffer;
         int ySize = (1 << this.sequenceHeader.SuperblockSizeLog2) * (1 << this.sequenceHeader.SuperblockSizeLog2);
         int inverseQuantizationSize = ySize +
@@ -83,7 +80,7 @@ internal class Av1BlockDecoder
             }
         }
 
-        partitionInfo.PopulateModeInfoNeighbors(this.frameInfo, colorConfig);
+        partitionInfo.PopulateModeInfoNeighbors(colorConfig);
 
         int maxBlocksWide = partitionInfo.GetMaxBlockWide(blockSize, false);
         int maxBlocksHigh = partitionInfo.GetMaxBlockHigh(blockSize, false);
@@ -111,11 +108,11 @@ internal class Av1BlockDecoder
             int transformInfoIndex = plane switch
             {
                 2 => superblockInfo.TransformInfoIndexUv + modeInfo.FirstTransformLocation[plane - 1] + chromaTransformUnitCount,
-                1 => superblockInfo.TransformInfoIndexY + modeInfo.FirstTransformLocation[plane],
+                1 => superblockInfo.TransformInfoIndexUv + modeInfo.FirstTransformLocation[plane],
                 0 => superblockInfo.TransformInfoIndexY + modeInfo.FirstTransformLocation[plane],
                 _ => throw new InvalidImageContentException("Maximum of 3 color planes")
             };
-            Span<Av1TransformInfo> transformInfo = this.frameInfo.GetSuperblockTransform(plane, superblockInfo.Position)[transformInfoIndex..];
+            Span<Av1TransformInfo> transformInfo = superblockInfo.GetTransformInfo(plane)[transformInfoIndex..];
             Guard.NotNull(transformInfo[0]);
 
             if (isLosslessBlock)

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

@@ -5,6 +5,7 @@ using SixLabors.ImageSharp.Formats.Heif.Av1;
 using SixLabors.ImageSharp.Formats.Heif.Av1.OpenBitstreamUnit;
 using SixLabors.ImageSharp.Formats.Heif.Av1.Pipeline;
 using SixLabors.ImageSharp.Formats.Heif.Av1.Tiling;
+using SixLabors.ImageSharp.Formats.Heif.Av1.Transform;
 using SixLabors.ImageSharp.Memory;
 using SixLabors.ImageSharp.PixelFormats;
 
@@ -166,6 +167,75 @@ public class Av1TilingTests
         Assert.True(parsedModeInfoCount > 16);
     }
 
+    [Fact]
+    public void ParsedCoefficientsRemainAvailablePerSuperblockAndPlane()
+    {
+        string filePath = Path.Combine(TestEnvironment.InputImagesDirectoryFullPath, TestImages.Heif.XnConvert);
+        byte[] content = File.ReadAllBytes(filePath);
+        const int dataOffset = 0x010E;
+        const int dataSize = 0x03CC;
+        const int tileOffset = 18;
+        Span<byte> headerSpan = content.AsSpan(dataOffset, dataSize);
+        Span<byte> tileSpan = content.AsSpan(tileOffset, dataSize - tileOffset);
+        Av1BitStreamReader bitStreamReader = new(headerSpan);
+        IAv1TileReader stub = new Av1TileDecoderStub();
+        ObuReader obuReader = new();
+        obuReader.ReadAll(ref bitStreamReader, dataSize, () => stub);
+        Av1TileReader tileReader = new(Configuration.Default, obuReader.SequenceHeader, obuReader.FrameHeader);
+
+        tileReader.ReadTile(tileSpan, 0);
+
+        int codedTransformCount = 0;
+        int superblockSize = obuReader.SequenceHeader.SuperblockModeInfoSize;
+        for (int row = 0; row < obuReader.FrameHeader.ModeInfoRowCount; row += superblockSize)
+        {
+            for (int column = 0; column < obuReader.FrameHeader.ModeInfoColumnCount; column += superblockSize)
+            {
+                Point superblockPosition = new(column / superblockSize, row / superblockSize);
+                Av1SuperblockInfo superblockInfo = tileReader.FrameInfo.GetSuperblock(superblockPosition);
+                int[] coefficientIndices = new int[Av1Constants.MaxPlanes];
+
+                foreach (Av1BlockModeInfo modeInfo in superblockInfo.GetModeInfos())
+                {
+                    Point modeInfoPosition = new(column + modeInfo.PositionInSuperblock.X, row + modeInfo.PositionInSuperblock.Y);
+                    bool hasChroma = Av1TileReader.HasChroma(obuReader.SequenceHeader, modeInfoPosition, modeInfo.BlockSize);
+
+                    for (int plane = 0; plane < obuReader.SequenceHeader.ColorConfig.PlaneCount; plane++)
+                    {
+                        if (plane != 0 && !hasChroma)
+                        {
+                            continue;
+                        }
+
+                        int transformUnitCount = modeInfo.TransformUnitsCount[Math.Min(plane, 1)];
+                        int transformInfoIndex = modeInfo.FirstTransformLocation[Math.Min(plane, 1)];
+                        if (plane == (int)Av1Plane.V)
+                        {
+                            transformInfoIndex += transformUnitCount;
+                        }
+
+                        Span<Av1TransformInfo> transformInfos = superblockInfo.GetTransformInfo(plane)[transformInfoIndex..];
+                        Span<int> coefficients = superblockInfo.GetCoefficients((Av1Plane)plane);
+                        for (int i = 0; i < transformUnitCount; i++)
+                        {
+                            if (!transformInfos[i].CodeBlockFlag)
+                            {
+                                continue;
+                            }
+
+                            int endOfBlock = coefficients[coefficientIndices[plane]];
+                            Assert.InRange(endOfBlock, 1, transformInfos[i].Size.GetWidth() * transformInfos[i].Size.GetHeight());
+                            coefficientIndices[plane] += endOfBlock + 1;
+                            codedTransformCount++;
+                        }
+                    }
+                }
+            }
+        }
+
+        Assert.True(codedTransformCount > 3);
+    }
+
     [Theory]
     [InlineData(TestImages.Heif.XnConvert, 0x010E, 0x03CC, 18, 16)]
     [InlineData(TestImages.Heif.Orange4x4, 0x010E, 0x001d, 21, 1)]