Pool AV1 parse neighbor contexts

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

@@ -13,7 +13,7 @@ namespace SixLabors.ImageSharp.Formats.Heif.Av1;
 /// <summary>
 /// Decodes one AV1 still-image elementary stream into an ImageSharp image.
 /// </summary>
-internal class Av1Decoder : IAv1TileReader
+internal sealed class Av1Decoder : IAv1TileReader, IDisposable
 {
     /// <summary>
     /// The open-bitstream-unit parser for the current image item.
@@ -185,4 +185,13 @@ internal class Av1Decoder : IAv1TileReader
 
         this.tileReader.ReadTile(tileData, tileNum);
     }
+
+    /// <summary>
+    /// Releases the tile reader and its frame-scoped parsing storage.
+    /// </summary>
+    public void Dispose()
+    {
+        this.tileReader?.Dispose();
+        this.tileReader = null;
+    }
 }

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

@@ -777,8 +777,8 @@ internal ref struct Av1SymbolDecoder
     public int ReadCoefficients(
         Av1BlockModeInfo modeInfo,
         Point blockPosition,
-        int[] aboveContexts,
+        Span<int> aboveContexts,
-        int[] leftContexts,
+        Span<int> leftContexts,
         int aboveOffset,
         int leftOffset,
         int plane,
@@ -1176,8 +1176,8 @@ internal ref struct Av1SymbolDecoder
     /// <param name="modeBlockToRightEdge">The signed distance from the mode block to the right frame edge.</param>
     /// <param name="modeBlockToBottomEdge">The signed distance from the mode block to the bottom frame edge.</param>
     private static void UpdateCoefficientContext(
-        int[] aboveContexts,
+        Span<int> aboveContexts,
-        int[] leftContexts,
+        Span<int> leftContexts,
         int blocksWide,
         int blocksHigh,
         Av1TransformSize transformSize,
@@ -1194,23 +1194,23 @@ internal ref struct Av1SymbolDecoder
         if (modeBlockToRightEdge < 0)
         {
             int aboveContextCount = Math.Min(transformSizeWide, blocksWide - blockPosition.X);
-            Array.Fill(aboveContexts, culLevel, aboveOffset, aboveContextCount);
+            aboveContexts.Slice(aboveOffset, aboveContextCount).Fill(culLevel);
-            Array.Fill(aboveContexts, 0, aboveOffset + aboveContextCount, transformSizeWide - aboveContextCount);
+            aboveContexts.Slice(aboveOffset + aboveContextCount, transformSizeWide - aboveContextCount).Clear();
         }
         else
         {
-            Array.Fill(aboveContexts, culLevel, aboveOffset, transformSizeWide);
+            aboveContexts.Slice(aboveOffset, transformSizeWide).Fill(culLevel);
         }
 
         if (modeBlockToBottomEdge < 0)
         {
             int leftContextCount = Math.Min(transformSizeHigh, blocksHigh - blockPosition.Y);
-            Array.Fill(leftContexts, culLevel, leftOffset, leftContextCount);
+            leftContexts.Slice(leftOffset, leftContextCount).Fill(culLevel);
-            Array.Fill(leftContexts, 0, leftOffset + leftContextCount, transformSizeHigh - leftContextCount);
+            leftContexts.Slice(leftOffset + leftContextCount, transformSizeHigh - leftContextCount).Clear();
         }
         else
         {
-            Array.Fill(leftContexts, culLevel, leftOffset, transformSizeHigh);
+            leftContexts.Slice(leftOffset, transformSizeHigh).Fill(culLevel);
         }
     }
 

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

@@ -1,65 +1,85 @@
 // Copyright (c) Six Labors.
 // Licensed under the Six Labors Split License.
 
+using System.Buffers;
 using SixLabors.ImageSharp.Formats.Heif.Av1.OpenBitstreamUnit;
 using SixLabors.ImageSharp.Formats.Heif.Av1.Transform;
+using SixLabors.ImageSharp.Memory;
 
 namespace SixLabors.ImageSharp.Formats.Heif.Av1.Tiling;
 
 /// <summary>
 /// Stores entropy, partition, and transform contexts for 4-by-4 blocks above the current block.
 /// </summary>
-internal class Av1ParseAboveNeighbor4x4Context
+internal sealed class Av1ParseAboveNeighbor4x4Context : IDisposable
 {
     /// <summary>
-    /// Stores DC-sign and cumulative coefficient-level contexts for each plane above the current block.
+    /// The region containing transform-width contexts.
     /// </summary>
-    private readonly int[][] aboveContext = new int[Av1Constants.MaxPlanes][];
+    private const int TransformWidthRegionIndex = 0;
 
     /// <summary>
-    /// Stores segmentation-prediction contexts from the preceding 4x4 row.
+    /// The region containing partition-width contexts.
     /// </summary>
-    private readonly int[] aboveSegmentIdPredictionContext;
+    private const int PartitionWidthRegionIndex = 1;
 
     /// <summary>
-    /// Stores compound-reference group contexts from the preceding 4x4 row.
+    /// The first region containing a color plane's coefficient contexts.
     /// </summary>
-    private readonly int[] aboveCompGroupIndex;
+    private const int PlaneContextRegionStart = 2;
+
+    /// <summary>
+    /// Owns the contiguous above-neighbor storage until this instance is disposed.
+    /// </summary>
+    private IMemoryOwner<int>? memory;
+
+    /// <summary>
+    /// The number of mode-information columns stored in each logical region.
+    /// </summary>
+    private readonly int contextLength;
 
     /// <summary>
     /// Initializes a new instance of the <see cref="Av1ParseAboveNeighbor4x4Context"/> class.
     /// </summary>
+    /// <param name="configuration">The configuration providing the memory allocator.</param>
     /// <param name="planesCount">The number of color planes.</param>
     /// <param name="modeInfoColumnCount">The frame width in 4x4 mode-information columns.</param>
-    public Av1ParseAboveNeighbor4x4Context(int planesCount, int modeInfoColumnCount)
+    public Av1ParseAboveNeighbor4x4Context(Configuration configuration, int planesCount, int modeInfoColumnCount)
     {
-        this.AboveTransformWidth = new int[modeInfoColumnCount];
+        this.contextLength = modeInfoColumnCount;
-        this.AbovePartitionWidth = new int[modeInfoColumnCount];
+        int regionCount = PlaneContextRegionStart + planesCount;
-        for (int i = 0; i < planesCount; i++)
+        int totalLength = checked(regionCount * modeInfoColumnCount);
-        {
-            this.aboveContext[i] = new int[modeInfoColumnCount];
-        }
 
-        this.aboveSegmentIdPredictionContext = new int[modeInfoColumnCount];
+        // Every region spans the same aligned frame width and shares the tile-reader lifetime.
-        this.aboveCompGroupIndex = new int[modeInfoColumnCount];
+        // One clean rent replaces the jagged array and its per-region arrays while preserving zero initialization.
+        this.memory = configuration.MemoryAllocator.Allocate<int>(totalLength, AllocationOptions.Clean);
     }
 
     /// <summary>
     /// Gets a buffer holding the partition context of the previous 4x4 block row.
     /// </summary>
-    public int[] AbovePartitionWidth { get; }
+    public Span<int> AbovePartitionWidth => this.GetRegion(PartitionWidthRegionIndex);
 
     /// <summary>
     /// Gets a buffer holding the transform sizes of the previous 4x4 block row.
     /// </summary>
-    public int[] AboveTransformWidth { get; }
+    public Span<int> AboveTransformWidth => this.GetRegion(TransformWidthRegionIndex);
 
     /// <summary>
     /// Gets the coefficient context row for the specified plane.
     /// </summary>
     /// <param name="plane">The zero-based plane index.</param>
     /// <returns>The coefficient contexts for the plane.</returns>
-    public int[] GetContext(int plane) => this.aboveContext[plane];
+    public Span<int> GetContext(int plane) => this.GetRegion(PlaneContextRegionStart + plane);
+
+    /// <summary>
+    /// Returns the above-neighbor storage to the configured memory allocator.
+    /// </summary>
+    public void Dispose()
+    {
+        this.memory?.Dispose();
+        this.memory = null;
+    }
 
     /// <summary>
     /// Resets above-neighbor state for the active tile-column range.
@@ -71,15 +91,12 @@ internal class Av1ParseAboveNeighbor4x4Context
     {
         int planeCount = sequenceHeader.ColorConfig.PlaneCount;
         int width = modeInfoColumnEnd - modeInfoColumnStart;
-        Array.Fill(this.AboveTransformWidth, Av1TransformSize.Size64x64.GetWidth(), 0, width);
+        this.AboveTransformWidth[..width].Fill(Av1TransformSize.Size64x64.GetWidth());
-        Array.Fill(this.AbovePartitionWidth, 0, 0, width);
+        this.AbovePartitionWidth[..width].Clear();
         for (int i = 0; i < planeCount; i++)
         {
-            Array.Fill(this.aboveContext[i], 0, 0, width);
+            this.GetContext(i)[..width].Clear();
         }
-
-        Array.Fill(this.aboveSegmentIdPredictionContext, 0, 0, width);
-        Array.Fill(this.aboveCompGroupIndex, 0, 0, width);
     }
 
     /// <summary>
@@ -97,7 +114,7 @@ internal class Av1ParseAboveNeighbor4x4Context
         int value = Av1PartitionContext.GetAboveContext(subSize);
 
         DebugGuard.MustBeLessThanOrEqualTo(startIndex, this.AboveTransformWidth.Length - bw, nameof(startIndex));
-        Array.Fill(this.AbovePartitionWidth, value, startIndex, bw);
+        this.AbovePartitionWidth.Slice(startIndex, bw).Fill(value);
     }
 
     /// <summary>
@@ -120,7 +137,7 @@ internal class Av1ParseAboveNeighbor4x4Context
         }
 
         DebugGuard.MustBeLessThanOrEqualTo(startIndex, this.AboveTransformWidth.Length - n4w, nameof(startIndex));
