Add AV1 palette map entropy coding

HEIF_IMPLEMENTATION_PLAN.md

@@ -852,6 +852,7 @@ Encoder verification contract:
 - [x] Empty-transform block skip now compares the complete live rate of the two decoder-identical syntax choices after luma and every coded chroma plane have been selected. Current libaom forces all-intra blocks to non-skip; this encoder retains that behavior for every non-empty block and for equal-cost empty blocks, but emits block skip when its adapted context cost is strictly lower than non-skip plus all empty-transform coefficient costs. Costing and writing share the same above-and-left skip-context calculation, and the coefficient estimator returns after the transform-block-skip symbol without reading coefficient storage. This adds no allocation, copy, or persistent state. A focused adapted-CDF regression proves both outcomes through the production decision helper, the two production all-zero fixtures still prove the default real block path, the exact net11 Release rebuild remains at 1,005 warnings and zero errors, all 8,975 HEIF/AV1 tests pass, and current-main `aomdec` at `a40ed1ea9e4ecc3df58a5bccb76623f2c94ae727` accepts all 29 regenerated payloads.
 - [~] Palette entropy coding now mirrors current libaom's adaptive luma-mode, chroma-mode, palette-size, and spatial color-index distributions, together with its truncated-binary uniform code used by palette colors. The complete mutable palette probability graph is created once on first palette search or write, so the current palette-disabled frame path retains zero palette allocations. Three focused regressions cover every legal 2-through-8 color alphabet and every defined mode, size, and color-index context; all 1,928 entropy cases and all 8,978 HEIF/AV1 cases pass direct net11 Release VSTest. The exact Release rebuild remains at 1,005 warnings and zero errors. This checkpoint adds the exact entropy foundation only: palette candidate generation, retained color and index storage, mode decision, map tokenization, and production syntax remain incomplete, and no generated payload changed.
 - [~] Luma and chroma palette-color coding now matches current libaom's neighbor-cache flags, sorted delta representation, wrapped V-plane deltas, strict delta-versus-raw V selection, and fixed-point color-rate model at 8, 10, and 12 bits. Encoder costing and emission use only fixed stack spans, including explicitly initialized cache-membership state, and steady-state color costing allocates zero managed bytes. The decoder consumes the same bounded color-syntax primitive after the tile reader derives its neighbor cache, removing duplicated color parsing without changing retained palette ownership. Nine focused syntax, exact palette decode, constrained-allocation, truncation, presentation, and allocation cases pass; all 1,933 entropy cases and all 8,983 HEIF/AV1 cases pass direct net11 Release VSTest. The exact Release rebuild remains at 1,005 warnings and zero errors. Retained encoder palette colors, neighbor caches, color-index maps, candidate generation, and production palette selection remain incomplete, and the compact 8-byte frame mode entries were not enlarged.
