Implement mode info indexing in frame buffer

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

@@ -9,19 +9,22 @@ internal class Av1BlockModeInfo
 {
     private int[] paletteSize;
 
-    public Av1BlockModeInfo(int numPlanes, Av1BlockSize blockSize, Point position)
+    public Av1BlockModeInfo(int numPlanes, Av1BlockSize blockSize, Point positionInSuperblock)
     {
         this.BlockSize = blockSize;
-        this.PositionInSuperblock = position;
+        this.PositionInSuperblock = positionInSuperblock;
         this.AngleDelta = new int[numPlanes - 1];
         this.paletteSize = new int[numPlanes - 1];
         this.FilterIntraModeInfo = new();
         this.FirstTransformLocation = new int[numPlanes - 1];
-        this.TusCount = new int[numPlanes - 1];
+        this.TransformUnitsCount = new int[numPlanes - 1];
     }
 
     public Av1BlockSize BlockSize { get; }
 
+    /// <summary>
+    /// Gets or sets the <see cref="Av1PredictionMode"/> for the luminance channel.
+    /// </summary>
     public Av1PredictionMode YMode { get; set; }
 
     public bool Skip { get; set; }
@@ -32,6 +35,9 @@ internal class Av1BlockModeInfo
 
     public int SegmentId { get; set; }
 
+    /// <summary>
+    /// Gets or sets the <see cref="Av1PredictionMode"/> for the chroma channels.
+    /// </summary>
     public Av1PredictionMode UvMode { get; set; }
 
     public bool UseUltraBlockCopy { get; set; }
@@ -42,13 +48,19 @@ internal class Av1BlockModeInfo
 
     public int[] AngleDelta { get; set; }
 
-    public Point PositionInSuperblock { get; set; }
+    /// <summary>
+    /// Gets the position relative to the Superblock, counted in mode info (4x4 pixels).
+    /// </summary>
+    public Point PositionInSuperblock { get; }
 
     public Av1IntraFilterModeInfo FilterIntraModeInfo { get; internal set; }
 
+    /// <summary>
+    /// Gets the index of the first <see cref="Av1TransformInfo"/> of this Mode Info in the <see cref="Av1FrameBuffer"/>.
+    /// </summary>
     public int[] FirstTransformLocation { get; }
 
-    public int[] TusCount { get; internal set; }
+    public int[] TransformUnitsCount { get; internal set; }
 
     public int GetPaletteSize(Av1PlaneType planeType) => this.paletteSize[(int)planeType];
 

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

@@ -1,13 +1,14 @@
 // Copyright (c) Six Labors.
 // Licensed under the Six Labors Split License.
 
+using System;
 using System.Reflection.Metadata.Ecma335;
 using SixLabors.ImageSharp.Formats.Heif.Av1.OpenBitstreamUnit;
 using SixLabors.ImageSharp.Formats.Heif.Av1.Transform;
 
 namespace SixLabors.ImageSharp.Formats.Heif.Av1.Symbol;
 
-internal class Av1FrameBuffer
+internal partial class Av1FrameBuffer
 {
     // Number of Coefficients in a single ModeInfo 4x4 block of pixels (1 length + 4 x 4).
     public const int CoefficientCountPerModeInfo = 1 + 16;
@@ -22,6 +23,7 @@ internal class Av1FrameBuffer
     private readonly int subsamplingFactor;
     private readonly Av1SuperblockInfo[] superblockInfos;
     private readonly Av1BlockModeInfo[] modeInfos;
+    private readonly Av1FrameModeInfoMap modeInfoMap;
     private readonly Av1TransformInfo[] transformInfosY;
     private readonly Av1TransformInfo[] transformInfosUv;
     private readonly int[] deltaQ;
@@ -46,30 +48,18 @@ internal class Av1FrameBuffer
         // Allocate the arrays.
         this.superblockInfos = new Av1SuperblockInfo[superblockCount];
         this.modeInfos = new Av1BlockModeInfo[superblockCount * this.modeInfoCountPerSuperblock];
+        this.modeInfoMap = new Av1FrameModeInfoMap(new Size(this.modeInfoCountPerSuperblock * this.superblockColumnCount, this.modeInfoCountPerSuperblock * this.superblockRowCount), superblockSizeLog2);
         this.transformInfosY = new Av1TransformInfo[superblockCount * this.modeInfoCountPerSuperblock];
         this.transformInfosUv = new Av1TransformInfo[2 * superblockCount * this.modeInfoCountPerSuperblock];
 
         // Initialize the arrays.
         int i = 0;
-        int j = 0;
-        int k = 0;
         for (int y = 0; y < this.superblockRowCount; y++)
         {
             for (int x = 0; x < this.superblockColumnCount; x++)
             {
                 Point point = new(x, y);
                 this.superblockInfos[i] = new(this, point);
-                for (int u = 0; u < this.modeInfoSizePerSuperblock; u++)
-                {
-                    for (int v = 0; v < this.modeInfoSizePerSuperblock; v++)
-                    {
-                        this.modeInfos[k] = new Av1BlockModeInfo(numPlanes, Av1BlockSize.Block4x4, new Point(u, v));
-                        k++;
-                    }
-
-                    j++;
-                }
-
                 i++;
             }
         }
@@ -110,10 +100,9 @@ internal class Av1FrameBuffer
 
     public Av1BlockModeInfo GetModeInfo(Point superblockIndex, Point modeInfoIndex)
     {
-        Span<Av1BlockModeInfo> span = this.modeInfos;
-        int superblock = (superblockIndex.Y * this.superblockColumnCount) + superblockIndex.X;
-        int modeInfo = (modeInfoIndex.Y * this.modeInfoSizePerSuperblock) + modeInfoIndex.X;
-        return span[(superblock * this.modeInfoCountPerSuperblock) + modeInfo];
+        Point location = this.GetModeInfoPosition(superblockIndex, modeInfoIndex);
+        int index = this.modeInfoMap[location];
+        return this.modeInfos[index];
     }
 
     public ref Av1TransformInfo GetTransformY(int index)
@@ -216,4 +205,17 @@ internal class Av1FrameBuffer
             1 or 2 => this.GetTransformUv(transformInfoIndex),
             _ => null,
         };
+
+    public void UpdateModeInfo(Av1BlockModeInfo modeInfo, Av1SuperblockInfo superblockInfo)
+    {
+        this.modeInfos[this.modeInfoMap.NextIndex] = modeInfo;
+        this.modeInfoMap.Update(this.GetModeInfoPosition(superblockInfo.Position, modeInfo.PositionInSuperblock), modeInfo.BlockSize);
+    }
+
+    private Point GetModeInfoPosition(Point superblockPosition, Point positionInSuperblock)
+    {
+        int x = (superblockPosition.X * this.modeInfoCountPerSuperblock) + positionInSuperblock.X;
+        int y = (superblockPosition.Y * this.modeInfoCountPerSuperblock) + positionInSuperblock.Y;
+        return new Point(x, y);
+    }
 }

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

