Encode partition symbol

src/ImageSharp/Formats/Heif/Av1/Av1Constants.cs

@@ -180,4 +180,6 @@ internal static class Av1Constants
     /// Maximum number of stages in a 1-dimensioanl transform function.
     /// </summary>
     public const int MaxTransformStageNumber = 12;
+
+    public const int PartitionProbabilitySet = 4;
 }

src/ImageSharp/Formats/Heif/Av1/Av1Math.cs

@@ -166,4 +166,9 @@ internal static class Av1Math
         DebugGuard.MustBeGreaterThanOrEqualTo(bit, 1, nameof(bit));
         return (int)((value + (1L << (bit - 1))) >> bit);
     }
+
+    /// <summary>
+    /// <paramref name="a"/> implies <paramref name="b"/>.
+    /// </summary>
+    internal static bool Implies(bool a, bool b) => !a || b;
 }

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

@@ -64,42 +64,24 @@ internal ref struct Av1SymbolDecoder
         return (Av1PartitionType)r.ReadSymbol(this.tilePartitionTypes[context]);
     }
 
-    public bool ReadSplitOrHorizontal(Av1BlockSize blockSize, int context)
+    /// <summary>
+    /// SVT: partition_gather_vert_alike
+    /// </summary>
+    public Av1PartitionType ReadSplitOrHorizontal(Av1BlockSize blockSize, int context)
     {
-        Av1Distribution input = this.tilePartitionTypes[context];
-        uint p = Av1Distribution.ProbabilityTop;
-        p -= GetElementProbability(input, Av1PartitionType.Horizontal);
-        p -= GetElementProbability(input, Av1PartitionType.Split);
-        p -= GetElementProbability(input, Av1PartitionType.HorizontalA);
-        p -= GetElementProbability(input, Av1PartitionType.HorizontalB);
-        p -= GetElementProbability(input, Av1PartitionType.VerticalA);
-        if (blockSize != Av1BlockSize.Block128x128)
-        {
-            p -= GetElementProbability(input, Av1PartitionType.Horizontal4);
-        }
-
-        Av1Distribution distribution = new(Av1Distribution.ProbabilityTop - p);
+        Av1Distribution distribution = GetSplitOrHorizontalDistribution(this.tilePartitionTypes, blockSize, context);
         ref Av1SymbolReader r = ref this.reader;
-        return r.ReadSymbol(distribution) > 0;
+        return r.ReadSymbol(distribution) > 0 ? Av1PartitionType.Split : Av1PartitionType.Horizontal;
     }
 
-    public bool ReadSplitOrVertical(Av1BlockSize blockSize, int context)
+    /// <summary>
+    /// SVT: partition_gather_horz_alike
+    /// </summary>
+    public Av1PartitionType ReadSplitOrVertical(Av1BlockSize blockSize, int context)
     {
-        Av1Distribution input = this.tilePartitionTypes[context];
-        uint p = Av1Distribution.ProbabilityTop;
-        p -= GetElementProbability(input, Av1PartitionType.Vertical);
-        p -= GetElementProbability(input, Av1PartitionType.Split);
-        p -= GetElementProbability(input, Av1PartitionType.HorizontalA);
-        p -= GetElementProbability(input, Av1PartitionType.VerticalA);
-        p -= GetElementProbability(input, Av1PartitionType.VerticalB);
-        if (blockSize != Av1BlockSize.Block128x128)
-        {
-            p -= GetElementProbability(input, Av1PartitionType.Vertical4);
-        }
-
-        Av1Distribution distribution = new(Av1Distribution.ProbabilityTop - p);
+        Av1Distribution distribution = GetSplitOrVerticalDistribution(this.tilePartitionTypes, blockSize, context);
         ref Av1SymbolReader r = ref this.reader;
-        return r.ReadSymbol(distribution) > 0;
+        return r.ReadSymbol(distribution) > 0 ? Av1PartitionType.Split : Av1PartitionType.Vertical;
     }
 
     public Av1PredictionMode ReadYMode(Av1BlockModeInfo? aboveModeInfo, Av1BlockModeInfo? leftModeInfo)
