Fix bitstream writer

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

@@ -33,42 +33,28 @@ internal ref struct Av1BitStreamWriter(Stream stream)
 
     public void WriteLiteral(uint value, int bitCount)
     {
-        int shift = 24;
-        uint padded = value << ((32 - bitCount) - this.bitOffset);
-        while ((bitCount + this.bitOffset) >= 8)
+        for (int bit = bitCount - 1; bit >= 0; bit--)
         {
-            byte current = (byte)(((padded >> shift) & 0xff) | this.buffer);
-            this.stream.WriteByte(current);
-            shift -= 8;
-            bitCount -= 8;
-            this.buffer = 0;
-            this.bitOffset = 0;
-        }
-
-        if (bitCount > 0)
-        {
-            this.buffer = (byte)(((padded >> shift) & 0xff) | this.buffer);
-            this.bitOffset += bitCount;
+            this.WriteBit((byte)((value >> bit) & 0x1));
         }
     }
 
     internal void WriteBoolean(bool value)
     {
         byte boolByte = value ? (byte)1 : (byte)0;
-        this.buffer = (byte)(((boolByte << (7 - this.bitOffset)) & 0xff) | this.buffer);
-        this.bitOffset++;
-        if (this.bitOffset == WordSize)
-        {
-            this.stream.WriteByte(this.buffer);
-            this.bitOffset = 0;
-        }
+        this.WriteBit(boolByte);
     }
 
     public void WriteSignedFromUnsigned(int signedValue, int n)
     {
         // See section 4.10.6 of the AV1-Specification
-        uint value = unchecked((uint)signedValue);
-        this.WriteLiteral(value, n + 1);
+        ulong value = (ulong)signedValue;
+        if (signedValue < 0)
+        {
+            value += 1UL << n;
+        }
+
+        this.WriteLiteral((uint)value, n);
     }
 
     public void WriteLittleEndianBytes128(uint value)
@@ -79,12 +65,47 @@ internal ref struct Av1BitStreamWriter(Stream stream)
         }
         else if (value < 0x8000U)
         {
+            this.WriteLiteral((value & 0x7FU) | 0x80U, 8);
             this.WriteLiteral(value >> 7, 8);
-            this.WriteLiteral(value & 0x80, 8);
         }
         else
         {
             throw new NotImplementedException("No such large values yet.");
         }
     }
+
+    internal void WriteNonSymmetric(uint value, uint numberOfSymbols)
+    {
+        // See section 4.10.7 of the AV1-Specification
+        if (numberOfSymbols <= 1)
+        {
+            return;
+        }
+
+        int w = (int)(Av1Math.FloorLog2(numberOfSymbols) + 1);
+        uint m = (uint)((1 << w) - numberOfSymbols);
+        if (value < m)
+        {
+            this.WriteLiteral(value, w - 1);
+        }
+        else
+        {
+            uint extraBit = ((value + m) >> 1) - value;
+            uint k = (value + m - extraBit) >> 1;
+            this.WriteLiteral(k, w - 1);
+            this.WriteLiteral(extraBit, 1);
+        }
+    }
+
+    private void WriteBit(byte value)
+    {
+        this.buffer = (byte)(((value << (7 - this.bitOffset)) & 0xff) | this.buffer);
+        this.bitOffset++;
+        if (this.bitOffset == WordSize)
+        {
+            this.stream.WriteByte(this.buffer);
+            this.buffer = 0;
+            this.bitOffset = 0;
+        }
+    }
 }

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

@@ -83,4 +83,186 @@ public class Av1BitsStreamTests
         uint actual = reader.ReadLiteral(bitCount);
         Assert.Equal(value, actual);
     }