@@ -0,0 +1,63 @@
+// Copyright (c) Six Labors.
+// Licensed under the Six Labors Split License.
+
+namespace SixLabors.ImageSharp.Formats.Heif.Av1.Symbol;
+
+internal partial class Av1FrameBuffer
+{
+    /// <summary>
+    /// Mapping of <see cref="Av1BlockModeInfo"/> instances, from position to index into the <see cref="Av1FrameBuffer"/>.
+    /// </summary>
+    /// <remarks>
+    /// For a visual representation of how this map looks in practice, see <seealso href="https://gitlab.com/AOMediaCodec/SVT-AV1/-/blob/v2.1.0/Docs/svt-av1-decoder-design.md?ref_type=tags#blockmodeinfo"/>
+    /// </remarks>
+    public class Av1FrameModeInfoMap
+    {
+        private readonly int[] offsets;
+        private Size alignedModeInfoCount;
+
+        public Av1FrameModeInfoMap(Size modeInfoCount, int superblockSizeLog2)
+        {
+            this.alignedModeInfoCount = new Size(
+                modeInfoCount.Width * (1 << (superblockSizeLog2 - Av1Constants.ModeInfoSizeLog2)),
+                modeInfoCount.Height * (1 << (superblockSizeLog2 - Av1Constants.ModeInfoSizeLog2)));
+            this.NextIndex = 0;
+            this.offsets = new int[this.alignedModeInfoCount.Width * this.alignedModeInfoCount.Height];
+        }
+
+        /// <summary>
+        /// Gets the next index to use.
+        /// </summary>
+        public int NextIndex { get; private set; }
+
+        /// <summary>
+        /// Gets the mapped index for the given location.
+        /// </summary>
+        public int this[Point location]
+        {
+            get
+            {
+                int index = (location.Y * this.alignedModeInfoCount.Width) + location.X;
+                return this.offsets[index];
+            }
+        }
+
+        public void Update(Point modeInfoLocation, Av1BlockSize blockSize)
+        {
+            // Equivalent in SVT-Av1: EbDecNbr.c svt_aom_update_block_nbrs
+            int bw4 = blockSize.Get4x4WideCount();
+            int bh4 = blockSize.Get4x4HighCount();
+            DebugGuard.MustBeGreaterThanOrEqualTo(modeInfoLocation.Y, 0, nameof(modeInfoLocation));
+            DebugGuard.MustBeLessThanOrEqualTo(modeInfoLocation.Y + bh4, this.alignedModeInfoCount.Height, nameof(modeInfoLocation));
+            DebugGuard.MustBeGreaterThanOrEqualTo(modeInfoLocation.X, 0, nameof(modeInfoLocation));
+            DebugGuard.MustBeLessThanOrEqualTo(modeInfoLocation.X + bw4, this.alignedModeInfoCount.Width, nameof(modeInfoLocation));
+            /* Update 4x4 nbr offset map */
+            for (int i = modeInfoLocation.Y; i < modeInfoLocation.Y + bh4; i++)
+            {
+                Array.Fill(this.offsets, this.NextIndex, (i * this.alignedModeInfoCount.Width) + modeInfoLocation.X, bw4);
+            }
+
+            this.NextIndex++;
+        }
+    }
+}

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