-        Array.Fill(this.AboveTransformWidth, transformWidth, startIndex, n4w);
+        this.AboveTransformWidth.Slice(startIndex, n4w).Fill(transformWidth);
     }
 
     /// <summary>
@@ -129,6 +146,17 @@ internal class Av1ParseAboveNeighbor4x4Context
     /// <param name="plane">The zero-based plane index.</param>
     /// <param name="offset">The first context index to clear.</param>
     /// <param name="length">The number of context entries to clear.</param>
-    internal void ClearContext(int plane, int offset, int length)
+    public void ClearContext(int plane, int offset, int length)
-        => Array.Fill(this.aboveContext[plane], 0, offset, length);
+        => this.GetContext(plane).Slice(offset, length).Clear();
+
+    /// <summary>
+    /// Gets one logical row from the contiguous above-neighbor allocation.
+    /// </summary>
+    /// <param name="regionIndex">The zero-based logical region index.</param>
+    /// <returns>The requested context row.</returns>
+    private Span<int> GetRegion(int regionIndex)
+    {
+        ObjectDisposedException.ThrowIf(this.memory is null, this);
+        return this.memory.Memory.Span.Slice(regionIndex * this.contextLength, this.contextLength);
+    }
 }

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

@@ -1,60 +1,80 @@
 // Copyright (c) Six Labors.
 // Licensed under the Six Labors Split License.
 