+- [~] Palette color-index map coding now shares the exact current-libaom neighbor weights, stable color ordering, five context classes, first-index uniform code, and diagonal wavefront between encoder costing, encoder writing, and decoder parsing. The decoder's stack-allocated context scores are explicitly cleared before accumulation, removing an invalid dependency on uninitialized stack contents. Costing and writing use a closed generic operation while the shared driver owns traversal and context derivation, so the semantic operations remain independent of map layout and tail handling. The path adds no retained state or per-call managed allocation. Twelve focused map, exact palette decode, padding, trailing-bit, and allocation cases pass; all 1,941 entropy cases and all 8,991 HEIF/AV1 cases pass direct net11 Release VSTest. The exact Release rebuild remains at 1,005 warnings and zero errors. Production payloads remain unchanged because palette selection is still disabled; retained colors, neighbor caches, index-map storage, candidate generation, and production palette mode decision remain incomplete.
 - [x] The expanded checkpoint exposed a pre-existing transform-block test that asserted uninitialized pooled padding was zero. The test now initializes the complete physical luma plane with a sentinel and proves the block operation leaves both adjacent padding samples unchanged. The exact net11 Release rebuild remains at 1,005 baseline warnings and zero errors, the focused allocator-order set passes 30 of 30 cases, and the complete HEIF/AV1 namespace passes 8,859 of 8,859 direct VSTest cases with zero failures or skips.
 - [x] Combined-frame OBU output now counts the byte-aligned frame and tile-group headers, non-final tile-size fields, and owned tile payloads before emitting the OBU size. It retains only the small allocator-owned header scratch and writes each entropy-coded tile span directly from its detached owner, removing the second file-sized allocator rent and complete-payload copy. A 64 KiB regression proves exactly one sub-payload-sized byte rent with a balanced return and verifies the exact streamed tile tail; the existing two-tile round trip proves size-prefix and ordering parity. The focused writer and production-frame set passes 32 of 32 direct net11 VSTest cases, current-main `aomdec` accepts all 29 generated native-format payloads, and the complete HEIF/AV1 namespace passes 8,860 of 8,860 cases with zero failures or skips.
 - [x] Finalized fixed-block decisions now set the block-level transform-skip flag only when every retained luma and coded chroma transform has zero EOB, matching current libaom's conjunction of per-plane skip state. The previous always-false flag produced legal but redundant non-skip and zero-coefficient syntax. Monochrome and 4:2:0 regressions prove both branches from actual coefficient state; the focused decision and production-frame set passes 32 of 32 direct net11 VSTest cases. Current-main `aomdec` accepts all 29 regenerated payloads, the recorded decoded-frame MD5s are unchanged, and affected 16x16 constant 8-bit and 10-bit payloads are one byte smaller. The complete HEIF/AV1 namespace passes 8,862 of 8,862 cases with zero failures or skips.

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

@@ -7,6 +7,7 @@ using SixLabors.ImageSharp.Formats.Heif.Av1.Prediction.ChromaFromLuma;
 using SixLabors.ImageSharp.Formats.Heif.Av1.Prediction.Inter;
 using SixLabors.ImageSharp.Formats.Heif.Av1.Tiling;
 using SixLabors.ImageSharp.Formats.Heif.Av1.Transform;
+using SixLabors.ImageSharp.Memory;
 
 namespace SixLabors.ImageSharp.Formats.Heif.Av1.Entropy;
 
@@ -245,6 +246,45 @@ internal ref struct Av1SymbolDecoder
         }
     }
 
+    /// <summary>
+    /// Reads a palette color-index map in AV1 diagonal wavefront order.
+    /// </summary>
+    /// <param name="paletteSize">The number of colors in the palette.</param>
+    /// <param name="planeType">The luma or chroma plane class.</param>
+    /// <param name="rows">The number of coded map rows.</param>
+    /// <param name="columns">The number of coded map columns.</param>
+    /// <param name="colorIndexMap">The row-addressable destination map.</param>
+    public void ReadPaletteColorMap(
+        int paletteSize,
+        Av1PlaneType planeType,
+        int rows,
+        int columns,
+        Buffer2DRegion<byte> colorIndexMap)
+    {
+        colorIndexMap.DangerousGetRowSpan(0)[0] = (byte)this.ReadUniform(paletteSize);
+        Span<byte> colorOrder = stackalloc byte[Av1Constants.PaletteMaxSize];
+        for (int diagonal = 1; diagonal < rows + columns - 1; diagonal++)
+        {
+            int firstColumn = Math.Min(diagonal, columns - 1);
+            int lastColumn = Math.Max(0, diagonal - rows + 1);
+            for (int column = firstColumn; column >= lastColumn; column--)
+            {
+                int row = diagonal - column;
+                int colorContext = Av1PaletteColorMap.GetContext(
+                    colorIndexMap,
+                    row,
+                    column,
+                    paletteSize,
+                    colorIndex: -1,
+                    colorOrder,
+                    out _);
+
+                int colorOrderIndex = this.ReadPaletteColorIndex(paletteSize, colorContext, planeType);
+                colorIndexMap.DangerousGetRowSpan(row)[column] = colorOrder[colorOrderIndex];
+            }
+        }
+    }
+
     /// <summary>
     /// Reads a finite subexponential value recentered around a preceding value.
     /// </summary>

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

