Add a fast path for primitive slices to all arrays

Author: mmastrac

gel-db-protocol/src/arrays.rs

@@ -315,6 +315,10 @@ mod tests {
         let mut buf = &data[..];
         let result = Array::<u32, u32>::decode_for(&mut buf);
         assert!(result.is_err());
+
+        let mut buf = [].as_slice();
+        let result = Array::<u32, u32>::decode_for(&mut buf);
+        assert!(result.is_err());
     }
 
     #[test]
@@ -334,6 +338,19 @@ mod tests {
         assert_eq!(collected, vec![1, 2, 3]);
     }
 
+    #[test]
+    fn test_zt_array_u32() {
+        // Unlikely, but helps test our primitive fast path
+        let data = vec![0xff, 0xff, 0xff, 0xff, 0xfe, 0xfe, 0xfe, 0xfe, 0];
+
+        let mut buf = &data[..];
+        let array = ZTArray::<u32>::decode_for(&mut buf).unwrap();
+
+        assert_eq!(array.len(), 2);
+        assert!(!array.is_empty());
+        assert_eq!(buf.len(), 0);
+    }
+
     #[test]
     fn test_zt_array_string() {
         let data = vec![
@@ -359,6 +376,12 @@ mod tests {
         let mut buf = &data[..];
         let result = ZTArray::<u8>::decode_for(&mut buf);
         assert!(result.is_err());
+
+        // Test with empty arrays
+        let mut buf = [].as_slice();
+        assert!(ZTArray::<u8>::decode_for(&mut buf).is_err());
+        assert!(ZTArray::<u32>::decode_for(&mut buf).is_err());
+        assert!(ZTArray::<ZTString>::decode_for(&mut buf).is_err());
     }
 
     #[test]

gel-db-protocol/src/encoding.rs

@@ -152,6 +152,16 @@ where
     fn decode_for(buf: &mut &'a [u8]) -> Result<Self, ParseError> {
         let len = L::decode_usize(buf)?;
         let orig_buf = *buf;
+        // Primitive types can skip the decode_for call.
+        if T::META.is_primitive {
+            let constant_size = T::META.constant_size.unwrap();
+            let byte_len = constant_size.saturating_mul(len);
+            if buf.len() < byte_len {
+                return Err(ParseError::TooShort);
+            }
+            *buf = &buf[byte_len..];
+            return Ok(Array::new(&orig_buf[..byte_len], len as _));
+        }
         for _ in 0..len {
             T::decode_for(buf)?;
         }
@@ -189,6 +199,25 @@ where
     fn decode_for(buf: &mut &'a [u8]) -> Result<Self, ParseError> {
         let mut orig_buf = *buf;
         let mut len = 0;
+
+        // Primitive types can skip the decode_for call and hunt for the 0 byte.
+        if T::META.is_primitive {
+            let constant_size = T::META.constant_size.unwrap();
+            loop {
+                if buf.is_empty() {
+                    return Err(ParseError::TooShort);
+                }
+                if buf[0] == 0 {
+                    break;
+                }
+                *buf = &buf[constant_size..];
+                len += 1;
+            }
+            *buf = &buf[1..];
+            orig_buf = &orig_buf[0..orig_buf.len() - buf.len() - 1];
+            return Ok(ZTArray::new(&orig_buf, len));
+        }
+
         loop {
             if buf.is_empty() {
                 return Err(crate::prelude::ParseError::TooShort);
@@ -233,6 +262,17 @@ where
 {
     fn decode_for(buf: &mut &'a [u8]) -> Result<Self, ParseError> {
         let orig_buf = *buf;
+        // Primitive types can skip the decode_for call and compute the number of elements
+        // until the end of the buffer.
+        if T::META.is_primitive {
+            let constant_size = T::META.constant_size.unwrap();
+            let len = buf.len() / constant_size;
+            if buf.len() % constant_size != 0 {
+                return Err(ParseError::TooShort);
+            }
+            *buf = &[];
+            return Ok(RestArray::new(orig_buf, len as _));
+        }
         let mut len = 0;
         while !buf.is_empty() {
             T::decode_for(buf)?;
@@ -542,10 +582,12 @@ declare_type!(u8);
 declare_type!(u16);
 declare_type!(u32);
 declare_type!(u64);
+declare_type!(u128);
 declare_type!(i8);
 declare_type!(i16);
 declare_type!(i32);
 declare_type!(i64);
+declare_type!(i128);
 
 declare_type!(f32);
 declare_type!(f64);

gel-db-protocol/src/macros.rs

@@ -4,7 +4,7 @@ macro_rules! declare_type {
     // Primitive types (no lifetime, fixed size)
     ($ty:ident) =>
     {
-        $crate::declare_type!($crate::prelude::DataType, $ty, {
+        $crate::declare_type!($crate::prelude::DataType, $ty, flags=[primitive], {
             fn to_usize(value: usize) -> $ty {
                 value as $ty
             }

gel-db-protocol/src/structs.rs

@@ -62,6 +62,7 @@ pub struct StructFieldMeta {
     pub is_enum: bool,
     pub is_struct: bool,
     pub is_array: bool,
+    pub is_primitive: bool,
 }
 
 impl StructFieldMeta {
@@ -73,6 +74,7 @@ impl StructFieldMeta {
             is_enum: false,
             is_struct: false,
             is_array: false,
+            is_primitive: false,
         }
     }
 
@@ -103,6 +105,13 @@ impl StructFieldMeta {
             ..self
         }
     }
+
+    pub const fn set_primitive(self) -> Self {
+        Self {
+            is_primitive: true,
+            ..self
+        }
+    }
 }
 
 #[derive(Clone, Copy, Debug)]