+using System.Buffers;
 using SixLabors.ImageSharp.Formats.Heif.Av1.OpenBitstreamUnit;
 using SixLabors.ImageSharp.Formats.Heif.Av1.Transform;
+using SixLabors.ImageSharp.Memory;
 
 namespace SixLabors.ImageSharp.Formats.Heif.Av1.Tiling;
 
 /// <summary>
 /// Stores entropy, partition, and transform contexts for 4-by-4 blocks left of the current block.
 /// </summary>
-internal class Av1ParseLeftNeighbor4x4Context
+internal sealed class Av1ParseLeftNeighbor4x4Context : IDisposable
 {
     /// <summary>
-    /// Stores DC-sign and cumulative coefficient-level contexts for each plane left of the current block.
+    /// The region containing transform-height contexts.
     /// </summary>
-    private readonly int[][] leftContext = new int[Av1Constants.MaxPlanes][];
+    private const int TransformHeightRegionIndex = 0;
 
     /// <summary>
-    /// Stores segmentation-prediction contexts for the current superblock row.
+    /// The region containing partition-height contexts.
     /// </summary>
-    private readonly int[] leftSegmentIdPredictionContext;
+    private const int PartitionHeightRegionIndex = 1;
 
     /// <summary>
-    /// Stores compound-reference group contexts for the current superblock row.
+    /// The first region containing a color plane's coefficient contexts.
     /// </summary>
-    private readonly int[] leftCompGroupIndex;
+    private const int PlaneContextRegionStart = 2;
+
+    /// <summary>
+    /// Owns the contiguous left-neighbor storage until this instance is disposed.
+    /// </summary>
+    private IMemoryOwner<int>? memory;
+
+    /// <summary>
+    /// The number of superblock-row entries stored in each logical region.
+    /// </summary>
+    private readonly int contextLength;
 
     /// <summary>
     /// Initializes a new instance of the <see cref="Av1ParseLeftNeighbor4x4Context"/> class.
     /// </summary>
+    /// <param name="configuration">The configuration providing the memory allocator.</param>
     /// <param name="planesCount">The number of color planes.</param>
     /// <param name="superblockModeInfoSize">The superblock height in 4x4 mode-information rows.</param>
-    public Av1ParseLeftNeighbor4x4Context(int planesCount, int superblockModeInfoSize)
+    public Av1ParseLeftNeighbor4x4Context(Configuration configuration, int planesCount, int superblockModeInfoSize)
     {
-        this.LeftTransformHeight = new int[superblockModeInfoSize];
+        this.contextLength = superblockModeInfoSize;
-        this.LeftPartitionHeight = new int[superblockModeInfoSize];
+        int regionCount = PlaneContextRegionStart + planesCount;
-        for (int i = 0; i < planesCount; i++)
+        int totalLength = checked(regionCount * superblockModeInfoSize);
-        {
-            this.leftContext[i] = new int[superblockModeInfoSize];
-        }
 
-        this.leftSegmentIdPredictionContext = new int[superblockModeInfoSize];
+        // Every region spans the same superblock height and shares the tile-reader lifetime.
-        this.leftCompGroupIndex = new int[superblockModeInfoSize];
+        // One clean rent replaces the jagged array and its per-region arrays while preserving zero initialization.
+        this.memory = configuration.MemoryAllocator.Allocate<int>(totalLength, AllocationOptions.Clean);
     }
 
     /// <summary>
     /// Gets a buffer holding the partition context of the left 4x4 blocks corresponding
     /// to the current super block row.
     /// </summary>
-    public int[] LeftPartitionHeight { get; }
+    public Span<int> LeftPartitionHeight => this.GetRegion(PartitionHeightRegionIndex);
 
     /// <summary>
     /// Gets a buffer holding the transform sizes of the left 4x4 blocks corresponding
     /// to the current super block row.
     /// </summary>
-    public int[] LeftTransformHeight { get; }
+    public Span<int> LeftTransformHeight => this.GetRegion(TransformHeightRegionIndex);
+
+    /// <summary>
+    /// Returns the left-neighbor storage to the configured memory allocator.
+    /// </summary>
+    public void Dispose()
+    {
+        this.memory?.Dispose();
+        this.memory = null;
+    }
 
     /// <summary>
     /// Resets all left-neighbor state for a new superblock row.
@@ -64,16 +84,12 @@ internal class Av1ParseLeftNeighbor4x4Context
     {
         int blockCount = sequenceHeader.SuperblockModeInfoSize;
         int planeCount = sequenceHeader.ColorConfig.PlaneCount;
-        int neighbor4x4Count = sequenceHeader.SuperblockModeInfoSize;
+        this.LeftTransformHeight[..blockCount].Fill(Av1TransformSize.Size64x64.GetHeight());
-        Array.Fill(this.LeftTransformHeight, Av1TransformSize.Size64x64.GetHeight(), 0, blockCount);
+        this.LeftPartitionHeight[..blockCount].Clear();
-        Array.Fill(this.LeftPartitionHeight, 0, 0, blockCount);
         for (int i = 0; i < planeCount; i++)
         {
-            Array.Fill(this.leftContext[i], 0, 0, blockCount);
+            this.GetContext(i)[..blockCount].Clear();
         }
-
-        Array.Fill(this.leftSegmentIdPredictionContext, 0, 0, blockCount);
-        Array.Fill(this.leftCompGroupIndex, 0, 0, blockCount);
     }
 
     /// <summary>
@@ -90,7 +106,7 @@ internal class Av1ParseLeftNeighbor4x4Context
         int bh = blockSize.Get4x4HighCount();
         int value = Av1PartitionContext.GetLeftContext(subSize);
         DebugGuard.MustBeLessThanOrEqualTo(startIndex, this.LeftPartitionHeight.Length - bh, nameof(startIndex));
-        Array.Fill(this.LeftPartitionHeight, value, startIndex, bh);
+        this.LeftPartitionHeight.Slice(startIndex, bh).Fill(value);
     }
 
     /// <summary>
