From 05dbf7de721d04e5c2a1b639fbdbd83d324ff15c Mon Sep 17 00:00:00 2001
From: Luca Cominardi <luca.cominardi@gmail.com>
Date: Mon, 18 Mar 2024 15:10:28 +0100
Subject: [PATCH] ZInt codec takes at most 9 bytes

---
 commons/zenoh-codec/src/core/zint.rs | 40 +++++++++++++---------------
 commons/zenoh-codec/tests/codec.rs   | 21 ++++++++++++++-
 2 files changed, 38 insertions(+), 23 deletions(-)

diff --git a/commons/zenoh-codec/src/core/zint.rs b/commons/zenoh-codec/src/core/zint.rs
index d5e1c158ec..0daff7348b 100644
--- a/commons/zenoh-codec/src/core/zint.rs
+++ b/commons/zenoh-codec/src/core/zint.rs
@@ -17,7 +17,7 @@ use zenoh_buffers::{
     writer::{DidntWrite, Writer},
 };
 
-const VLE_LEN: usize = 10;
+const VLE_LEN: usize = 9;
 
 impl LCodec<u64> for Zenoh080 {
     fn w_len(self, x: u64) -> usize {
@@ -29,7 +29,6 @@ impl LCodec<u64> for Zenoh080 {
         const B6: u64 = u64::MAX << (7 * 6);
         const B7: u64 = u64::MAX << (7 * 7);
         const B8: u64 = u64::MAX << (7 * 8);
-        const B9: u64 = u64::MAX << (7 * 9);
 
         if (x & B1) == 0 {
             1
@@ -47,10 +46,8 @@ impl LCodec<u64> for Zenoh080 {
             7
         } else if (x & B8) == 0 {
             8
-        } else if (x & B9) == 0 {
-            9
         } else {
-            10
+            9
         }
     }
 }
@@ -123,14 +120,20 @@ where
                 len += 1;
                 x >>= 7;
             }
-            // SAFETY: buffer is guaranteed to be VLE_LEN long where VLE_LEN is
-            //         the maximum number of bytes a VLE can take once encoded.
-            //         I.e.: x is shifted 7 bits to the right every iteration,
-            //         the loop is at most VLE_LEN iterations.
-            unsafe {
-                *buffer.get_unchecked_mut(len) = x as u8;
+            // In case len == VLE_LEN then all the bits have already been written in the latest iteration.
+            // Else we haven't written all the necessary bytes yet.
+            if len != VLE_LEN {
+                // SAFETY: buffer is guaranteed to be VLE_LEN long where VLE_LEN is
+                //         the maximum number of bytes a VLE can take once encoded.
+                //         I.e.: x is shifted 7 bits to the right every iteration,
+                //         the loop is at most VLE_LEN iterations.
+                unsafe {
+                    *buffer.get_unchecked_mut(len) = x as u8;
+                }
+                len += 1;
             }
-            len + 1
+            // The number of written bytes
+            len
         })?;
         Ok(())
     }
@@ -144,24 +147,17 @@ where
 
     fn read(self, reader: &mut R) -> Result<u64, Self::Error> {
         let mut b = reader.read_u8()?;
-        if (b & 0x80_u8) == 0 {
-            return Ok(b as u64);
-        }
 
         let mut v = 0;
         let mut i = 0;
         // 7 * VLE_LEN is beyond the maximum number of shift bits
-        while (b & 0x80_u8) != 0 && i != 7 * VLE_LEN {
+        while (b & 0x80_u8) != 0 && i != 7 * (VLE_LEN - 1) {
             v |= ((b & 0x7f_u8) as u64) << i;
             b = reader.read_u8()?;
             i += 7;
         }
-        if i != 7 * VLE_LEN {
-            v |= (b as u64) << i;
-            Ok(v)
-        } else {
-            Err(DidntRead)
-        }
+        v |= (b as u64) << i;
+        Ok(v)
     }
 }
 
diff --git a/commons/zenoh-codec/tests/codec.rs b/commons/zenoh-codec/tests/codec.rs
index 3fdb95e1b5..7b402a5645 100644
--- a/commons/zenoh-codec/tests/codec.rs
+++ b/commons/zenoh-codec/tests/codec.rs
@@ -121,10 +121,28 @@ macro_rules! run {
 // Core
 #[test]
 fn codec_zint() {
+    run!(u8, { u8::MIN });
+    run!(u8, { u8::MAX });
     run!(u8, { thread_rng().gen::<u8>() });
+
+    run!(u16, { u16::MIN });
+    run!(u16, { u16::MAX });
     run!(u16, { thread_rng().gen::<u16>() });
+
+    run!(u32, { u32::MIN });
+    run!(u32, { u32::MAX });
     run!(u32, { thread_rng().gen::<u32>() });
+
+    run!(u64, { u64::MIN });
+    run!(u64, { u64::MAX });
+    let codec = Zenoh080::new();
+    for i in 1..=codec.w_len(u64::MAX) {
+        run!(u64, { 1 << (7 * i) });
+    }
     run!(u64, { thread_rng().gen::<u64>() });
+
+    run!(usize, { usize::MIN });
+    run!(usize, { usize::MAX });
     run!(usize, thread_rng().gen::<usize>());
 }
 
@@ -138,11 +156,12 @@ fn codec_zint_len() {
     codec.write(&mut writer, n).unwrap();
     assert_eq!(codec.w_len(n), buff.len());
 
-    for i in 1..=9 {
+    for i in 1..=codec.w_len(u64::MAX) {
         let mut buff = vec![];
         let mut writer = buff.writer();
         let n: u64 = 1 << (7 * i);
         codec.write(&mut writer, n).unwrap();
+        println!("ZInt len: {} {:02x?}", n, buff);
         assert_eq!(codec.w_len(n), buff.len());
     }