Consolidate AV1 decoder block scratch

HEIF_IMPLEMENTATION_PLAN.md

@@ -721,6 +721,13 @@ Final decoder allocation, lifetime, precision, architecture, and test-validity a
   boundaries. Frame planes enforce their contiguous single-span invariant before allocation; palette,
   transform, film-grain, super-resolution, color-conversion, and alpha workspaces remain bounded and
   allocator owned. No per-block managed allocation remains in reconstruction.
+- [x] Block reconstruction now uses one exact-size signed-short owner for inverse quantization, inverse
+  transform, compound prediction, convolution, and chroma-from-luma scratch. Even-length slices provide
+  the integer workspaces without another rent. Monochrome reserves no chroma coefficients, and 4:2:0,
+  4:2:2, and 4:4:4 reserve two symmetric chroma planes at their coded subsampling. This replaces three
+  constructor rents and their catch-all rollback path; exact allocation length, coefficient span length,
+  and exactly-once return pass for all four layouts, with 549 adjacent reconstruction tests passing direct
+  net11 VSTest in Release.
 - [x] Valid unsupported tile-list syntax is rejected explicitly. Reserved and metadata OBUs are consumed
   only after bounded framing and trailing-bit validation. Eight-, ten-, and twelve-bit reconstruction,
   presentation, alpha, restoration, and film-grain paths retain native precision.

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

@@ -54,22 +54,27 @@ internal sealed class Av1BlockDecoder : IDisposable
     private readonly Av1ReferenceFrameStore referenceFrames;
 
     /// <summary>
-    /// Owns the reusable raster-order inverse-quantization buffer.
+    /// Owns all reusable inverse-quantization, transform, and prediction storage.
     /// </summary>
-    private readonly IMemoryOwner<int> inverseQuantizationOwner;
+    private readonly IMemoryOwner<short> workspaceOwner;
 
     /// <summary>
-    /// Owns the reusable two-dimensional inverse-transform workspace.
+    /// The inverse-quantization prefix length in signed-short storage elements.
     /// </summary>
-    private readonly IMemoryOwner<int> transformWorkspaceOwner;
+    private readonly int inverseQuantizationStorageLength;
 
     /// <summary>
-    /// Owns the reusable directional and filter-intra prediction workspace.
+    /// The inverse-transform workspace offset in signed-short storage elements.
     /// </summary>
-    private readonly IMemoryOwner<short> predictionScratchOwner;
+    private readonly int transformWorkspaceOffset;
 
     /// <summary>
-    /// The reusable predictor portion of <see cref="predictionScratchOwner"/>, excluding compound and chroma-from-luma storage.
+    /// The prediction workspace offset in signed-short storage elements.
+    /// </summary>
+    private readonly int predictionScratchOffset;
+
+    /// <summary>
+    /// The reusable predictor portion of <see cref="workspaceOwner"/>, excluding compound and chroma-from-luma storage.
     /// </summary>
     private readonly int predictorWorkingLength;
 
@@ -122,76 +127,56 @@ internal sealed class Av1BlockDecoder : IDisposable
 
         // One scratch plane is reused for every transform unit. Its maximum size must cover a complete superblock
         // across all coded planes, with chroma dimensions reduced independently by their subsampling axes.