@@ -8,6 +8,7 @@ using SixLabors.ImageSharp.Formats.Heif.Av1.Prediction;
 using SixLabors.ImageSharp.Formats.Heif.Av1.Prediction.ChromaFromLuma;
 using SixLabors.ImageSharp.Formats.Heif.Av1.Tiling;
 using SixLabors.ImageSharp.Formats.Heif.Av1.Transform;
+using SixLabors.ImageSharp.Memory;
 
 namespace SixLabors.ImageSharp.Formats.Heif.Av1.Entropy;
 
@@ -240,6 +241,40 @@ internal class Av1SymbolEncoder : IDisposable
             int bitCount);
     }
 
+    /// <summary>
+    /// Defines how the shared palette-map traversal handles its uniform first index and adaptive remaining indices.
+    /// </summary>
+    private interface IPaletteColorMapOperation
+    {
+        /// <summary>
+        /// Handles the first uniformly coded palette index.
+        /// </summary>
+        /// <param name="encoder">The tile symbol encoder.</param>
+        /// <param name="paletteSize">The number of colors in the palette.</param>
+        /// <param name="colorIndex">The first palette index.</param>
+        /// <returns>The index's rate contribution.</returns>
+        public static abstract int ProcessFirstIndex(
+            Av1SymbolEncoder encoder,
+            int paletteSize,
+            int colorIndex);
+
+        /// <summary>
+        /// Handles one context-adaptive palette color-order index.
+        /// </summary>
+        /// <param name="encoder">The tile symbol encoder.</param>
+        /// <param name="paletteSize">The number of colors in the palette.</param>
+        /// <param name="planeType">The luma or chroma plane class.</param>
+        /// <param name="colorContext">The spatial color-index context.</param>
+        /// <param name="colorOrderIndex">The index in the context-specific color order.</param>
+        /// <returns>The index's rate contribution.</returns>
+        public static abstract int ProcessColorIndex(
+            Av1SymbolEncoder encoder,
+            int paletteSize,
+            Av1PlaneType planeType,
+            int colorContext,
+            int colorOrderIndex);
+    }
+
     /// <summary>
     /// Writes an unsigned fixed-width literal to the tile entropy stream.
     /// </summary>
@@ -579,6 +614,51 @@ internal class Av1SymbolEncoder : IDisposable
         }
     }
 
+    /// <summary>
+    /// Gets the current fixed-point rate of a complete palette color-index map.
+    /// </summary>
+    /// <param name="paletteSize">The number of colors in the palette.</param>
+    /// <param name="planeType">The luma or chroma plane class.</param>
+    /// <param name="rows">The number of coded map rows.</param>
+    /// <param name="columns">The number of coded map columns.</param>
+    /// <param name="colorIndexMap">The complete row-addressable color-index map.</param>
+    /// <returns>The rate cost in 1/512-bit units.</returns>
+    public int GetPaletteColorMapCost(
+        int paletteSize,
+        Av1PlaneType planeType,
+        int rows,
+        int columns,
+        Buffer2DRegion<byte> colorIndexMap)
+        => this.ProcessPaletteColorMap<PaletteColorMapCostOperation>(
+            paletteSize,
+            planeType,
+            rows,
+            columns,
+            colorIndexMap);
+
+    /// <summary>
+    /// Writes a complete palette color-index map in AV1 diagonal wavefront order.
+    /// </summary>
+    /// <param name="paletteSize">The number of colors in the palette.</param>
+    /// <param name="planeType">The luma or chroma plane class.</param>
+    /// <param name="rows">The number of coded map rows.</param>
+    /// <param name="columns">The number of coded map columns.</param>
+    /// <param name="colorIndexMap">The complete row-addressable color-index map.</param>
+    public void WritePaletteColorMap(
+        int paletteSize,
+        Av1PlaneType planeType,
+        int rows,
+        int columns,
+        Buffer2DRegion<byte> colorIndexMap)
+    {
+        _ = this.ProcessPaletteColorMap<PaletteColorMapWriteOperation>(
+            paletteSize,
+            planeType,
+            rows,
+            columns,
+            colorIndexMap);
+    }
+
     /// <summary>
     /// Writes the frame-local intra-block-copy flag.
     /// </summary>