@@ -259,6 +241,40 @@ internal ref struct Av1SymbolDecoder
         return r.ReadSymbol(this.chromeForLumaAlpha[context]);
     }
 
+    internal static Av1Distribution GetSplitOrHorizontalDistribution(Av1Distribution[] inputs, Av1BlockSize blockSize, int context)
+    {
+        Av1Distribution input = inputs[context];
+        uint p = Av1Distribution.ProbabilityTop;
+        p -= GetElementProbability(input, Av1PartitionType.Horizontal);
+        p -= GetElementProbability(input, Av1PartitionType.Split);
+        p -= GetElementProbability(input, Av1PartitionType.HorizontalA);
+        p -= GetElementProbability(input, Av1PartitionType.HorizontalB);
+        p -= GetElementProbability(input, Av1PartitionType.VerticalA);
+        if (blockSize != Av1BlockSize.Block128x128)
+        {
+            p -= GetElementProbability(input, Av1PartitionType.Horizontal4);
+        }
+
+        return new(Av1Distribution.ProbabilityTop - p);
+    }
+
+    internal static Av1Distribution GetSplitOrVerticalDistribution(Av1Distribution[] inputs, Av1BlockSize blockSize, int context)
+    {
+        Av1Distribution input = inputs[context];
+        uint p = Av1Distribution.ProbabilityTop;
+        p -= GetElementProbability(input, Av1PartitionType.Vertical);
+        p -= GetElementProbability(input, Av1PartitionType.Split);
+        p -= GetElementProbability(input, Av1PartitionType.HorizontalA);
+        p -= GetElementProbability(input, Av1PartitionType.VerticalA);
+        p -= GetElementProbability(input, Av1PartitionType.VerticalB);
+        if (blockSize != Av1BlockSize.Block128x128)
+        {
+            p -= GetElementProbability(input, Av1PartitionType.Vertical4);
+        }
+
+        return new(Av1Distribution.ProbabilityTop - p);
+    }
+
     private static uint GetElementProbability(Av1Distribution probability, Av1PartitionType element)
         => probability[(int)element - 1] - probability[(int)element];
 }

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

@@ -10,22 +10,52 @@ internal class Av1SymbolEncoder : IDisposable
     private readonly Av1Distribution tileIntraBlockCopy = Av1DefaultDistributions.IntraBlockCopy;
     private readonly Av1Distribution[] tilePartitionTypes = Av1DefaultDistributions.PartitionTypes;
 
-    private Av1SymbolWriter? writer;
+    private bool isDisposed;
+    private Av1SymbolWriter writer;
 
     public Av1SymbolEncoder(Configuration configuration, int initialSize)
         => this.writer = new(configuration, initialSize);
 
     public void WriteUseIntraBlockCopy(bool value)
-        => this.writer!.WriteSymbol(value ? 1 : 0, this.tileIntraBlockCopy);
+    {
+        ref Av1SymbolWriter w = ref this.writer;
+        w.WriteSymbol(value ? 1 : 0, this.tileIntraBlockCopy);
+    }
+
+    public void WritePartitionType(Av1PartitionType partitionType, int context)
+    {
+        ref Av1SymbolWriter w = ref this.writer;
+        w.WriteSymbol((int)partitionType, this.tilePartitionTypes[context]);
+    }
+
+    public void WriteSplitOrHorizontal(Av1PartitionType partitionType, Av1BlockSize blockSize, int context)
+    {
+        Av1Distribution distribution = Av1SymbolDecoder.GetSplitOrHorizontalDistribution(this.tilePartitionTypes, blockSize, context);
+        int value = (partitionType == Av1PartitionType.Split) ? 1 : 0;
+        ref Av1SymbolWriter w = ref this.writer;
+        w.WriteSymbol(value, distribution);
+    }
 