-        int inverseQuantizationSize = ySize +
-            (this.sequenceHeader.ColorConfig.SubSamplingX ? ySize >> 2 : ySize) +
-            (this.sequenceHeader.ColorConfig.SubSamplingY ? ySize >> 2 : ySize);
-
-        IMemoryOwner<int>? inverseQuantizationOwner = null;
-        IMemoryOwner<int>? transformWorkspaceOwner = null;
-        IMemoryOwner<short>? predictionScratchOwner = null;
-        try
-        {
-            inverseQuantizationOwner = this.frameBuffer.MemoryAllocator.Allocate<int>(inverseQuantizationSize);
-            transformWorkspaceOwner = this.frameBuffer.MemoryAllocator.Allocate<int>(Av1TransformWorkspace.MaximumLength);
-            int maximumBlockLength = 1 << sequenceHeader.SuperblockSizeLog2;
-            int maximumBlockArea = maximumBlockLength * maximumBlockLength;
-            int predictorWorkingLength = Math.Max(
-                Av1PredictionDecoder.ScratchLength,
-                Math.Max(
-                    Av1TranslationalInterPredictor.GetScratchLength(maximumBlockLength, maximumBlockLength),
-                    Av1ScaledInterPredictor.GetMaximumScaledScratchLength(maximumBlockLength, maximumBlockLength)));
-
-            int compoundMaskLength = (maximumBlockArea + 1) >> 1;
-            int predictorWorkingOffset = (2 * maximumBlockArea) + compoundMaskLength;
-            int chromaFromLumaOffset = predictorWorkingOffset + predictorWorkingLength;
-
-            // Compound prediction retains both high-precision reference planes plus the full-resolution luma mask.
-            // Keeping those planes, convolution workspace, and CfL surface in one owner avoids independent managed
-            // buffers while ensuring the two scratch consumers never overlap.
-            int predictionScratchLength = chromaFromLumaOffset + Av1ChromaFromLumaContext.BufferLength;
-            predictionScratchOwner = this.frameBuffer.MemoryAllocator.Allocate<short>(predictionScratchLength);
-
-            this.inverseQuantizationOwner = inverseQuantizationOwner;
-            this.transformWorkspaceOwner = transformWorkspaceOwner;
-            this.predictionScratchOwner = predictionScratchOwner;
-            this.predictorWorkingLength = predictorWorkingLength;
-            this.predictionDecoder = new(
-                sequenceHeader,
-                frameHeader,
-                predictionScratchOwner.Memory.Slice(predictorWorkingOffset, predictorWorkingLength),
-                paletteColorIndexMaps);
-            this.isLoopFilterEnabled = frameHeader.LoopFilterParameters.FilterLevel[0] != 0 ||
-                frameHeader.LoopFilterParameters.FilterLevel[1] != 0;
-
-            this.chromaFromLumaContext = new(
-                sequenceHeader.ColorConfig,
-                predictionScratchOwner.Memory.Slice(chromaFromLumaOffset, Av1ChromaFromLumaContext.BufferLength));
-        }
-        catch
-        {
-            // A constructor that does not return transfers no ownership to its caller. Unwind successful rents in
-            // reverse order so allocator diagnostics and pooled buffers remain balanced after any later allocation.
-            predictionScratchOwner?.Dispose();
-            transformWorkspaceOwner?.Dispose();
-            inverseQuantizationOwner?.Dispose();
-            throw;
-        }
+        ObuColorConfig colorConfig = this.sequenceHeader.ColorConfig;
+        int chromaSubsampling = (colorConfig.SubSamplingX ? 1 : 0) + (colorConfig.SubSamplingY ? 1 : 0);
+        int chromaSize = ySize >> chromaSubsampling;
+        int inverseQuantizationSize = colorConfig.IsMonochrome ? ySize : ySize + (2 * chromaSize);
+        int maximumBlockLength = 1 << sequenceHeader.SuperblockSizeLog2;
+        int maximumBlockArea = maximumBlockLength * maximumBlockLength;
+        int predictorWorkingLength = Math.Max(
+            Av1PredictionDecoder.ScratchLength,
+            Math.Max(
+                Av1TranslationalInterPredictor.GetScratchLength(maximumBlockLength, maximumBlockLength),
+                Av1ScaledInterPredictor.GetMaximumScaledScratchLength(maximumBlockLength, maximumBlockLength)));
+
+        int compoundMaskLength = (maximumBlockArea + 1) >> 1;
+        int predictorWorkingOffset = (2 * maximumBlockArea) + compoundMaskLength;
+        int chromaFromLumaOffset = predictorWorkingOffset + predictorWorkingLength;
+        int predictionScratchLength = chromaFromLumaOffset + Av1ChromaFromLumaContext.BufferLength;
+        this.inverseQuantizationStorageLength = inverseQuantizationSize * 2;
+        this.transformWorkspaceOffset = this.inverseQuantizationStorageLength;
+        this.predictionScratchOffset = this.transformWorkspaceOffset + (Av1TransformWorkspace.MaximumLength * 2);
+
+        // Integer workspaces occupy even signed-short slices so one allocator owner can retain the complete block
+        // lifetime while prediction still receives the Memory<short> contract needed by its reusable context.
+        this.workspaceOwner = this.frameBuffer.MemoryAllocator.Allocate<short>(
+            this.predictionScratchOffset + predictionScratchLength);
+
+        Memory<short> predictionScratch = this.workspaceOwner.Memory[this.predictionScratchOffset..];
+        this.predictorWorkingLength = predictorWorkingLength;
+        this.predictionDecoder = new(
+            sequenceHeader,
+            frameHeader,
+            predictionScratch.Slice(predictorWorkingOffset, predictorWorkingLength),
+            paletteColorIndexMaps);
+        this.isLoopFilterEnabled = frameHeader.LoopFilterParameters.FilterLevel[0] != 0 ||
+            frameHeader.LoopFilterParameters.FilterLevel[1] != 0;
+
+        this.chromaFromLumaContext = new(
+            sequenceHeader.ColorConfig,
+            predictionScratch.Slice(chromaFromLumaOffset, Av1ChromaFromLumaContext.BufferLength));
     }
 
     /// <summary>
     /// Gets the reusable raster-order coefficient buffer populated by inverse quantization.
     /// </summary>