@@ -13,11 +13,14 @@ internal class Av1SuperblockInfo
         this.frameBuffer = frameBuffer;
     }
 
+    /// <summary>
+    /// Gets the position of this superblock inside the tile, counted in superblocks.
+    /// </summary>
     public Point Position { get; }
 
     public ref int SuperblockDeltaQ => ref this.frameBuffer.GetDeltaQuantizationIndex(this.Position);
 
-    public Av1BlockModeInfo SuperblockModeInfo => this.GetModeInfo(default);
+    public Av1BlockModeInfo SuperblockModeInfo => this.GetModeInfo(new Point(0, 0));
 
     public Span<int> CoefficientsY => this.frameBuffer.GetCoefficientsY(this.Position);
 

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

@@ -115,12 +115,12 @@ internal class Av1TileDecoder : IAv1TileDecoder
                 Av1SuperblockInfo superblockInfo = this.FrameBuffer.GetSuperblock(superblockPosition);
 
                 // this.ClearBlockDecodedFlags(modeInfoLocation, superBlock4x4Size);
-                Point modeInfoLocation = new(column, row);
+                Point modeInfoPosition = new(column, row);
                 this.FrameBuffer.ClearCdef(superblockPosition);
                 this.firstTransformOffset[0] = 0;
                 this.firstTransformOffset[1] = 0;
-                this.ReadLoopRestoration(modeInfoLocation, superBlockSize);
-                this.ParsePartition(ref reader, modeInfoLocation, superBlockSize, superblockInfo, tileInfo);
+                this.ReadLoopRestoration(modeInfoPosition, superBlockSize);
+                this.ParsePartition(ref reader, modeInfoPosition, superBlockSize, superblockInfo, tileInfo);
             }
         }
     }
@@ -291,7 +291,9 @@ internal class Av1TileDecoder : IAv1TileDecoder
         int block4x4Width = blockSize.Get4x4WideCount();
         int block4x4Height = blockSize.Get4x4HighCount();
         int planesCount = this.SequenceHeader.ColorConfig.ChannelCount;
