MySqlArguments: Replace manual bitmap code with NullBitMap helper type

launchbadge · Apr 26, 2024 · e53367a · e53367a
1 parent a88438b
commit e53367a
Show file tree

Hide file tree

Showing 2 changed files with 56 additions and 7 deletions.
diff --git a/sqlx-mysql/src/arguments.rs b/sqlx-mysql/src/arguments.rs
@@ -3,13 +3,14 @@ use crate::types::Type;
 use crate::{MySql, MySqlTypeInfo};
 pub(crate) use sqlx_core::arguments::*;
 use sqlx_core::error::BoxDynError;
+use std::ops::Deref;
 
 /// Implementation of [`Arguments`] for MySQL.
 #[derive(Debug, Default, Clone)]
 pub struct MySqlArguments {
     pub(crate) values: Vec<u8>,
     pub(crate) types: Vec<MySqlTypeInfo>,
-    pub(crate) null_bitmap: Vec<u8>,
+    pub(crate) null_bitmap: NullBitMap,
 }
 
 impl MySqlArguments {
@@ -18,14 +19,11 @@ impl MySqlArguments {
         T: Encode<'q, MySql> + Type<MySql>,
     {
         let ty = value.produces().unwrap_or_else(T::type_info);
-        let index = self.types.len();
 
         self.types.push(ty);
-        self.null_bitmap.resize((index / 8) + 1, 0);
 
-        if let IsNull::Yes = value.encode(&mut self.values)? {
-            self.null_bitmap[index / 8] |= (1 << (index % 8)) as u8;
-        }
+        let is_null = value.encode(&mut self.values)?;
+        self.null_bitmap.push(is_null);
 
         Ok(())
     }
@@ -50,3 +48,53 @@ impl<'q> Arguments<'q> for MySqlArguments {
         self.types.len()
     }
 }
+
+#[derive(Debug, Default, Clone)]
+pub(crate) struct NullBitMap {
+    bytes: Vec<u8>,
+    length: usize,
+}
+
+impl NullBitMap {
+    fn push(&mut self, is_null: IsNull) {
+        let byte_index = self.length / (u8::BITS as usize);
+        let bit_offset = self.length % (u8::BITS as usize);
+
+        if bit_offset == 0 {
+            self.bytes.push(0);
+        }
+
+        self.bytes[byte_index] |= u8::from(is_null.is_null()) << bit_offset;
+        self.length += 1;
+    }
+}
+
+impl Deref for NullBitMap {
+    type Target = [u8];
+
+    fn deref(&self) -> &Self::Target {
+        &self.bytes
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use super::*;
+
+    #[test]
+    fn null_bit_map_should_push_is_null() {
+        let mut bit_map = NullBitMap::default();
+
+        bit_map.push(IsNull::Yes);
+        bit_map.push(IsNull::No);
+        bit_map.push(IsNull::Yes);
+        bit_map.push(IsNull::No);
+        bit_map.push(IsNull::Yes);
+        bit_map.push(IsNull::No);
+        bit_map.push(IsNull::Yes);
+        bit_map.push(IsNull::No);
+        bit_map.push(IsNull::Yes);
+
+        assert_eq!([0b01010101, 0b1].as_slice(), bit_map.deref());
+    }
+}
diff --git a/sqlx-mysql/src/protocol/statement/execute.rs b/sqlx-mysql/src/protocol/statement/execute.rs
@@ -2,6 +2,7 @@ use crate::io::Encode;
 use crate::protocol::text::ColumnFlags;
 use crate::protocol::Capabilities;
 use crate::MySqlArguments;
+use std::ops::Deref;
 
 // https://dev.mysql.com/doc/dev/mysql-server/8.0.12/page_protocol_com_stmt_execute.html
 
@@ -19,7 +20,7 @@ impl<'q> Encode<'_, Capabilities> for Execute<'q> {
         buf.extend(&1_u32.to_le_bytes()); // iterations (always 1): int<4>
 
         if !self.arguments.types.is_empty() {
-            buf.extend(&*self.arguments.null_bitmap);
+            buf.extend(&*self.arguments.null_bitmap.deref());
             buf.push(1); // send type to server
 
             for ty in &self.arguments.types {