+
+    [Theory]
+    [InlineData(3)]
+    [InlineData(4)]
+    [InlineData(5)]
+    [InlineData(6)]
+    [InlineData(7)]
+    [InlineData(8)]
+    [InlineData(16)]
+    public void ReadLiteralRainbowArray(int bitCount)
+    {
+        uint[] values = Enumerable.Range(0, (1 << bitCount) - 1).Select(i => (uint)i).ToArray();
+        using MemoryStream stream = new(280);
+        Av1BitStreamWriter writer = new(stream);
+        for (int i = 0; i < values.Length; i++)
+        {
+            writer.WriteLiteral(values[i], bitCount);
+        }
+
+        writer.Flush();
+
+        // Read the written value back.
+        Av1BitStreamReader reader = new(stream.GetBuffer());
+        uint[] actuals = new uint[values.Length];
+        for (int i = 0; i < values.Length; i++)
+        {
+            uint actual = reader.ReadLiteral(bitCount);
+            actuals[i] = actual;
+        }
+
+        Assert.Equal(values, actuals);
+    }
+
+    [Theory]
+    [InlineData(4, 6, 4, 9, 14)]
+    [InlineData(8, 42, 8, 189, 63)]
+    [InlineData(8, 52, 18, 255, 241)]
+    [InlineData(16, 4050, 16003, 503, 814)]
+    public void ReadWriteAsLiteralArray(int bitCount, uint val1, uint val2, uint val3, uint val4)
+    {
+        uint[] values = [val1, val2, val3, val4];
+        using MemoryStream stream = new(80);
+        Av1BitStreamWriter writer = new(stream);
+        for (int i = 0; i < values.Length; i++)
+        {
+            writer.WriteLiteral(values[i], bitCount);
+        }
+
+        writer.Flush();
+
+        // Read the written value back.
+        Av1BitStreamReader reader = new(stream.GetBuffer());
+        for (int i = 0; i < values.Length; i++)
+        {
+            uint actual = reader.ReadLiteral(bitCount);
+            Assert.NotEqual(0U, actual);
+            Assert.Equal(values[i], actual);
+        }
+    }
+
+    [Theory]
+
+    [InlineData(4, 0, 1, 2, 3)]
+    [InlineData(5, 0, 1, 2, 3)]
+    [InlineData(8, 0, 1, 2, 3)]
+    [InlineData(8, 4, 5, 6, 7)]
+    [InlineData(16, 15, 0, 5, 8)]
+    public void ReadWriteAsNonSymmetricArray(uint numberOfSymbols, uint val1, uint val2, uint val3, uint val4)
+    {
+        uint[] values = [val1, val2, val3, val4];
+        using MemoryStream stream = new(80);
+        Av1BitStreamWriter writer = new(stream);
+        for (int i = 0; i < values.Length; i++)
+        {
+            writer.WriteNonSymmetric(values[i], numberOfSymbols);
+        }
+
+        writer.Flush();
+
+        // Read the written value back.
+        Av1BitStreamReader reader = new(stream.GetBuffer());
+        uint[] actuals = new uint[4];
+        for (int i = 0; i < values.Length; i++)
+        {
+            ulong actual = reader.ReadNonSymmetric(numberOfSymbols);
+            actuals[i] = (uint)actual;
+            // Assert.NotEqual(0UL, actual);
+        }
+
+        Assert.Equal(values, actuals);
+    }
+
+    [Theory]
+    [InlineData(3)]
+    [InlineData(4)]
+    [InlineData(5)]
+    [InlineData(7)]
+    [InlineData(8)]
+    public void ReadSignedRainbowArray(int bitCount)
+    {
+        int maxValue = (1 << (bitCount - 1)) - 1;
+        int[] values = Enumerable.Range(-maxValue, maxValue).ToArray();
+        using MemoryStream stream = new(280);
+        Av1BitStreamWriter writer = new(stream);
+        for (int i = 0; i < values.Length; i++)
+        {
+            writer.WriteSignedFromUnsigned(values[i], bitCount);
+        }
+
+        writer.Flush();
+
+        // Read the written value back.
+        Av1BitStreamReader reader = new(stream.GetBuffer());
+        int[] actuals = new int[values.Length];
+        for (int i = 0; i < values.Length; i++)
+        {
+            int actual = reader.ReadSignedFromUnsigned(bitCount);
+            actuals[i] = actual;
+        }
+
+        Assert.Equal(values, actuals);
+    }
+
+    [Theory]
+    [InlineData(5, 6, 4, -7, -2)]
+    [InlineData(7, 26, -8, -19, -26)]
+    [InlineData(8, 52, 127, -127, -21)]
+    [InlineData(16, -4050, -16003, -503, 8414)]
+    public void ReadWriteSignedArray(int bitCount, int val1, int val2, int val3, int val4)
+    {
+        int[] values = [val1, val2, val3, val4];
+        using MemoryStream stream = new(80);
+        Av1BitStreamWriter writer = new(stream);
+        for (int i = 0; i < values.Length; i++)
+        {
+            writer.WriteSignedFromUnsigned(values[i], bitCount);
+        }
+
+        writer.Flush();
+
+        // Read the written value back.
+        Av1BitStreamReader reader = new(stream.GetBuffer());
+        int[] actuals = new int[4];
+        for (int i = 0; i < values.Length; i++)
+        {
+            int actual = reader.ReadSignedFromUnsigned(bitCount);
+            actuals[i] = actual;
+            // Assert.NotEqual(0, actual);
+        }
+
+        Assert.Equal(values, actuals);
+    }
+
+    [Theory]
+    [InlineData(4, 6, 4, 9, 14)]
+    [InlineData(8, 42, 8, 189, 63)]
+    [InlineData(8, 52, 18, 255, 241)]
+    [InlineData(16, 4050, 16003, 503, 8414)]
+    public void ReadWriteLittleEndianBytes128Array(uint val0, uint val1, uint val2, uint val3, uint val4)
+    {
+        uint[] values = [val0, val1, val2, val3, val4];
+        using MemoryStream stream = new(80);
+        Av1BitStreamWriter writer = new(stream);
+        for (int i = 0; i < values.Length; i++)
+        {
+            writer.WriteLittleEndianBytes128(values[i]);
+        }
+
+        writer.Flush();
+
+        // Read the written value back.
+        Av1BitStreamReader reader = new(stream.GetBuffer());
+        uint[] actuals = new uint[5];
+        for (int i = 0; i < values.Length; i++)
+        {
+            ulong actual = reader.ReadLittleEndianBytes128(out int length);
+            actuals[i] = (uint)actual;
+            Assert.NotEqual(0UL, actual);
+        }
+
+        Assert.Equal(values, actuals);
+    }
 }