@@ -1527,6 +1607,56 @@ internal class Av1SymbolEncoder : IDisposable
         }
     }
 
+    /// <summary>
+    /// Traverses a palette color-index map once for either live rate costing or entropy emission.
+    /// </summary>
+    /// <typeparam name="TOperation">The closed map-symbol operation.</typeparam>
+    /// <param name="paletteSize">The number of colors in the palette.</param>
+    /// <param name="planeType">The luma or chroma plane class.</param>
+    /// <param name="rows">The number of coded map rows.</param>
+    /// <param name="columns">The number of coded map columns.</param>
+    /// <param name="colorIndexMap">The complete row-addressable color-index map.</param>
+    /// <returns>The rate cost in 1/512-bit units, or zero while writing.</returns>
+    private int ProcessPaletteColorMap<TOperation>(
+        int paletteSize,
+        Av1PlaneType planeType,
+        int rows,
+        int columns,
+        Buffer2DRegion<byte> colorIndexMap)
+        where TOperation : struct, IPaletteColorMapOperation
+    {
+        int colorIndex = colorIndexMap.DangerousGetRowSpan(0)[0];
+        int cost = TOperation.ProcessFirstIndex(this, paletteSize, colorIndex);
+        Span<byte> colorOrder = stackalloc byte[Av1Constants.PaletteMaxSize];
+        for (int diagonal = 1; diagonal < rows + columns - 1; diagonal++)
+        {
+            int firstColumn = Math.Min(diagonal, columns - 1);
+            int lastColumn = Math.Max(0, diagonal - rows + 1);
+            for (int column = firstColumn; column >= lastColumn; column--)
+            {
+                int row = diagonal - column;
+                colorIndex = colorIndexMap.DangerousGetRowSpan(row)[column];
+                int colorContext = Av1PaletteColorMap.GetContext(
+                    colorIndexMap,
+                    row,
+                    column,
+                    paletteSize,
+                    colorIndex,
+                    colorOrder,
+                    out int colorOrderIndex);
+
+                cost += TOperation.ProcessColorIndex(
+                    this,
+                    paletteSize,
+                    planeType,
+                    colorContext,
+                    colorOrderIndex);
+            }
+        }
+
+        return cost;
+    }
+
     /// <summary>
     /// Separates palette colors selected from the neighbor cache from colors that require literal coding.
     /// </summary>
@@ -1738,6 +1868,61 @@ internal class Av1SymbolEncoder : IDisposable
             => Av1ProbabilityCost.GetLiteralCost(bitCount);
     }
 
+    /// <summary>
+    /// Emits palette-map syntax and reports no estimated rate.
+    /// </summary>
+    private readonly struct PaletteColorMapWriteOperation : IPaletteColorMapOperation
+    {
+        public static int ProcessFirstIndex(
+            Av1SymbolEncoder encoder,
+            int paletteSize,
+            int colorIndex)
+        {
+            encoder.WriteUniform(paletteSize, colorIndex);
+            return 0;
+        }
+
+        public static int ProcessColorIndex(
+            Av1SymbolEncoder encoder,
+            int paletteSize,
+            Av1PlaneType planeType,
+            int colorContext,
+            int colorOrderIndex)
+        {
+            encoder.WritePaletteColorIndex(
+                colorOrderIndex,
+                paletteSize,
+                colorContext,
+                planeType);
+
+            return 0;
+        }
+    }
+
+    /// <summary>
+    /// Measures palette-map syntax against the live tile distributions without changing them.
+    /// </summary>
+    private readonly struct PaletteColorMapCostOperation : IPaletteColorMapOperation
+    {
+        public static int ProcessFirstIndex(
+            Av1SymbolEncoder encoder,
+            int paletteSize,
+            int colorIndex)
+            => GetUniformCost(paletteSize, colorIndex);
+
+        public static int ProcessColorIndex(
+            Av1SymbolEncoder encoder,
+            int paletteSize,
+            Av1PlaneType planeType,
+            int colorContext,
+            int colorOrderIndex)
+            => encoder.GetPaletteColorIndexCost(
+                colorOrderIndex,
+                paletteSize,
+                colorContext,
+                planeType);
+    }
+
     /// <summary>
     /// Owns the adaptive distributions used only by AV1 palette syntax.
     /// </summary>

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