@@ -113,7 +129,7 @@ internal class Av1ParseLeftNeighbor4x4Context
         }
 
         DebugGuard.MustBeLessThanOrEqualTo(startIndex, this.LeftTransformHeight.Length - n4h, nameof(startIndex));
-        Array.Fill(this.LeftTransformHeight, transformHeight, startIndex, n4h);
+        this.LeftTransformHeight.Slice(startIndex, n4h).Fill(transformHeight);
     }
 
     /// <summary>
@@ -122,13 +138,24 @@ internal class Av1ParseLeftNeighbor4x4Context
     /// <param name="plane">The zero-based plane index.</param>
     /// <param name="offset">The first context index to clear.</param>
     /// <param name="length">The number of context entries to clear.</param>
-    internal void ClearContext(int plane, int offset, int length)
+    public void ClearContext(int plane, int offset, int length)
-        => Array.Fill(this.leftContext[plane], 0, offset, length);
+        => this.GetContext(plane).Slice(offset, length).Clear();
 
     /// <summary>
     /// Gets the coefficient context column for the specified plane.
     /// </summary>
     /// <param name="plane">The zero-based plane index.</param>
     /// <returns>The coefficient contexts for the plane.</returns>
-    internal int[] GetContext(int plane) => this.leftContext[plane];
+    public Span<int> GetContext(int plane) => this.GetRegion(PlaneContextRegionStart + plane);
+
+    /// <summary>
+    /// Gets one logical column from the contiguous left-neighbor allocation.
+    /// </summary>
+    /// <param name="regionIndex">The zero-based logical region index.</param>
+    /// <returns>The requested context column.</returns>
+    private Span<int> GetRegion(int regionIndex)
+    {
+        ObjectDisposedException.ThrowIf(this.memory is null, this);
+        return this.memory.Memory.Span.Slice(regionIndex * this.contextLength, this.contextLength);
+    }
 }

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