-    public void WritePartitionType(Av1PartitionType value, int context)
-        => this.writer!.WriteSymbol((int)value, this.tilePartitionTypes[context]);
+    public void WriteSplitOrVertical(Av1PartitionType partitionType, Av1BlockSize blockSize, int context)
+    {
+        Av1Distribution distribution = Av1SymbolDecoder.GetSplitOrVerticalDistribution(this.tilePartitionTypes, blockSize, context);
+        int value = (partitionType == Av1PartitionType.Split) ? 1 : 0;
+        ref Av1SymbolWriter w = ref this.writer;
+        w.WriteSymbol(value, distribution);
+    }
 
-    public IMemoryOwner<byte> Exit() => this.writer!.Exit();
+    public IMemoryOwner<byte> Exit()
+    {
+        ref Av1SymbolWriter w = ref this.writer;
+        return w.Exit();
+    }
 
     public void Dispose()
     {
-        this.writer?.Dispose();
-        this.writer = null;
+        if (!this.isDisposed)
+        {
+            this.writer.Dispose();
+            this.isDisposed = true;
+        }
     }
 }

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

@@ -14,7 +14,6 @@ internal class Av1TileReader : IAv1TileReader
 {
     private static readonly int[] SgrprojXqdMid = [-32, 31];
     private static readonly int[] WienerTapsMid = [3, -7, 15];
-    private const int PartitionProbabilitySet = 4;
     private static readonly int[] Signs = [0, -1, 1];
     private static readonly int[] DcSignContexts = [
         1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0,
@@ -182,13 +181,11 @@ internal class Av1TileReader : IAv1TileReader
             }
             else if (hasColumns)
             {
-                bool splitOrVertical = reader.ReadSplitOrVertical(blockSize, ctx);
-                partitionType = splitOrVertical ? Av1PartitionType.Split : Av1PartitionType.Horizontal;
+                partitionType = reader.ReadSplitOrHorizontal(blockSize, ctx);
             }
             else if (hasRows)
             {
-                bool splitOrHorizontal = reader.ReadSplitOrHorizontal(blockSize, ctx);
-                partitionType = splitOrHorizontal ? Av1PartitionType.Split : Av1PartitionType.Vertical;
+                partitionType = reader.ReadSplitOrVertical(blockSize, ctx);
             }
         }
 
@@ -1886,7 +1883,7 @@ internal class Av1TileReader : IAv1TileReader
         int left = (leftCtx >> blockSizeLog) & 0x1;
         DebugGuard.IsTrue(blockSize.Get4x4WidthLog2() == blockSize.Get4x4HeightLog2(), "Blocks should be square");
         DebugGuard.MustBeGreaterThanOrEqualTo(blockSizeLog, 0, nameof(blockSizeLog));
-        return ((left << 1) + above) + (blockSizeLog * PartitionProbabilitySet);
+        return ((left << 1) + above) + (blockSizeLog * Av1Constants.PartitionProbabilitySet);
     }
 
     private void UpdatePartitionContext(Point modeInfoLocation, Av1TileInfo tileLoc, Av1SuperblockInfo superblockInfo, Av1BlockSize subSize, Av1BlockSize blockSize, Av1PartitionType partition)

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

@@ -2,9 +2,11 @@
 // Licensed under the Six Labors Split License.
 
 using System.Buffers;
+using System.Reflection;
 using SixLabors.ImageSharp.Formats.Heif.Av1;
 using SixLabors.ImageSharp.Formats.Heif.Av1.Tiling;
 using SixLabors.ImageSharp.Memory;
+using static System.Runtime.InteropServices.JavaScript.JSType;
 
 namespace SixLabors.ImageSharp.Tests.Formats.Heif.Av1;
 
@@ -217,6 +219,80 @@ public class SymbolTests
         Assert.Equal(values, actuals);
     }
 
