Fix pad_to_size to work correctly misaligned

Author: KredeGC

include/bitstream/stream/bit_reader.h

@@ -178,24 +178,23 @@ namespace bitstream
 
 			BS_ASSERT(num_bytes * 8U >= m_NumBitsRead);
 
-			uint32_t offset = m_NumBitsRead / 32;
+			uint32_t remainder = (num_bytes * 8U - m_NumBitsRead) % 32U;
 			uint32_t zero;
 
-			// Test for zeros in padding
-			for (uint32_t i = offset; i < num_bytes / 4; i++)
+			// Test the last word more carefully, as it may have data
+			if (remainder != 0U)
 			{
-				bool status = serialize_bits(zero, 32);
-
+				bool status = serialize_bits(zero, remainder);
 				BS_ASSERT(status && zero == 0);
 			}
 
-			uint32_t remainder = num_bytes * 8U - m_NumBitsRead;
+			uint32_t offset = m_NumBitsRead / 32;
+            uint32_t max = num_bytes / 4;
 
-			// Test the last word more carefully, as it may have data
-			if (remainder % 32U != 0U)
+			// Test for zeros in padding
+			for (uint32_t i = offset; i < max; i++)
 			{
-				bool status = serialize_bits(zero, remainder);
-
+				bool status = serialize_bits(zero, 32);
 				BS_ASSERT(status && zero == 0);
 			}
 

include/bitstream/stream/bit_writer.h

@@ -214,20 +214,33 @@ namespace bitstream
 			BS_ASSERT(num_bytes * 8U <= m_TotalBits);
 
 			BS_ASSERT(num_bytes * 8U >= m_NumBitsWritten);
+            
+            if (m_NumBitsWritten == 0)
+            {
+                std::memset(m_Buffer, 0, num_bytes);
+                
+                m_NumBitsWritten = num_bytes * 8;
+                m_Scratch = 0;
+                m_ScratchBits = 0;
+                m_WordIndex = num_bytes / 4;
+                
+                return true;
+            }
+            
+            uint32_t remainder = (num_bytes * 8U - m_NumBitsWritten) % 32U;
+			uint32_t zero = 0;
+
+			// Align to byte
+			if (remainder != 0U)
+				BS_ASSERT(serialize_bits(zero, remainder));
 
 			uint32_t offset = m_NumBitsWritten / 32;
-			uint32_t zero = 0;
+            uint32_t max = num_bytes / 4;
 
 			// Serialize words
-			for (uint32_t i = offset; i < num_bytes / 4; i++)
+			for (uint32_t i = offset; i < max; i++)
 				BS_ASSERT(serialize_bits(zero, 32));
 
-			uint32_t remainder = num_bytes * 8U - m_NumBitsWritten;
-
-			// Align to byte
-			if (remainder % 32U != 0U)
-				BS_ASSERT(serialize_bits(zero, remainder));
-
 			return true;
 		}
 

src/test/stream_test.cpp

@@ -155,6 +155,46 @@ namespace bitstream::test::stream
 		BS_TEST_ASSERT(out_value2 == in_value2);
 	}
 
+	BS_ADD_TEST(test_serialize_padding_empty)
+	{
+		// Write some initial value, pad it, and then write another value
+		byte_buffer<32> buffer;
+		bit_writer writer(buffer);
+		BS_TEST_ASSERT(writer.pad_to_size(30));
+		uint32_t num_bytes = writer.flush();
+
+		BS_TEST_ASSERT(num_bytes == 30);
+
+		// Read the values and validate padding
+		bit_reader reader(buffer, num_bytes);
+
+		BS_TEST_ASSERT(reader.pad_to_size(30));
+	}
+
+	BS_ADD_TEST(test_serialize_padding_full)
+	{
+		// Test padding
+		uint32_t in_value = 3;
+
+		// Write some initial value, pad it, and then write another value
+		byte_buffer<32> buffer;
+		bit_writer writer(buffer);
+		BS_TEST_ASSERT(writer.serialize_bits(in_value, 3));
+		BS_TEST_ASSERT(writer.pad_to_size(32));
+		uint32_t num_bytes = writer.flush();
+
+		BS_TEST_ASSERT(num_bytes == 32);
+
+		// Read the values and validate padding
+		uint32_t out_value;
+		bit_reader reader(buffer, num_bytes);
+
+		BS_TEST_ASSERT(reader.serialize_bits(out_value, 3));
+		BS_TEST_ASSERT(reader.pad_to_size(32));
+
+		BS_TEST_ASSERT(out_value == in_value);
+	}
+
 	void test_serialize_align()
 	{
 		// Test align