@@ -0,0 +1,101 @@
+// Copyright (c) Six Labors.
+// Licensed under the Six Labors Split License.
+
+using SixLabors.ImageSharp.Memory;
+
+namespace SixLabors.ImageSharp.Formats.Heif.Av1.Tiling;
+
+/// <summary>
+/// Derives AV1 palette color-index ordering and entropy contexts from preceding spatial indices.
+/// </summary>
+internal static class Av1PaletteColorMap
+{
+    /// <summary>
+    /// Derives the palette color order, current color-order index, and entropy context for one map position.
+    /// </summary>
+    /// <param name="colorIndexMap">The partially or completely populated color-index map.</param>
+    /// <param name="row">The current map row.</param>
+    /// <param name="column">The current map column.</param>
+    /// <param name="paletteSize">The number of colors in the palette.</param>
+    /// <param name="colorIndex">The current palette index, or a negative value while decoding.</param>
+    /// <param name="colorOrder">The destination color order for the current context.</param>
+    /// <param name="colorOrderIndex">The current index in <paramref name="colorOrder" />, or a negative value while decoding.</param>
+    /// <returns>The color-index entropy context in the range from zero through four.</returns>
+    public static int GetContext(
+        Buffer2DRegion<byte> colorIndexMap,
+        int row,
+        int column,
+        int paletteSize,
+        int colorIndex,
+        Span<byte> colorOrder,
+        out int colorOrderIndex)
+    {
+        Span<int> scores = stackalloc int[Av1Constants.PaletteMaxSize];
+        scores.Clear();
+        ReadOnlySpan<byte> currentRow = colorIndexMap.DangerousGetRowSpan(row);
+        if (column > 0)
+        {
+            scores[currentRow[column - 1]] += 2;
+        }
+
+        if (row > 0)
+        {
+            ReadOnlySpan<byte> aboveRow = colorIndexMap.DangerousGetRowSpan(row - 1);
+            if (column > 0)
+            {
+                scores[aboveRow[column - 1]]++;
+            }
+
+            scores[aboveRow[column]] += 2;
+        }
+
+        Span<int> inverseColorOrder = stackalloc int[Av1Constants.PaletteMaxSize];
+        for (int i = 0; i < Av1Constants.PaletteMaxSize; i++)
+        {
+            colorOrder[i] = (byte)i;
+            inverseColorOrder[i] = i;
+        }
+
+        // Stable descending score order keeps lower palette indices ahead when neighboring scores tie.
+        for (int i = 0; i < 3; i++)
+        {
+            int maximumScore = scores[i];
+            int maximumIndex = i;
+            for (int j = i + 1; j < paletteSize; j++)
+            {
+                if (scores[j] > maximumScore)
+                {
+                    maximumScore = scores[j];
+                    maximumIndex = j;
+                }
+            }
+
+            if (maximumIndex != i)
+            {
+                byte maximumColor = colorOrder[maximumIndex];
+                for (int j = maximumIndex; j > i; j--)
+                {
+                    scores[j] = scores[j - 1];
+                    colorOrder[j] = colorOrder[j - 1];
+                    inverseColorOrder[colorOrder[j]] = j;
+                }
+
+                scores[i] = maximumScore;
+                colorOrder[i] = maximumColor;
+                inverseColorOrder[maximumColor] = i;
+            }
+        }
+
+        colorOrderIndex = colorIndex < 0 ? -1 : inverseColorOrder[colorIndex];
+        int contextHash = scores[0] + (2 * scores[1]) + (2 * scores[2]);
+        return contextHash switch
+        {
+            2 => 0,
+            5 => 4,
+            6 => 3,
+            7 => 2,
+            8 => 1,
+            _ => -1
+        };
+    }
+}

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

@@ -423,11 +423,6 @@ internal sealed class Av1TileReader : IAv1TileReader, IDisposable
         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,
         2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2];
 
