Fix pad_to_size to work correctly with bytes less than 4

Fix bit_writer::flush to return 0 when the buffer is empty

Author: KredeGC

include/bitstream/stream/bit_reader.h

@@ -139,21 +139,24 @@ namespace bitstream
 		*/
 		bool serialize_checksum(uint32_t protocol_version) noexcept
 		{
+			BS_ASSERT(m_NumBitsRead == 0);
+
+			BS_ASSERT(can_serialize_bits(32U));
+
 			uint32_t num_bytes = (m_TotalBits - 1U) / 8U + 1U;
 
 			// Read the checksum
-			uint32_t checksum;
-			std::memcpy(&checksum, m_Buffer, sizeof(uint32_t));
+			uint32_t checksum = *m_Buffer;
 
 			// Copy protocol version to buffer
 			uint32_t* buffer = const_cast<uint32_t*>(m_Buffer); // Somewhat of a hack, but it's faster to change the checksum twice than allocate memory for it
-			std::memcpy(buffer, &protocol_version, sizeof(uint32_t));
+			*buffer = protocol_version;
 
 			// Generate checksum to compare against
 			uint32_t generated_checksum = utility::crc_uint32(reinterpret_cast<uint8_t*>(buffer), num_bytes);
 
 			// Write the checksum back, just in case
-			std::memcpy(buffer, &checksum, sizeof(uint32_t));
+			*buffer = checksum;
 
 			// Advance the reader by the size of the checksum (32 bits / 1 word)
 			m_WordIndex++;
@@ -174,30 +177,23 @@ namespace bitstream
 
 			BS_ASSERT(num_bytes * 8U >= m_NumBitsRead);
 
-			// Align with word size
-			uint32_t remainder = m_NumBitsRead % 32;
-			if (remainder != 0)
-			{
-				uint32_t zero;
-				bool status = serialize_bits(zero, 32 - remainder);
-
-				BS_ASSERT(status && zero == 0);
-			}
+			uint32_t offset = m_NumBitsRead / 32;
+			uint32_t zero;
 
 			// Test for zeros in padding
-			for (uint32_t i = m_WordIndex; i < num_bytes / 4; i++)
+			for (uint32_t i = offset; i < num_bytes / 4; i++)
 			{
-				uint32_t zero = 0;
 				bool status = serialize_bits(zero, 32);
 
 				BS_ASSERT(status && zero == 0);
 			}
 
+			uint32_t remainder = num_bytes * 8U - m_NumBitsRead;
+
 			// Test the last word more carefully, as it may have data
-			if (num_bytes % 4 != 0)
+			if (remainder % 32U != 0U)
 			{
-				uint32_t zero = 0;
-				bool status = serialize_bits(zero, (num_bytes % 4) * 8);
+				bool status = serialize_bits(zero, remainder);
 
 				BS_ASSERT(status && zero == 0);
 			}

include/bitstream/stream/bit_writer.h

@@ -159,7 +159,10 @@ namespace bitstream
 				m_WordIndex++;
 			}
 
-			return (m_NumBitsWritten - 1U) / 8U + 1U;
+			if (m_NumBitsWritten > 0U)
+				return (m_NumBitsWritten - 1U) / 8U + 1U;
+			else
+				return 0U;
 		}
 
 		/**
@@ -189,13 +192,13 @@ namespace bitstream
 			uint32_t num_bytes = flush();
 
 			// Copy protocol version to buffer
-			std::memcpy(m_Buffer, &protocol_version, sizeof(uint32_t));
+			*m_Buffer = protocol_version;
 
 			// Generate checksum of version + data
 			uint32_t checksum = utility::crc_uint32(reinterpret_cast<uint8_t*>(m_Buffer), num_bytes);
 
 			// Put checksum at beginning
-			std::memcpy(m_Buffer, &checksum, sizeof(uint32_t));
+			*m_Buffer = checksum;
 
 			return num_bytes;
 		}
@@ -209,18 +212,20 @@ namespace bitstream
 		{
 			BS_ASSERT(num_bytes * 8U <= m_TotalBits);
 
-			flush();
-
 			BS_ASSERT(num_bytes * 8U >= m_NumBitsWritten);
 
-			// Set to the padding to 0
-			std::memset(m_Buffer + m_WordIndex, 0, num_bytes - m_WordIndex * 4U);
+			uint32_t offset = m_NumBitsWritten / 32;
+			uint32_t zero = 0;
+
+			// Serialize words
+			for (uint32_t i = offset; i < num_bytes / 4; i++)
+				BS_ASSERT(serialize_bits(zero, 32));
 
-			m_NumBitsWritten = num_bytes * 8U;
+			uint32_t remainder = num_bytes * 8U - m_NumBitsWritten;
 
-			m_Scratch = 0ULL;
-			m_ScratchBits = (num_bytes % 4U) * 8U;
-			m_WordIndex = num_bytes / 4U;
+			// Align to byte
+			if (remainder % 32U != 0U)
+				BS_ASSERT(serialize_bits(zero, remainder));
 
 			return true;
 		}

src/test/stream_test.cpp

@@ -77,27 +77,56 @@ namespace bitstream::test::stream
 	{
 		// Test checksum
 		uint32_t protocol_version = 0xDEADBEEF;
-		uint32_t value = 3;
+		uint32_t value = 5;
 
 		// Write some initial values and finish with a checksum
 		byte_buffer<16> buffer;
 		bit_writer writer(buffer);
 
 		writer.prepend_checksum();
-		BS_TEST_ASSERT(writer.serialize_bits(value, 2));
+		BS_TEST_ASSERT(writer.serialize_bits(value, 3));
 		uint16_t num_bytes = writer.serialize_checksum(protocol_version);
 
 		// Read the checksum and validate
 		uint32_t out_value;
 		bit_reader reader(buffer, num_bytes);
 
 		BS_TEST_ASSERT(reader.serialize_checksum(protocol_version));
-		BS_TEST_ASSERT(reader.serialize_bits(out_value, 2));
+		BS_TEST_ASSERT(reader.serialize_bits(out_value, 3));
 
 		BS_TEST_ASSERT(out_value == value);
 	}
 
-	BS_ADD_TEST(test_serialize_padding)
+	BS_ADD_TEST(test_serialize_padding_small)
+	{
+		// Test padding
+		uint32_t in_value1 = 3;
+		uint32_t in_value2 = 6;
+
+		// 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_value1, 3));
+		BS_TEST_ASSERT(writer.pad_to_size(2));
+		BS_TEST_ASSERT(writer.serialize_bits(in_value2, 5));
+		uint32_t num_bytes = writer.flush();
+
+		BS_TEST_ASSERT(num_bytes == 3);
+
+		// Read the values and validate padding
+		uint32_t out_value1;
+		uint32_t out_value2;
+		bit_reader reader(buffer, num_bytes);
+
+		BS_TEST_ASSERT(reader.serialize_bits(out_value1, 3));
+		BS_TEST_ASSERT(reader.pad_to_size(2));
+		BS_TEST_ASSERT(reader.serialize_bits(out_value2, 5));
+
+		BS_TEST_ASSERT(out_value1 == in_value1);
+		BS_TEST_ASSERT(out_value2 == in_value2);
+	}
+
+	BS_ADD_TEST(test_serialize_padding_large)
 	{
 		// Test padding
 		uint32_t in_value1 = 3;