Add coefficients unit test for multiple transform types

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

@@ -8,16 +8,6 @@ namespace SixLabors.ImageSharp.Formats.Heif.Av1.Entropy;
 
 internal static class Av1NzMap
 {
-    private static readonly int[] ClipMax3 = [
-        0, 1, 2, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3,
-        3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3,
-        3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3,
-        3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3,
-        3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3,
-        3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3,
-        3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3
-    ];
-
     // SIG_COEF_CONTEXTS_2D = 26
     private const int NzMapContext0 = 26;
     private const int NzMapContext5 = NzMapContext0 + 5;
@@ -351,35 +341,35 @@ internal static class Av1NzMap
         Span<byte> row2 = levels.GetRow(position.Y + 2)[position.X..];
 
         // Note: AOMMIN(level, 3) is useless for decoder since level < 3.
-        mag = ClipMax3[row0[1]]; // { 0, 1 }
-        mag += ClipMax3[row1[0]]; // { 1, 0 }
+        mag = ClipMax3(row0[1]); // { 0, 1 }
+        mag += ClipMax3(row1[0]); // { 1, 0 }
 
         switch (transformClass)
         {
             case Av1TransformClass.Class2D:
-                mag += ClipMax3[row1[1]]; // { 1, 1 }
-                mag += ClipMax3[row0[2]]; // { 0, 2 }
-                mag += ClipMax3[row2[0]]; // { 2, 0 }
+                mag += ClipMax3(row1[1]); // { 1, 1 }
+                mag += ClipMax3(row0[2]); // { 0, 2 }
+                mag += ClipMax3(row2[0]); // { 2, 0 }
                 break;
 
             case Av1TransformClass.ClassVertical:
                 Span<byte> row3 = levels.GetRow(position.Y + 3)[position.X..];
                 Span<byte> row4 = levels.GetRow(position.Y + 4)[position.X..];
-                mag += ClipMax3[row2[0]]; // { 2, 0 }
-                mag += ClipMax3[row3[0]]; // { 3, 0 }
-                mag += ClipMax3[row4[0]]; // { 4, 0 }
+                mag += ClipMax3(row2[0]); // { 2, 0 }
+                mag += ClipMax3(row3[0]); // { 3, 0 }
+                mag += ClipMax3(row4[0]); // { 4, 0 }
                 break;
             case Av1TransformClass.ClassHorizontal:
-                mag += ClipMax3[row0[2]]; // { 0, 2 }
-                mag += ClipMax3[row0[3]]; // { 0, 3 }
-                mag += ClipMax3[row0[4]]; // { 0, 4 }
+                mag += ClipMax3(row0[2]); // { 0, 2 }
+                mag += ClipMax3(row0[3]); // { 0, 3 }
+                mag += ClipMax3(row0[4]); // { 0, 4 }
                 break;
         }
 
         return mag;
     }
 
-    public static int GetNzMapContextFromStats(int stats, Av1LevelBuffer levels, Point position, Av1TransformSize transformSize, Av1TransformClass transformClass)
+    public static int GetNzMapContextFromStats(int stats, Point position, Av1TransformSize transformSize, Av1TransformClass transformClass)
     {
         // tx_class == 0(TX_CLASS_2D)
         if (position.Y == 0 && ((int)transformClass | position.X) == 0)
@@ -418,4 +408,6 @@ internal static class Av1NzMap
     public static int GetNzMapContext(Av1TransformSize transformSize, Point pos) => GetNzMapContext(transformSize, pos.X + (pos.Y * transformSize.GetWidth()));
 
     public static int GetNzMapContext(Av1TransformSize transformSize, int pos) => NzMapContextOffset[(int)transformSize][pos];
+
+    private static int ClipMax3(int value) => Math.Min(value, 3);
 }

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

@@ -123,6 +123,11 @@ internal static class Av1SymbolContextHelper
         return ctx + Av1NzMap.GetNzMapContext(transformSize, position);
     }
 
