feat: add basic codec support for binary le

pilota/Cargo.toml

@@ -44,3 +44,7 @@ rand = "0.8"
 [[bench]]
 name = "faststr"
 harness = false
+
+[[bench]]
+name = "thrift_binary"
+harness = false

pilota/benches/thrift_binary.rs

@@ -0,0 +1,178 @@
+#![allow(clippy::redundant_clone)]
+
+use bytes::BytesMut;
+use criterion::{black_box, criterion_group, criterion_main};
+use pilota::thrift::{TInputProtocol, TOutputProtocol};
+use rand::{self, Rng};
+
+fn binary_bench(c: &mut criterion::Criterion) {
+    let size = std::env::var("SIZE")
+        .unwrap_or("10000".to_string())
+        .parse()
+        .unwrap();
+    let mut group = c.benchmark_group("Bench Thrift Binary");
+    let mut v: Vec<i64> = Vec::with_capacity(size);
+    for _ in 0..size {
+        v.push(rand::thread_rng().gen());
+    }
+    let mut buf = BytesMut::new();
+
+    let mut p = pilota::thrift::binary::TBinaryProtocol::new(&mut buf, true);
+    for i in &v {
+        p.write_i64(*i).unwrap();
+    }
+    drop(p);
+    assert_eq!(buf.len(), 8 * size);
+
+    let mut buf_le = BytesMut::new();
+
+    let mut p = pilota::thrift::binary_le::TBinaryProtocol::new(&mut buf_le, true);
+    for i in &v {
+        p.write_i64(*i).unwrap();
+    }
+    drop(p);
+    assert_eq!(buf_le.len(), 8 * size);
+
+    let b = buf_le.clone();
+    let mut v2: Vec<i64> = Vec::with_capacity(size);
+    let src = b.as_ptr();
+    let dst = v2.as_mut_ptr();
+    unsafe {
+        std::ptr::copy_nonoverlapping(src, dst as *mut u8, size * 8);
+        v2.set_len(size);
+    }
+    assert_eq!(v, v2);
+
+    group.bench_function("big endian decode vec i64", |b| {
+        b.iter(|| {
+            black_box({
+                let b = buf_le.clone();
+                black_box(read_be(b, size));
+            });
+        })
+    });
+    group.bench_function("big endian decode vec i64 unsafe", |b| {
+        b.iter(|| {
+            black_box({
+                let b = buf_le.clone();
+                black_box(read_be_unsafe(b, size));
+            });
+        })
+    });
+    group.bench_function("little endian decode vec i64", |b| {
+        b.iter(|| {
+            black_box({
+                let b = buf_le.clone();
+                black_box(read_le(b, size));
+            });
+        })
+    });
+    group.bench_function("little endian decode vec i64 unsafe", |b| {
+        b.iter(|| {
+            black_box({
+                let b = buf_le.clone();
+                black_box(read_le_unsafe(b, size));
+            });
+        })
+    });
+    group.bench_function("little endian decode vec i64 optimized", |b| {
+        b.iter(|| {
+            black_box({
+                let b = buf_le.clone();
+                black_box(read_le_optimized(b, size));
+            });
+        })
+    });
+
+    group.bench_function("alloc vec", |b| {
+        b.iter(|| {
+            let mut b = buf_le.clone();
+            let _p = pilota::thrift::binary_le::TBinaryProtocol::new(&mut b, true);
+            let _: Vec<i64> = black_box(Vec::with_capacity(size));
+        })
+    });
+
+    group.finish();
+}
+
+#[inline(never)]
+fn read_be(mut b: BytesMut, size: usize) -> Vec<i64> {
+    let mut p = pilota::thrift::binary::TBinaryProtocol::new(&mut b, true);
+    let mut v = Vec::with_capacity(size);
+    for _ in 0..size {
+        v.push(p.read_i64().unwrap());
+    }
+    v
+}
+
+#[inline(never)]
+fn read_be_unsafe(b: BytesMut, size: usize) -> Vec<i64> {
+    unsafe {
+        let buf: &[u8] = b.as_ref();
+        assert!(buf.len() >= size * 8);
+        let mut index = 0;
+
+        let mut v = Vec::with_capacity(size);
+        for i in 0..size {
+            *v.get_unchecked_mut(i) = i64::from_be_bytes(
+                buf.get_unchecked(index..index + 8)
+                    .try_into()
+                    .unwrap_unchecked(),
+            );
+            index += 8;
+        }
+        v.set_len(size);
+        v
+    }
+}
+
+#[inline(never)]
+fn read_le(mut b: BytesMut, size: usize) -> Vec<i64> {
+    let mut p = pilota::thrift::binary_le::TBinaryProtocol::new(&mut b, true);
+
+    let mut v = Vec::with_capacity(size);
+    for _ in 0..size {
+        v.push(p.read_i64().unwrap());
+    }
+    v
+}
+
+// cargo asm -p pilota --bench thrift_binary --native --full-name --keep-labels --simplify --rust
+#[inline(never)]
+fn read_le_unsafe(b: BytesMut, size: usize) -> Vec<i64> {
+    unsafe {
+        let buf: &[u8] = b.as_ref();
+        assert!(buf.len() >= size * 8);
+        let mut index = 0;
+
+        let mut v = Vec::with_capacity(size);
+        for i in 0..size {
+            *v.get_unchecked_mut(i) = i64::from_le_bytes(
+                buf.get_unchecked(index..index + 8)
+                    .try_into()
+                    .unwrap_unchecked(),
+            );
+            index += 8;
+        }
+        v.set_len(size);
+        v
+    }
+}
+
+#[inline(never)]
+fn read_le_optimized(mut b: BytesMut, size: usize) -> Vec<i64> {
+    let _p = pilota::thrift::binary_le::TBinaryProtocol::new(&mut b, true);
+    let mut v: Vec<i64> = Vec::with_capacity(size);
+    let _ = black_box({
+        let src = b.as_ptr();
+        let dst = v.as_mut_ptr();
+        unsafe {
+            std::ptr::copy_nonoverlapping(src, dst as *mut u8, size * 8);
+            v.set_len(size);
+        }
+    });
+    v
+}
+
+criterion_group!(benches, binary_bench);
+criterion_main!(benches);