-    public Span<int> CurrentInverseQuantizationCoefficients => this.inverseQuantizationOwner.Memory.Span;
+    public Span<int> CurrentInverseQuantizationCoefficients
+        => MemoryMarshal.Cast<short, int>(this.workspaceOwner.Memory.Span[..this.inverseQuantizationStorageLength]);
 
     /// <summary>
     /// Releases the pooled reconstruction workspaces owned by this decoder.
     /// </summary>
-    public void Dispose()
-    {
-        this.predictionScratchOwner.Dispose();
-        this.transformWorkspaceOwner.Dispose();
-        this.inverseQuantizationOwner.Dispose();
-    }
+    public void Dispose() => this.workspaceOwner.Dispose();
 
     /// <summary>
     /// Resets the per-plane packed coefficient cursors before reconstructing a superblock.
@@ -216,7 +201,10 @@ internal sealed class Av1BlockDecoder : IDisposable
     /// <param name="tileInfo">The tile boundaries used to determine neighbor availability.</param>
     public void DecodeBlock(Av1BlockModeInfo modeInfo, Point modeInfoPosition, Av1BlockSize blockSize, Av1SuperblockInfo superblockInfo, Av1TileInfo tileInfo)
     {
-        Span<int> transformWorkspace = this.transformWorkspaceOwner.Memory.Span;
+        Span<int> transformWorkspace = MemoryMarshal.Cast<short, int>(
+            this.workspaceOwner.Memory.Span.Slice(
+                this.transformWorkspaceOffset,
+                Av1TransformWorkspace.MaximumLength * 2));
 
         ObuColorConfig colorConfig = this.sequenceHeader.ColorConfig;
         Av1TransformType transformType;
@@ -351,7 +339,7 @@ internal sealed class Av1BlockDecoder : IDisposable
                 int maximumBlockLength = 1 << this.sequenceHeader.SuperblockSizeLog2;
                 int maximumBlockArea = maximumBlockLength * maximumBlockLength;
                 int compoundMaskStorageLength = (maximumBlockArea + 1) >> 1;
-                Span<short> predictionStorage = this.predictionScratchOwner.Memory.Span;
+                Span<short> predictionStorage = this.workspaceOwner.Memory.Span[this.predictionScratchOffset..];
                 Span<short> secondPredictionStorage = predictionStorage[..maximumBlockArea];
                 Span<ushort> firstCompoundPrediction = MemoryMarshal.Cast<short, ushort>(
                     predictionStorage.Slice(maximumBlockArea, maximumBlockArea));

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

@@ -5,6 +5,9 @@ using SixLabors.ImageSharp.Formats.Heif.Av1;
 using SixLabors.ImageSharp.Formats.Heif.Av1.OpenBitstreamUnit;
 using SixLabors.ImageSharp.Formats.Heif.Av1.Pipeline.LoopFilter;
 using SixLabors.ImageSharp.Formats.Heif.Av1.Pipeline.Quantizers;
+using SixLabors.ImageSharp.Formats.Heif.Av1.Prediction;
+using SixLabors.ImageSharp.Formats.Heif.Av1.Prediction.ChromaFromLuma;
+using SixLabors.ImageSharp.Formats.Heif.Av1.Prediction.Inter;
 using SixLabors.ImageSharp.Formats.Heif.Av1.ReferenceFrames;
 using SixLabors.ImageSharp.Formats.Heif.Av1.Tiling;
 using SixLabors.ImageSharp.Formats.Heif.Av1.Transform;
@@ -115,16 +118,20 @@ public class Av1FrameBufferTests
     }
 
     /// <summary>
-    /// Verifies that each later block-decoder workspace failure releases every workspace rented earlier.
+    /// Verifies that block reconstruction uses one exact-size owner across monochrome and chroma plane layouts.
     /// </summary>
     [Theory]
