Fix reading of larger literals

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

@@ -1,19 +1,35 @@
 // Copyright (c) Six Labors.
 // Licensed under the Six Labors Split License.
 
+using System;
 using System.Buffers.Binary;
+using System.Runtime.InteropServices;
 
 namespace SixLabors.ImageSharp.Formats.Heif.Av1;
 
-internal ref struct Av1BitStreamReader(Span<byte> data)
+internal ref struct Av1BitStreamReader
 {
-    private const int WordSize = sizeof(byte) * 8;
-    private readonly Span<byte> data = data;
+    private const int WordSize = 32;
+
+    private readonly Span<uint> data;
+    private uint currentWord;
+    private uint nextWord;
     private int wordPosition = 0;
     private int bitOffset = 0;
 
-    public readonly int BitPosition => (this.wordPosition * WordSize) + this.bitOffset;
+    public Av1BitStreamReader(Span<byte> data)
+    {
+        this.data = MemoryMarshal.Cast<byte, uint>(data);
+        this.wordPosition = -1;
+        this.AdvanceToNextWord();
+        this.AdvanceToNextWord();
+    }
 
+    public readonly int BitPosition => ((this.wordPosition - 1) * WordSize) + this.bitOffset;
+
+    /// <summary>
+    /// Gets the number of bytes in the readers buffer.
+    /// </summary>
     public readonly int Length => this.data.Length;
 
     public void Reset()
@@ -34,22 +50,34 @@ internal ref struct Av1BitStreamReader(Span<byte> data)
 
     public uint ReadLiteral(int bitCount)
     {
-        uint bits = BinaryPrimitives.ReadUInt32BigEndian(this.data[this.wordPosition..]);
-        this.Skip(bitCount);
-        return bits >> (32 - bitCount);
-    }
+        DebugGuard.MustBeBetweenOrEqualTo(bitCount, 1, 32, nameof(bitCount));
 
-    internal bool ReadBoolean()
-    {
-        bool bit = (this.data[this.wordPosition] & (1 << (WordSize - this.bitOffset - 1))) > 0;
-        this.Skip(1);
-        return bit;
+        uint bits = (this.currentWord << this.bitOffset) >> (WordSize - bitCount);
+        this.bitOffset += bitCount;
+        if (this.bitOffset > WordSize)
+        {
+            int overshoot = WordSize + WordSize - this.bitOffset;
+            if (overshoot < 32)
+            {
+                bits |= this.nextWord >> overshoot;
+            }
+        }
+
+        if (this.bitOffset >= WordSize)
+        {
+            this.AdvanceToNextWord();
+            this.bitOffset -= WordSize;
+        }
+
+        return bits;
     }
 
+    internal bool ReadBoolean() => this.ReadLiteral(1) > 0;
+
     public ulong ReadLittleEndianBytes128(out int length)
     {
         // See section 4.10.5 of the AV1-Specification
-        DebugGuard.IsTrue((this.bitOffset & (WordSize - 1)) == 0, "Reading of Little Endian 128 value only allowed on byte alignment");
+        DebugGuard.IsTrue((this.bitOffset & 0xF7) == 0, "Reading of Little Endian 128 value only allowed on byte alignment");
 
         ulong value = 0;
         length = 0;
@@ -58,7 +86,7 @@ internal ref struct Av1BitStreamReader(Span<byte> data)
             uint leb128Byte = this.ReadLiteral(8);
             value |= (leb128Byte & 0x7FUL) << i;
             length++;
-            if ((leb128Byte & 0x80U) != 0x80U)
+            if ((leb128Byte & 0x80U) == 0)
             {
                 break;
             }
@@ -137,4 +165,12 @@ internal ref struct Av1BitStreamReader(Span<byte> data)
 
         return t;
     }
+
+    public void AdvanceToNextWord()
+    {
+        this.currentWord = this.nextWord;
+        this.wordPosition++;
+        uint temp = this.data[this.wordPosition];
+        this.nextWord = (temp << 24) | ((temp & 0x0000ff00U) << 8) | ((temp & 0x00ff0000U) >> 8) | (temp >> 24);
+    }
 }

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

@@ -12,11 +12,11 @@ public class Av1BitsStreamTests
     [Theory]
     [InlineData(42, new bool[] { false, false, true, false, true, false, true, false })]
     [InlineData(52, new bool[] { false, false, true, true, false, true, false, false })]
-    public void ReadAsBoolean(int value, bool[] bits)
+    public void ReadAsBoolean(byte value, bool[] bits)
     {
         int bitCount = bits.Length;
-        byte[] buffer = new byte[4];
-        BinaryPrimitives.WriteInt32BigEndian(buffer, value << (32 - bitCount));
+        byte[] buffer = new byte[8];
+        buffer[0] = value;
         Av1BitStreamReader reader = new(buffer);
         bool[] actual = new bool[bitCount];
         for (int i = 0; i < bitCount; i++)
@@ -34,7 +34,7 @@ public class Av1BitsStreamTests
     [InlineData(4050, 16)]
     public void ReadAsLiteral(uint expected, int bitCount)
     {
-        byte[] buffer = new byte[4];
+        byte[] buffer = new byte[8];
         BinaryPrimitives.WriteUInt32BigEndian(buffer, expected << (32 - bitCount));
         Av1BitStreamReader reader = new(buffer);
         uint actual = reader.ReadLiteral(bitCount);
@@ -46,7 +46,7 @@ public class Av1BitsStreamTests
     [InlineData(new bool[] { false, true, false, true })]
     public void WriteAsBoolean(bool[] booleans)
     {
-        using MemoryStream stream = new(4);
+        using MemoryStream stream = new(8);
         Av1BitStreamWriter writer = new(stream);
         for (int i = 0; i < booleans.Length; i++)
         {
@@ -73,7 +73,7 @@ public class Av1BitsStreamTests
     [InlineData(4050, 16)]
     public void WriteAsLiteral(uint value, int bitCount)
     {
-        using MemoryStream stream = new(4);
+        using MemoryStream stream = new(8);
         Av1BitStreamWriter writer = new(stream);
         writer.WriteLiteral(value, bitCount);
         writer.Flush();