@@ -14,7 +14,7 @@ namespace SixLabors.ImageSharp.Formats.Heif.Av1.Tiling;
 /// <summary>
 /// Parses partition, mode, transform, and coefficient syntax for one AV1 tile.
 /// </summary>
-internal class Av1TileReader : IAv1TileReader
+internal sealed class Av1TileReader : IAv1TileReader, IDisposable
 {
     /// <summary>
     /// The default self-guided restoration projection coefficients for each color plane.
@@ -180,13 +180,24 @@ internal class Av1TileReader : IAv1TileReader
             Av1Math.AlignPowerOf2(sequenceHeader.MaxFrameWidth, sequenceHeader.SuperblockSizeLog2) >> sequenceHeader.SuperblockSizeLog2;
         int modeInfoWideColumnCount = superblockColumnCount * sequenceHeader.SuperblockModeInfoSize;
         modeInfoWideColumnCount = Av1Math.AlignPowerOf2(modeInfoWideColumnCount, sequenceHeader.SuperblockSizeLog2 - Av1Constants.ModeInfoSizeLog2);
-        this.aboveNeighborContext = new Av1ParseAboveNeighbor4x4Context(planesCount, modeInfoWideColumnCount);
-        this.leftNeighborContext = new Av1ParseLeftNeighbor4x4Context(planesCount, sequenceHeader.SuperblockModeInfoSize);
         this.transformUnitCount = new int[Av1Constants.MaxPlanes][];
         this.transformUnitCount[0] = new int[this.FrameInfo.ModeInfoCount];
         this.transformUnitCount[1] = new int[this.FrameInfo.ModeInfoCount];
         this.transformUnitCount[2] = new int[this.FrameInfo.ModeInfoCount];
         this.coefficientIndex = new int[Av1Constants.MaxPlanes];
+
+        this.aboveNeighborContext = new Av1ParseAboveNeighbor4x4Context(configuration, planesCount, modeInfoWideColumnCount);
+        try
+        {
+            this.leftNeighborContext = new Av1ParseLeftNeighbor4x4Context(configuration, planesCount, sequenceHeader.SuperblockModeInfoSize);
+        }
+        catch
+        {
+            // The reader is not returned when its second context allocation fails, so release
+            // the first rent here rather than relying on an owner that the caller cannot reach.
+            this.aboveNeighborContext.Dispose();
+            throw;
+        }
     }
 
     /// <summary>