-    /// <summary>
-    /// Gets the color-index entropy context for each weighted palette-neighbor score hash.
-    /// </summary>
-    private static ReadOnlySpan<int> PaletteColorIndexContexts => [-1, -1, 0, -1, -1, 4, 3, 2, 1];
-
     /// <summary>
     /// Gets the frame header whose tile syntax is being parsed.
     /// </summary>
@@ -2814,26 +2809,7 @@ internal sealed class Av1TileReader : IAv1TileReader, IDisposable
         int columns,
         Buffer2DRegion<byte> colorIndexMap)
     {
-        colorIndexMap.DangerousGetRowSpan(0)[0] = (byte)reader.ReadUniform(paletteSize);
+        reader.ReadPaletteColorMap(paletteSize, planeType, rows, columns, colorIndexMap);
-        Span<byte> colorOrder = stackalloc byte[Av1Constants.PaletteMaxSize];
-        for (int diagonal = 1; diagonal < rows + columns - 1; diagonal++)
-        {
-            int firstColumn = Math.Min(diagonal, columns - 1);
-            int lastColumn = Math.Max(0, diagonal - rows + 1);
-            for (int column = firstColumn; column >= lastColumn; column--)
-            {
-                int row = diagonal - column;
-                int colorContext = GetPaletteColorIndexContext(
-                    colorIndexMap,
-                    row,
-                    column,
-                    paletteSize,
-                    colorOrder);
-
-                int colorOrderIndex = reader.ReadPaletteColorIndex(paletteSize, colorContext, planeType);
-                colorIndexMap.DangerousGetRowSpan(row)[column] = colorOrder[colorOrderIndex];
-            }
-        }
 
         if (columns < planeWidth)
         {
@@ -2854,84 +2830,6 @@ internal sealed class Av1TileReader : IAv1TileReader, IDisposable
         }
     }
 
-    /// <summary>
-    /// Derives the palette color order and entropy context from the left, upper-left, and above indices.
-    /// </summary>
-    /// <param name="colorIndexMap">The partially decoded color-index map.</param>
-    /// <param name="row">The current map row.</param>
-    /// <param name="column">The current map column.</param>
-    /// <param name="paletteSize">The number of palette colors.</param>
-    /// <param name="colorOrder">The destination color order for the current context.</param>
-    /// <returns>The color-index entropy context in the range from zero through four.</returns>
-    private static int GetPaletteColorIndexContext(
-        Buffer2DRegion<byte> colorIndexMap,
-        int row,
-        int column,
-        int paletteSize,
-        Span<byte> colorOrder)
-    {
-        Span<int> neighborColors = stackalloc int[3];
-        ReadOnlySpan<byte> currentRow = colorIndexMap.DangerousGetRowSpan(row);
-        neighborColors[0] = column > 0 ? currentRow[column - 1] : -1;
-        if (row > 0)
-        {
-            ReadOnlySpan<byte> aboveRow = colorIndexMap.DangerousGetRowSpan(row - 1);
-            neighborColors[1] = column > 0 ? aboveRow[column - 1] : -1;
-            neighborColors[2] = aboveRow[column];
-        }
-        else
-        {
-            neighborColors[1] = -1;
-            neighborColors[2] = -1;
-        }
-
-        Span<int> scores = stackalloc int[Av1Constants.PaletteMaxSize];
-        ReadOnlySpan<int> neighborWeights = [2, 1, 2];
-        for (int i = 0; i < neighborColors.Length; i++)
-        {
-            if (neighborColors[i] >= 0)
-            {
-                scores[neighborColors[i]] += neighborWeights[i];
-            }
-        }
-
-        for (int i = 0; i < colorOrder.Length; i++)
-        {
-            colorOrder[i] = (byte)i;
-        }
-
-        // Stable descending score order keeps lower palette indices ahead when neighboring scores tie.
-        for (int i = 0; i < 3; i++)
-        {
-            int maximumScore = scores[i];
-            int maximumIndex = i;
-            for (int j = i + 1; j < paletteSize; j++)
-            {
-                if (scores[j] > maximumScore)
-                {
-                    maximumScore = scores[j];
-                    maximumIndex = j;
-                }
-            }
-
-            if (maximumIndex != i)
-            {
-                byte maximumColor = colorOrder[maximumIndex];
-                for (int j = maximumIndex; j > i; j--)
-                {
-                    scores[j] = scores[j - 1];
-                    colorOrder[j] = colorOrder[j - 1];
-                }
-
-                scores[i] = maximumScore;
-                colorOrder[i] = maximumColor;
-            }
-        }
-
-        int contextHash = scores[0] + (2 * scores[1]) + (2 * scores[2]);
-        return PaletteColorIndexContexts[contextHash];
-    }
-
     /// <summary>
     /// Reads the joint signs and nonzero alpha magnitudes for chroma-from-luma prediction.
     /// </summary>

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