+    /// <summary>
+    /// Section 8.3.2 in the spec, under coeff_br. Optimized for end of block based
+    /// on the fact that {0, 1}, {1, 0}, {1, 1}, {0, 2} and {2, 0} will all be 0 in
+    /// the end of block case.
+    /// </summary>
     internal static int GetBaseRangeContextEndOfBlock(Point pos, Av1TransformClass transformClass)
     {
         if (pos.X == 0 && pos.Y == 0)
@@ -143,6 +148,7 @@ internal static class Av1SymbolContextHelper
     /// <summary>
     /// SVT: get_br_ctx
     /// </summary>
+    /// <remarks>Spec section 8.2.3, under 'coeff_br'.</remarks>
     internal static int GetBaseRangeContext(Av1LevelBuffer levels, Point position, Av1TransformClass transformClass)
     {
         Span<byte> row0 = levels.GetRow(position.Y);
@@ -224,7 +230,7 @@ internal static class Av1SymbolContextHelper
     internal static int GetLowerLevelsContext(Av1LevelBuffer levels, Point position, Av1TransformSize transformSize, Av1TransformClass transformClass)
     {
         int stats = Av1NzMap.GetNzMagnitude(levels, position, transformClass);
-        return Av1NzMap.GetNzMapContextFromStats(stats, levels, position, transformSize, transformClass);
+        return Av1NzMap.GetNzMapContextFromStats(stats, position, transformSize, transformClass);
     }
 
     /// <summary>
@@ -265,18 +271,17 @@ internal static class Av1SymbolContextHelper
     internal static sbyte GetNzMapContext(
         Av1LevelBuffer levels,
         Point position,
-        int scan_idx,
-        bool is_eob,
+        bool isEndOfBlock,
         Av1TransformSize transformSize,
         Av1TransformClass transformClass)
     {
-        if (is_eob)
+        if (isEndOfBlock)
         {
             return (sbyte)GetLowerLevelContextEndOfBlock(levels, position);
         }
 
         int stats = Av1NzMap.GetNzMagnitude(levels, position, transformClass);
-        return (sbyte)Av1NzMap.GetNzMapContextFromStats(stats, levels, position, transformSize, transformClass);
+        return (sbyte)Av1NzMap.GetNzMapContextFromStats(stats, position, transformSize, transformClass);
     }
 
     /// <summary>
@@ -294,7 +299,7 @@ internal static class Av1SymbolContextHelper
         {
             int pos = scan[i];
             Point position = levels.GetPosition(pos);
-            coefficientContexts[pos] = GetNzMapContext(levels, position, i, i == eob - 1, transformSize, transformClass);
+            coefficientContexts[pos] = GetNzMapContext(levels, position, i == eob - 1, transformSize, transformClass);
         }
     }
 

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

@@ -562,6 +562,22 @@ internal ref struct Av1SymbolDecoder
         return r.ReadSymbol(this.coefficientsBase[(int)transformSizeContext][(int)planeType][coefficientContext]);
     }
 
+    private void ReadCoefficientsBaseRangeLoop(Av1TransformSize transformSizeContext, Av1PlaneType planeType, int baseRangeContext, ref int level)
+    {
+        ref Av1SymbolReader r = ref this.reader;
+        Av1TransformSize limitedTransformSizeContext = (Av1TransformSize)Math.Min((int)transformSizeContext, (int)Av1TransformSize.Size32x32);
+        Av1Distribution distribution = this.coefficientsBaseRange[(int)transformSizeContext][(int)planeType][baseRangeContext];
+        for (int idx = 0; idx < Av1Constants.CoefficientBaseRange; idx += Av1Constants.BaseRangeSizeMinus1)
+        {
+            int coefficientBaseRange = r.ReadSymbol(distribution);
+            level += coefficientBaseRange;
+            if (coefficientBaseRange < Av1Constants.BaseRangeSizeMinus1)
+            {
+                break;
+            }
+        }
+    }
+
     internal int ReadGolomb()
     {
         ref Av1SymbolReader r = ref this.reader;

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

@@ -103,6 +103,8 @@ internal static class Av1TransformSizeExtensions
         Av1TransformSize.Size64x64,  // TX_64X16
     ];
 