-        Av1BlockModeInfo blockModeInfo = superblockInfo.GetModeInfo(modeInfoLocation);
+        Point superblockLocation = superblockInfo.Position * this.SequenceHeader.SuperblockModeInfoSize;
+        Point locationInSuperblock = new Point(modeInfoLocation.X - superblockLocation.X, modeInfoLocation.Y - superblockLocation.Y);
+        Av1BlockModeInfo blockModeInfo = new(planesCount, blockSize, locationInSuperblock);
         blockModeInfo.PartitionType = partitionType;
         blockModeInfo.FirstTransformLocation[0] = this.firstTransformOffset[0];
         blockModeInfo.FirstTransformLocation[1] = this.firstTransformOffset[1];
@@ -332,6 +334,9 @@ internal class Av1TileDecoder : IAv1TileDecoder
         }
 
         this.Residual(ref reader, partitionInfo, superblockInfo, tileInfo, blockSize);
+
+        // Update the Frame buffer for this ModeInfo.
+        this.FrameBuffer.UpdateModeInfo(blockModeInfo, superblockInfo);
     }
 
     private void ResetSkipContext(Av1PartitionInfo partitionInfo)
@@ -367,7 +372,7 @@ internal class Av1TileDecoder : IAv1TileDecoder
         bool isLossless = this.FrameInfo.LosslessArray[partitionInfo.ModeInfo.SegmentId];
         bool isLosslessBlock = isLossless && (blockSize >= Av1BlockSize.Block64x64) && (blockSize <= Av1BlockSize.Block128x128);
         int subSampling = (this.SequenceHeader.ColorConfig.SubSamplingX ? 1 : 0) + (this.SequenceHeader.ColorConfig.SubSamplingY ? 1 : 0);
-        int chromaTusCount = isLosslessBlock ? ((maxBlocksWide * maxBlocksHigh) >> subSampling) : partitionInfo.ModeInfo.TusCount[(int)Av1PlaneType.Uv];
+        int chromaTusCount = isLosslessBlock ? ((maxBlocksWide * maxBlocksHigh) >> subSampling) : partitionInfo.ModeInfo.TransformUnitsCount[(int)Av1PlaneType.Uv];
         int[] transformInfoIndices = new int[3];
         transformInfoIndices[0] = superblockInfo.TransformInfoIndexY + partitionInfo.ModeInfo.FirstTransformLocation[(int)Av1PlaneType.Y];
         transformInfoIndices[1] = superblockInfo.TransformInfoIndexUv + partitionInfo.ModeInfo.FirstTransformLocation[(int)Av1PlaneType.Uv];
@@ -400,7 +405,7 @@ internal class Av1TileDecoder : IAv1TileDecoder
                     }
                     else
                     {
-                        totalTusCount = partitionInfo.ModeInfo.TusCount[Math.Min(1, plane)];
+                        totalTusCount = partitionInfo.ModeInfo.TransformUnitsCount[Math.Min(1, plane)];
                         tusCount = this.tusCount[plane][forceSplitCount];
 
                         DebugGuard.IsFalse(totalTusCount == 0, nameof(totalTusCount), string.Empty);
@@ -1358,8 +1363,8 @@ internal class Av1TileDecoder : IAv1TileDecoder
             }
         }
 
-        partitionInfo.ModeInfo.TusCount[(int)Av1PlaneType.Y] = totalLumaTusCount;
-        partitionInfo.ModeInfo.TusCount[(int)Av1PlaneType.Uv] = totalChromaTusCount;
+        partitionInfo.ModeInfo.TransformUnitsCount[(int)Av1PlaneType.Y] = totalLumaTusCount;
+        partitionInfo.ModeInfo.TransformUnitsCount[(int)Av1PlaneType.Uv] = totalChromaTusCount;
 
         this.firstTransformOffset[(int)Av1PlaneType.Y] += totalLumaTusCount;
         this.firstTransformOffset[(int)Av1PlaneType.Uv] += totalChromaTusCount << 1;