@@ -215,6 +226,15 @@ internal class Av1TileReader : IAv1TileReader
     /// </summary>
     public Av1FrameInfo FrameInfo { get; }
 
+    /// <summary>
+    /// Returns the tile-neighbor context storage to the configured memory allocator.
+    /// </summary>
+    public void Dispose()
+    {
+        this.aboveNeighborContext.Dispose();
+        this.leftNeighborContext.Dispose();
+    }
+
     /// <summary>
     /// Parses one tile's partition, mode, transform, coefficient, and filter syntax in superblock order.
     /// </summary>
@@ -984,8 +1004,8 @@ internal class Av1TileReader : IAv1TileReader
         int leftOffset)
     {
         Av1TransformBlockContext transformBlockContext = new();
-        ReadOnlySpan<int> aboveContext = this.aboveNeighborContext.GetContext(plane).AsSpan(aboveOffset);
+        ReadOnlySpan<int> aboveContext = this.aboveNeighborContext.GetContext(plane)[aboveOffset..];
-        ReadOnlySpan<int> leftContext = this.leftNeighborContext.GetContext(plane).AsSpan(leftOffset);
+        ReadOnlySpan<int> leftContext = this.leftNeighborContext.GetContext(plane)[leftOffset..];
         int dcSign = 0;
         int k = 0;
         int mask = (1 << Av1Constants.CoefficientContextBitCount) - 1;

src/ImageSharp/Formats/Heif/Av1HeifItemDecoder.cs

@@ -61,7 +61,7 @@ internal class Av1HeifItemDecoder<TPixel> : IHeifItemDecoder<TPixel>
             out HeifContentLightLevel? obuContentLightLevel,
             out HeifMasteringDisplayColorVolume? obuMasteringDisplayColorVolume);
 
-        Av1Decoder decoder = new(configuration);
+        using Av1Decoder decoder = new(configuration);
         Image<TPixel> image = decoder.Decode<TPixel>(data, colorProfile, codecConfiguration);
         HeifMetadata metadata = image.Metadata.GetHeifMetadata();
         metadata.CompressionMethod = this.CompressionMethod;