@@ -1233,6 +1233,173 @@ public class Av1EntropyTests
         Assert.Equal(before, after);
     }
 
+    [Fact]
+    public void PaletteColorMapCostMatchesCurrentDistributions()
+    {
+        using Buffer2D<byte> map = Configuration.Default.MemoryAllocator.Allocate2D<byte>(2, 2, AllocationOptions.Clean);
+        map.DangerousGetRowSpan(0)[0] = 2;
+        map.DangerousGetRowSpan(0)[1] = 0;
+        map.DangerousGetRowSpan(1)[0] = 1;
+        map.DangerousGetRowSpan(1)[1] = 2;
+        Buffer2DRegion<byte> region = new(map);
+        Av1Distribution[][] distributions = Av1DefaultDistributions.PaletteYColorIndex;
+        int expected = Av1SymbolEncoder.GetUniformCost(3, 2);
+        expected += Av1ProbabilityCost.GetSymbolCost(distributions[1][0], 1);
+        expected += Av1ProbabilityCost.GetSymbolCost(distributions[1][0], 2);
+        expected += Av1ProbabilityCost.GetSymbolCost(distributions[1][1], 2);
+        using Av1SymbolEncoder encoder = new(Configuration.Default, 64, BaseQIndex, updateCdf: false);
+
+        Assert.Equal(
+            expected,
+            encoder.GetPaletteColorMapCost(3, Av1PlaneType.Y, 2, 2, region));
+    }
+
+    [Theory]
+    [InlineData(1, 1, 0, 0, 0, 1, 4, 1)]
+    [InlineData(1, 1, 0, 1, 0, 1, 3, 1)]
+    [InlineData(1, 1, 0, 0, 1, 1, 2, 1)]
+    [InlineData(1, 1, 0, 2, 1, 2, 1, 2)]
+    [InlineData(0, 1, 0, 2, 0, 0, 0, 1)]
+    public void PaletteColorMapContextMatchesCurrentLibaom(
+        int row,
+        int column,
+        byte left,
+        byte upperLeft,
+        byte above,
+        byte current,
+        int expectedContext,
+        int expectedOrderIndex)
+    {
+        using Buffer2D<byte> map = Configuration.Default.MemoryAllocator.Allocate2D<byte>(2, 2, AllocationOptions.Clean);
+        map.DangerousGetRowSpan(0)[0] = upperLeft;
+        map.DangerousGetRowSpan(0)[1] = above;
+        map.DangerousGetRowSpan(1)[0] = left;
+        map.DangerousGetRowSpan(row)[column] = current;
+        Span<byte> colorOrder = stackalloc byte[Av1Constants.PaletteMaxSize];
+
+        int actualContext = Av1PaletteColorMap.GetContext(
+            new Buffer2DRegion<byte>(map),
+            row,
+            column,
+            paletteSize: 4,
+            current,
+            colorOrder,
+            out int actualOrderIndex);
+
+        Assert.Equal(expectedContext, actualContext);
+        Assert.Equal(expectedOrderIndex, actualOrderIndex);
+    }
+
+    [Fact]
+    public void RoundTripPaletteColorMaps()
+    {
+        const int Rows = 5;
+        const int Columns = 7;
+        const int Width = 9;
+        const int Height = 6;
+        Configuration configuration = Configuration.Default;
+        using Buffer2D<byte> source = configuration.MemoryAllocator.Allocate2D<byte>(Width, Height);
+        using Buffer2D<byte> decoded = configuration.MemoryAllocator.Allocate2D<byte>(Width, Height);
+        Buffer2DRegion<byte> sourceRegion = new(source);
+        Buffer2DRegion<byte> decodedRegion = new(decoded);
+        using Av1SymbolEncoder encoder = new(configuration, 512, BaseQIndex);
+        for (int paletteSize = 2; paletteSize <= Av1Constants.PaletteMaxSize; paletteSize++)
+        {
+            for (int plane = 0; plane < 2; plane++)
+            {
+                for (int row = 0; row < Rows; row++)
+                {
+                    Span<byte> sourceRow = source.DangerousGetRowSpan(row);
+                    for (int column = 0; column < Columns; column++)
+                    {
+                        sourceRow[column] = (byte)(((row * 3) + (column * 5) + plane) % paletteSize);
+                    }
+                }
+
+                encoder.WritePaletteColorMap(
+                    paletteSize,
+                    (Av1PlaneType)plane,
+                    Rows,
+                    Columns,
+                    sourceRegion);
+            }
+        }
+
+        using IMemoryOwner<byte> encoded = encoder.Exit();
+        Av1SymbolDecoder decoder = new(configuration, encoded.GetSpan(), BaseQIndex);
+        for (int paletteSize = 2; paletteSize <= Av1Constants.PaletteMaxSize; paletteSize++)
+        {
+            for (int plane = 0; plane < 2; plane++)
+            {
+                for (int row = 0; row < Height; row++)
+                {
+                    decoded.DangerousGetRowSpan(row).Fill(byte.MaxValue);
+                }
+
+                decoder.ReadPaletteColorMap(
+                    paletteSize,
+                    (Av1PlaneType)plane,
+                    Rows,
+                    Columns,
+                    decodedRegion);
+
+                for (int row = 0; row < Rows; row++)
+                {
+                    ReadOnlySpan<byte> decodedRow = decoded.DangerousGetRowSpan(row);
+                    for (int column = 0; column < Columns; column++)
+                    {
+                        Assert.Equal(
+                            (byte)(((row * 3) + (column * 5) + plane) % paletteSize),
+                            decodedRow[column]);
+                    }
+
+                    for (int column = Columns; column < Width; column++)
+                    {
+                        Assert.Equal(byte.MaxValue, decodedRow[column]);
+                    }
+                }
+
+                for (int row = Rows; row < Height; row++)
+                {
+                    ReadOnlySpan<byte> decodedRow = decoded.DangerousGetRowSpan(row);
+                    for (int column = 0; column < Width; column++)
+                    {
+                        Assert.Equal(byte.MaxValue, decodedRow[column]);
+                    }
+                }
+            }
+        }
+
+        decoder.ValidateTrailingBits();
+    }
+
+    [Fact]
+    public void PaletteColorMapCostDoesNotAllocateAfterEntropyInitialization()
+    {
+        using Buffer2D<byte> map = Configuration.Default.MemoryAllocator.Allocate2D<byte>(7, 5, AllocationOptions.Clean);
+        for (int row = 0; row < map.Height; row++)
+        {
+            Span<byte> mapRow = map.DangerousGetRowSpan(row);
+            for (int column = 0; column < map.Width; column++)
+            {
+                mapRow[column] = (byte)(((row * 3) + (column * 5)) % 4);
+            }
+        }
+
+        using Av1SymbolEncoder encoder = new(Configuration.Default, 64, BaseQIndex, updateCdf: false);
+        Buffer2DRegion<byte> region = new(map);
+        _ = encoder.GetPaletteColorMapCost(4, Av1PlaneType.Y, map.Height, map.Width, region);
+        long before = GC.GetAllocatedBytesForCurrentThread();
+
+        for (int i = 0; i < 1000; i++)
+        {
+            _ = encoder.GetPaletteColorMapCost(4, Av1PlaneType.Y, map.Height, map.Width, region);
+        }
+
+        long after = GC.GetAllocatedBytesForCurrentThread();
+        Assert.Equal(before, after);
+    }
+
     [Theory]
     [MemberData(nameof(GetRangeData), 20)]
     public void RoundTripPartitionType(int context)