-    [InlineData(3, 1)]
-    [InlineData(4, 2)]
-    public void BlockDecoderConstructorFailureReleasesEarlierWorkspaces(
-        int failureAllocationNumber,
-        int successfulWorkspaceCount)
+    [InlineData(true, false, false, 4096)]
+    [InlineData(false, true, true, 6144)]
+    [InlineData(false, true, false, 8192)]
+    [InlineData(false, false, false, 12288)]
+    public void BlockDecoderUsesOneContiguousWorkspaceOwner(
+        bool isMonochrome,
+        bool subsamplingX,
+        bool subsamplingY,
+        int expectedInverseQuantizationSize)
     {
-        FailingTestMemoryAllocator allocator = new(failureAllocationNumber);
+        TestMemoryAllocator allocator = new();
         Configuration configuration = Configuration.Default.Clone();
         configuration.MemoryAllocator = allocator;
         ObuSequenceHeader sequenceHeader = new()
@@ -134,12 +141,20 @@ public class Av1FrameBufferTests
             Use128x128Superblock = false,
             ColorConfig = new ObuColorConfig
             {
-                IsMonochrome = true,
+                IsMonochrome = isMonochrome,
+                SubSamplingX = subsamplingX,
+                SubSamplingY = subsamplingY,
                 BitDepth = Av1BitDepth.EightBit
             }
         };
 
-        using Av1FrameBuffer<byte> frameBuffer = new(configuration, sequenceHeader, Av1ColorFormat.Yuv400, false);
+        Av1ColorFormat colorFormat = isMonochrome
+            ? Av1ColorFormat.Yuv400
+            : subsamplingX
+                ? subsamplingY ? Av1ColorFormat.Yuv420 : Av1ColorFormat.Yuv422
+                : Av1ColorFormat.Yuv444;
+
+        using Av1FrameBuffer<byte> frameBuffer = new(configuration, sequenceHeader, colorFormat, false);
         ObuFrameHeader frameHeader = new()
         {
             ModeInfoColumnCount = 16,
@@ -152,28 +167,46 @@ public class Av1FrameBufferTests
         Av1InverseQuantizer inverseQuantizer = new(sequenceHeader, frameHeader);
         using Av1ReferenceFrameStore referenceFrames = new();
 
-        // The frame's luma plane is allocation attempt one. Resetting only the logs preserves that counter while
-        // isolating the block-decoder owners that must be returned when a later workspace rent fails.
+        // Reset the frame-plane logs so the following assertions describe only the block decoder's scratch owner.
         allocator.EnableNonThreadSafeLogging();
 
-        Assert.Throws<InvalidMemoryOperationException>(
-            () => new Av1BlockDecoder(
+        int maximumBlockLength = 1 << sequenceHeader.SuperblockSizeLog2;
+        int maximumBlockArea = maximumBlockLength * maximumBlockLength;
+        int predictorWorkingLength = Math.Max(
+            Av1PredictionDecoder.ScratchLength,
+            Math.Max(
+                Av1TranslationalInterPredictor.GetScratchLength(maximumBlockLength, maximumBlockLength),
+                Av1ScaledInterPredictor.GetMaximumScaledScratchLength(maximumBlockLength, maximumBlockLength)));
+
+        int predictionScratchLength =
+            (2 * maximumBlockArea) +
+            ((maximumBlockArea + 1) >> 1) +
+            predictorWorkingLength +
+            Av1ChromaFromLumaContext.BufferLength;
+
+        int expectedWorkspaceLength =
+            (expectedInverseQuantizationSize * 2) +
+            (Av1TransformWorkspace.MaximumLength * 2) +
+            predictionScratchLength;
+
+        TestMemoryAllocator.AllocationRequest workspaceAllocation;
+        using (Av1BlockDecoder blockDecoder = new(
                 sequenceHeader,
                 frameHeader,
                 frameBuffer,
                 loopFilterContext,
                 inverseQuantizer,
-                referenceFrames));
-
-        Assert.Equal(failureAllocationNumber, allocator.AllocationAttemptCount);
-        Assert.Equal(successfulWorkspaceCount, allocator.AllocationLog.Count);
-        Assert.Equal(successfulWorkspaceCount, allocator.ReturnLog.Count);
-        foreach (TestMemoryAllocator.AllocationRequest allocation in allocator.AllocationLog)
+                referenceFrames))
         {
-            Assert.Contains(
-                allocator.ReturnLog,
-                returned => returned.HashCodeOfBuffer == allocation.HashCodeOfBuffer);
+            workspaceAllocation = Assert.Single(allocator.AllocationLog);
+            Assert.Empty(allocator.ReturnLog);
+            Assert.Equal(typeof(short), workspaceAllocation.ElementType);
+            Assert.Equal(expectedWorkspaceLength, workspaceAllocation.Length);
+            Assert.Equal(expectedInverseQuantizationSize, blockDecoder.CurrentInverseQuantizationCoefficients.Length);
         }
+
+        TestMemoryAllocator.ReturnRequest returned = Assert.Single(allocator.ReturnLog);
+        Assert.Equal(workspaceAllocation.AllocationId, returned.AllocationId);
     }
 
     /// <summary>