+    [Theory]
+    [InlineData((int)Av1BlockSize.Block4x4, 7)]
+    [InlineData((int)Av1BlockSize.Block4x4, 5)]
+    [InlineData((int)Av1BlockSize.Block8x4, 7)]
+    [InlineData((int)Av1BlockSize.Block4x8, 7)]
+    [InlineData((int)Av1BlockSize.Block32x64, 7)]
+    [InlineData((int)Av1BlockSize.Block64x32, 7)]
+    [InlineData((int)Av1BlockSize.Block64x64, 7)]
+    public void RoundTripSplitOrHorizontalPartitionType(int blockSize, int context)
+    {
+        // Assign
+        Configuration configuration = Configuration.Default;
+        Av1SymbolEncoder encoder = new(configuration, 100 / 8);
+        Av1PartitionType[] values = [
+            Av1PartitionType.Split, Av1PartitionType.Split, Av1PartitionType.Split, Av1PartitionType.Horizontal,
+            Av1PartitionType.Split, Av1PartitionType.Split, Av1PartitionType.Horizontal, Av1PartitionType.Horizontal];
+        Av1PartitionType[] actuals = new Av1PartitionType[values.Length];
+
+        // Act
+        foreach (Av1PartitionType value in values)
+        {
+            encoder.WriteSplitOrHorizontal(value, (Av1BlockSize)blockSize, context);
+        }
+
+        using IMemoryOwner<byte> encoded = encoder.Exit();
+
+        Av1SymbolDecoder decoder = new(encoded.GetSpan(), 0);
+        Av1SymbolReader reader = new(encoded.GetSpan());
+        for (int i = 0; i < values.Length; i++)
+        {
+            actuals[i] = decoder.ReadSplitOrHorizontal((Av1BlockSize)blockSize, context);
+        }
+
+        // Assert
+        Assert.Equal(values, actuals);
+    }
+
+    [Theory]
+    [InlineData((int)Av1BlockSize.Block4x4, 7)]
+    [InlineData((int)Av1BlockSize.Block4x4, 5)]
+    [InlineData((int)Av1BlockSize.Block8x4, 7)]
+    [InlineData((int)Av1BlockSize.Block4x8, 7)]
+    [InlineData((int)Av1BlockSize.Block32x64, 7)]
+    [InlineData((int)Av1BlockSize.Block64x32, 7)]
+    [InlineData((int)Av1BlockSize.Block64x64, 7)]
+    public void RoundTripSplitOrVerticalPartitionType(int blockSize, int context)
+    {
+        // Assign
+        Configuration configuration = Configuration.Default;
+        Av1SymbolEncoder encoder = new(configuration, 100 / 8);
+        Av1PartitionType[] values = [
+            Av1PartitionType.Split, Av1PartitionType.Split, Av1PartitionType.Split, Av1PartitionType.Vertical,
+            Av1PartitionType.Split, Av1PartitionType.Split, Av1PartitionType.Vertical, Av1PartitionType.Vertical];
+        Av1PartitionType[] actuals = new Av1PartitionType[values.Length];
+
+        // Act
+        foreach (Av1PartitionType value in values)
+        {
+            encoder.WriteSplitOrVertical(value, (Av1BlockSize)blockSize, context);
+        }
+
+        using IMemoryOwner<byte> encoded = encoder.Exit();
+
+        Av1SymbolDecoder decoder = new(encoded.GetSpan(), 0);
+        Av1SymbolReader reader = new(encoded.GetSpan());
+        for (int i = 0; i < values.Length; i++)
+        {
+            actuals[i] = decoder.ReadSplitOrVertical((Av1BlockSize)blockSize, context);
+        }
+
+        // Assert
+        Assert.Equal(values, actuals);
+    }
+
     [Fact]
     public void RoundTripUseIntraBlockCopy()
     {
@@ -235,7 +311,6 @@ public class SymbolTests
         using IMemoryOwner<byte> encoded = encoder.Exit();
 
         Av1SymbolDecoder decoder = new(encoded.GetSpan(), 0);
-        Av1SymbolReader reader = new(encoded.GetSpan());
         for (int i = 0; i < values.Length; i++)
         {
             actuals[i] = decoder.ReadUseIntraBlockCopy();