From 665c90f0b326b125658267b2e50c4a6b43b3a42a Mon Sep 17 00:00:00 2001
From: Luca Cominardi <luca.cominardi@gmail.com>
Date: Mon, 18 Mar 2024 18:10:32 +0100
Subject: [PATCH] Optimize zint encode/decode (#838)

* Rebase on protocol_changes

* Fix rebase conflict
---
 commons/zenoh-codec/src/core/zint.rs | 52 ++++++++++++++++------------
 commons/zenoh-codec/tests/codec.rs   | 21 ++++++++++-
 2 files changed, 50 insertions(+), 23 deletions(-)

diff --git a/commons/zenoh-codec/src/core/zint.rs b/commons/zenoh-codec/src/core/zint.rs
index 1c2f5a28e4..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
         }
     }
 }
@@ -112,15 +109,31 @@ where
     fn write(self, writer: &mut W, mut x: u64) -> Self::Output {
         writer.with_slot(VLE_LEN, move |buffer| {
             let mut len = 0;
-            let mut b = x as u8;
-            while x > 0x7f {
-                buffer[len] = b | 0x80;
+            while (x & !0x7f_u64) != 0 {
+                // 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) | 0x80_u8;
+                }
                 len += 1;
                 x >>= 7;
-                b = x as u8;
             }
-            buffer[len] = b;
-            len + 1
+            // 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;
+            }
+            // The number of written bytes
+            len
         })?;
         Ok(())
     }
@@ -137,19 +150,14 @@ where
 
         let mut v = 0;
         let mut i = 0;
-        let mut k = VLE_LEN;
-        while b > 0x7f && k > 0 {
-            v |= ((b & 0x7f) as u64) << i;
-            i += 7;
+        // 7 * VLE_LEN is beyond the maximum number of shift bits
+        while (b & 0x80_u8) != 0 && i != 7 * (VLE_LEN - 1) {
+            v |= ((b & 0x7f_u8) as u64) << i;
             b = reader.read_u8()?;
-            k -= 1;
-        }
-        if k > 0 {
-            v |= ((b & 0x7f) as u64) << i;
-            Ok(v)
-        } else {
-            Err(DidntRead)
+            i += 7;
         }
+        v |= (b as u64) << i;
+        Ok(v)
     }
 }
 
diff --git a/commons/zenoh-codec/tests/codec.rs b/commons/zenoh-codec/tests/codec.rs
index 7f23214b49..3bca8b7489 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());
     }