Transform set refactoring, closer to spec

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

@@ -11,26 +11,31 @@ namespace SixLabors.ImageSharp.Formats.Heif.Av1.Entropy;
 internal static class Av1SymbolContextHelper
 {
     public static readonly int[][] ExtendedTransformIndices = [
-        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
-        [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
-        [1, 3, 4, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
-        [1, 5, 6, 4, 0, 0, 0, 0, 0, 0, 2, 3, 0, 0, 0, 0],
-        [3, 4, 5, 8, 6, 7, 9, 10, 11, 0, 1, 2, 0, 0, 0, 0],
-        [7, 8, 9, 12, 10, 11, 13, 14, 15, 0, 1, 2, 3, 4, 5, 6],
+        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], // DCT only
+        [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], // Inter set 3
+        [1, 3, 4, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], // Intra set 2
+        [1, 5, 6, 4, 0, 0, 0, 0, 0, 0, 2, 3, 0, 0, 0, 0], // Intra set 1
+        [3, 4, 5, 8, 6, 7, 9, 10, 11, 0, 1, 2, 0, 0, 0, 0], // Inter set 2
+        [7, 8, 9, 12, 10, 11, 13, 14, 15, 0, 1, 2, 3, 4, 5, 6], // All 16, inter set 1
     ];
 
-    public static readonly int[][] ExtendedTransformIndicesInverse = [
-        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
-        [9, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
-        [9, 0, 3, 1, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
-        [9, 0, 10, 11, 3, 1, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0],
-        [9, 10, 11, 0, 1, 2, 4, 5, 3, 6, 7, 8, 0, 0, 0, 0],
-        [9, 10, 11, 12, 13, 14, 15, 0, 1, 2, 4, 5, 3, 6, 7, 8],
+    // Maps tx set types to the distribution indices. INTRA values only
+    private static readonly int[] ExtendedTransformSetToIndex = [0, -1, 2, 1, -1, -1];
+
+    /// <summary>
+    /// Section 5.11.48: Transform type syntax
+    /// </summary>
+    public static readonly Av1TransformType[][] ExtendedTransformInverse = [
+        [Av1TransformType.DctDct], // DCT only
+        [], // Inter set 3
+        [Av1TransformType.Identity, Av1TransformType.DctDct, Av1TransformType.AdstAdst, Av1TransformType.AdstDct, Av1TransformType.DctAdst], // Intra set 2
+        [Av1TransformType.Identity, Av1TransformType.DctDct, Av1TransformType.VerticalDct, Av1TransformType.HorizontalDct, Av1TransformType.AdstAdst, Av1TransformType.AdstDct, Av1TransformType.DctAdst], // Intra set 1
+        [], // Inter set 2
+        [], // All 16, inter set 1
     ];
 
     public static readonly int[] EndOfBlockOffsetBits = [0, 0, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9];
     public static readonly int[] EndOfBlockGroupStart = [0, 1, 2, 3, 5, 9, 17, 33, 65, 129, 257, 513];
-    private static readonly int[] TransformCountInSet = [1, 2, 5, 7, 12, 16];
     private static readonly byte[] EndOfBlockToPositionSmall = [
         0, 1, 2, // 0-2
         3, 3, // 3-4
@@ -59,9 +64,6 @@ internal static class Av1SymbolContextHelper
         11 // 513-
     ];
 
-    // Maps tx set types to the indices. INTRA values only
-    private static readonly int[] ExtendedTransformSetToIndex = [0, -1, 2, 1, -1, -1];
-
     internal static Av1TransformSize GetTransformSizeContext(Av1TransformSize originalSize)
         => (Av1TransformSize)(((int)originalSize.GetSquareSize() + (int)originalSize.GetSquareUpSize() + 1) >> 1);
 
@@ -239,11 +241,11 @@ internal static class Av1SymbolContextHelper
 
         if (useReducedSet)
         {
-            return Av1TransformSetType.Dtt4Identity;
+            return Av1TransformSetType.IntraSet2;
         }
 
         Av1TransformSize squareSize = transformSize.GetSquareSize();
-        return squareSize == Av1TransformSize.Size16x16 ? Av1TransformSetType.Dtt4Identity : Av1TransformSetType.Dtt4Identity1dDct;
+        return squareSize == Av1TransformSize.Size16x16 ? Av1TransformSetType.IntraSet2 : Av1TransformSetType.IntraSet1;
     }
 
     internal static Av1TransformType ConvertIntraModeToTransformType(Av1BlockModeInfo modeInfo, Av1PlaneType planeType)
@@ -299,7 +301,7 @@ internal static class Av1SymbolContextHelper
     /// <summary>
     /// SVT: get_ext_tx_types
     /// </summary>
-    internal static int GetExtendedTransformTypeCount(Av1TransformSetType setType) => TransformCountInSet[(int)setType];
+    internal static int GetExtendedTransformTypeCount(Av1TransformSetType setType) => ExtendedTransformInverse[(int)setType].Length;
 
     /// <summary>
     /// SVT: get_ext_tx_set

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

@@ -250,21 +250,16 @@ internal ref struct Av1SymbolDecoder
         // Ignoring INTER blocks here, as these should not end up here.
         // int inter_block = is_inter_block_dec(mbmi);
         Av1TransformSetType transformSetType = Av1SymbolContextHelper.GetExtendedTransformSetType(transformSize, useReducedTransformSet);
-        if (Av1SymbolContextHelper.GetExtendedTransformTypeCount(transformSetType) > 1 && baseQIndex > 0)
+        if (transformSetType > Av1TransformSetType.DctOnly && baseQIndex > 0)
         {
             int extendedSet = Av1SymbolContextHelper.GetExtendedTransformSet(transformSetType);
-
-            // eset == 0 should correspond to a set with only DCT_DCT and
-            // there is no need to read the tx_type
-            Guard.IsFalse(extendedSet == 0, nameof(extendedSet), string.Empty);
-
             Av1TransformSize squareTransformSize = transformSize.GetSquareSize();
             Av1PredictionMode intraMode = useFilterIntra
                 ? filterIntraMode.ToIntraDirection()
                 : intraDirection;
             ref Av1SymbolReader r = ref this.reader;
             int symbol = r.ReadSymbol(this.intraExtendedTransform[extendedSet][(int)squareTransformSize][(int)intraMode]);
-            transformType = (Av1TransformType)Av1SymbolContextHelper.ExtendedTransformIndicesInverse[(int)transformSetType][symbol];
+            transformType = Av1SymbolContextHelper.ExtendedTransformInverse[(int)transformSetType][symbol];
         }
 
         return transformType;

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

@@ -6,32 +6,35 @@ namespace SixLabors.ImageSharp.Formats.Heif.Av1.Transform;
 internal enum Av1TransformSetType
 {
     /// <summary>
-    /// DCT only.
+    /// Allowed transforms: DCT only.
     /// </summary>
     DctOnly,
 
     /// <summary>
-    /// DCT + Identity only
+    /// Allowed transforms: DCT + Identity only
     /// </summary>
-    DctIdentity,
+    InterSet3,
 
     /// <summary>
-    /// Discrete Trig transforms w/o flip (4) + Identity (1)
+    /// Allowed transforms: Discrete Trig transforms w/o flip (4) + Identity (1)
     /// </summary>
-    Dtt4Identity,
+    /// <remarks>Referenced in spec as TX_SET_INTRA_2.</remarks>
+    IntraSet2,
 
     /// <summary>
-    /// Discrete Trig transforms w/o flip (4) + Identity (1) + 1D Hor/vert DCT (2)
+    /// Allowed transforms: Discrete Trig transforms w/o flip (4) + Identity (1) + 1D Hor/vert DCT (2)
     /// </summary>
-    Dtt4Identity1dDct,
+    /// <remarks>Referenced in spec as TX_SET_INTRA_1.</remarks>
+    IntraSet1,
 
     /// <summary>
-    /// Discrete Trig transforms w/ flip (9) + Identity (1) + 1D Hor/Ver DCT (2)
+    /// Allowed transforms: Discrete Trig transforms w/ flip (9) + Identity (1) + 1D Hor/Ver DCT (2)
     /// </summary>
-    Dtt9Identity1dDct,
+    InterSet2,
 
     /// <summary>
-    /// Discrete Trig transforms w/ flip (9) + Identity (1) + 1D Hor/Ver (6)
+    /// Allowed transforms: Discrete Trig transforms w/ flip (9) + Identity (1) + 1D Hor/Ver (6)
     /// </summary>
-    All16
+    InterSet1,
+    AllSets
 }

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

@@ -601,9 +601,28 @@ public class Av1EntropyTests
         TheoryData<int, int, int> result = [];
         for (Av1TransformSize transformSize = Av1TransformSize.Size4x4; transformSize < Av1TransformSize.AllSizes; transformSize++)
         {
-            // TODO: Figure out why larger sizes don't round trip correctly.
+            if (transformSize == Av1TransformSize.Size16x16)
+            {
+                for (Av1PredictionMode intraDirection = Av1PredictionMode.IntraModeStart; intraDirection < Av1PredictionMode.IntraModeEnd; intraDirection++)
+                {
+                    result.Add((int)transformSize, (int)Av1FilterIntraMode.AllFilterIntraModes, (int)intraDirection);
+                }
+
+                if (transformSize == Av1TransformSize.Size16x16)
+                {
+                    result.Add((int)transformSize, 0, 0);
+                    result.Add((int)transformSize, 1, 1);
+                    result.Add((int)transformSize, 2, 2);
+                    result.Add((int)transformSize, 3, 6);
+                    result.Add((int)transformSize, 4, 0);
+                }
+
+                continue;
+            }
+
             if (transformSize.GetSquareSize() >= Av1TransformSize.Size16x16 || transformSize is Av1TransformSize.Size32x8 or Av1TransformSize.Size8x32)
             {
+                // DctOnly, doesn't make sense to test.
                 continue;
             }
 

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

@@ -38,8 +38,8 @@ public class Av1SymbolContextTests
         Av1TransformSetType transformSetType = (Av1TransformSetType)setType;
 
         // Act
-        int transformType = Av1SymbolContextHelper.ExtendedTransformIndicesInverse[(int)transformSetType][index];
-        int actualIndex = Av1SymbolContextHelper.ExtendedTransformIndices[(int)transformSetType][transformType];
+        Av1TransformType transformType = Av1SymbolContextHelper.ExtendedTransformInverse[(int)transformSetType][index];
+        int actualIndex = Av1SymbolContextHelper.ExtendedTransformIndices[(int)transformSetType][(int)transformType];
 
         // Assert
         Assert.Equal(actualIndex, index);
@@ -66,10 +66,10 @@ public class Av1SymbolContextTests
     public static TheoryData<int, int> GetExtendedTransformIndicesData()
     {
         TheoryData<int, int> result = [];
-        for (Av1TransformSetType setType = Av1TransformSetType.DctOnly; setType <= Av1TransformSetType.All16; setType++)
+        for (Av1TransformSetType setType = Av1TransformSetType.DctOnly; setType < Av1TransformSetType.AllSets; setType++)
         {
             int count = Av1SymbolContextHelper.GetExtendedTransformTypeCount(setType);
-            for (int index = 1; index < count; index++)
+            for (int index = 0; index < count; index++)
             {
                 result.Add((int)setType, index);
             }