+    // This is computed as:
+    // min(transform_width_log2, 5) + min(transform_height_log2, 5) - 4.
     private static readonly int[] Log2Minus4 = [
         0, // TX_4X4
         2, // TX_8X8

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

@@ -100,25 +100,26 @@ public class Av1CoefficientsEntropyTests
         Assert.Equal(endOfBlock, actuals[0]);
     }
 
-    [Fact]
-    public void RoundTripFullCoefficients()
+    [Theory]
+    [MemberData(nameof(GetBlockSize4x4Data))]
+    public void RoundTripFullCoefficientsYSize4x4(int bSize, int txSize, int txType)
     {
         // Assign
         const ushort endOfBlock = 16;
-        const Av1BlockSize blockSize = Av1BlockSize.Block4x4;
-        const Av1TransformSize transformSize = Av1TransformSize.Size4x4;
-        const Av1TransformType transformType = Av1TransformType.Identity;
-        const Av1PredictionMode intraDirection = Av1PredictionMode.DC;
         const Av1ComponentType componentType = Av1ComponentType.Luminance;
-        const Av1FilterIntraMode filterIntraMode = Av1FilterIntraMode.DC;
+        Av1BlockSize blockSize = (Av1BlockSize)bSize;
+        Av1TransformSize transformSize = (Av1TransformSize)txSize;
+        Av1TransformType transformType = (Av1TransformType)txType;
+        Av1PredictionMode intraDirection = Av1PredictionMode.DC;
+        Av1FilterIntraMode filterIntraMode = Av1FilterIntraMode.DC;
         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];
+        int[] aboveContexts = new int[transformSize.Get4x4WideCount()];
+        int[] leftContexts = new int[transformSize.Get4x4HighCount()];
         Av1TransformBlockContext transformBlockContext = new();
         Configuration configuration = Configuration.Default;
         Av1SymbolEncoder encoder = new(configuration, 100 / 8, BaseQIndex);
-        Span<int> coefficientsBuffer = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15];
+        Span<int> coefficientsBuffer = Enumerable.Range(0, blockSize.GetHeight() * blockSize.GetWidth()).ToArray();
         Span<int> actuals = new int[16 + 1];
 
         // Act
@@ -134,4 +135,17 @@ public class Av1CoefficientsEntropyTests
         Assert.Equal(endOfBlock, actuals[0]);
         Assert.Equal(coefficientsBuffer[..endOfBlock], actuals[1..(endOfBlock + 1)]);
     }
+
+    public static TheoryData<int, int, int> GetBlockSize4x4Data()
+    {
+        TheoryData<int, int, int> result = [];
+        Av1BlockSize blockSize = Av1BlockSize.Block4x4;
+        Av1TransformSize transformSize = blockSize.GetMaximumTransformSize();
+        for (Av1TransformType transformType = Av1TransformType.DctDct; transformType < Av1TransformType.VerticalDct; transformType++)
+        {
+            result.Add((int)blockSize, (int)transformSize, (int)transformType);
+        }
+
+        return result;
+    }
 }

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

@@ -131,6 +131,21 @@ public class Av1TransformSizeTests
         Assert.Equal(transformHeight, blockHeight);
     }
 
+    [Theory]
+    [MemberData(nameof(GetAllSizes))]
+    internal void LogMinus4ReturnsReferenceValues(int s)
+    {
+        // Assign
+        Av1TransformSize transformSize = (Av1TransformSize)s;
+        int expected = ReferenceLog2Minus4(transformSize);
+
+        // Act
+        int actual = transformSize.GetLog2Minus4();
+
+        // Assert
+        Assert.Equal(expected, actual);
+    }
+
     public static TheoryData<int> GetAllSizes()
     {
         TheoryData<int> combinations = [];
@@ -149,4 +164,11 @@ public class Av1TransformSizeTests
         int ratio = width > height ? width / height : height / width;
         return ratio;
     }
+
+    private static int ReferenceLog2Minus4(Av1TransformSize transformSize)
+    {
+        int widthLog2 = Av1Math.Log2(transformSize.GetWidth());
+        int heightLog2 = Av1Math.Log2(transformSize.GetHeight());
+        return Math.Min(widthLog2, 5) + Math.Min(heightLog2, 5) - 4;
+    }
 }