Skip to content

Instantly share code, notes, and snippets.

@josephg
Last active February 3, 2025 01:14
Show Gist options
  • Save josephg/e2dd6e7baf0a764a21bd724f8a2e411b to your computer and use it in GitHub Desktop.
Save josephg/e2dd6e7baf0a764a21bd724f8a2e411b to your computer and use it in GitHub Desktop.
[package]
name = "writer-bench"
version = "0.1.0"
edition = "2021"
[dependencies]
criterion = "0.5.1"
[profile.release]
lto = true
codegen-units = 1
panic = "abort"
use std::io::Write;
use criterion::Criterion;
fn main() {
let mut c = Criterion::default()
.configure_from_args();
generic_writer(&mut c);
append_slice(&mut c);
append_slice_raw(&mut c);
}
struct WriteCtx {
times: usize,
data: &'static [u8],
expected: Vec<u8>,
arr: Vec<u8>,
}
impl WriteCtx {
fn new() -> Self {
let times = 1000;
let data = b"hello world\n";
// Pre-compute the expected output: `data` repeated `times` times.
let mut expected = Vec::with_capacity(data.len() * times);
for _ in 0..times {
expected.extend_from_slice(data);
}
// Pre-allocate the output vector with enough capacity.
let arr = Vec::with_capacity(expected.len());
WriteCtx {
times,
data,
expected,
arr,
}
}
fn clear(&mut self) {
// This retains the allocation capacity.
self.arr.clear();
}
}
fn generic_writer(c: &mut Criterion) {
c.bench_function("generic_writer", |b| {
let mut ctx = WriteCtx::new();
b.iter(|| {
ctx.clear();
for _ in 0..ctx.times {
// We use write_all so that all bytes from data are written.
ctx.arr.write_all(ctx.data).unwrap();
}
});
assert_eq!(ctx.arr.as_slice(), ctx.expected.as_slice());
});
}
fn append_slice(c: &mut Criterion) {
c.bench_function("append_slice", |b| {
let mut ctx = WriteCtx::new();
b.iter(|| {
ctx.clear();
for _ in 0..ctx.times {
ctx.arr.extend_from_slice(ctx.data);
}
});
assert_eq!(ctx.arr.as_slice(), ctx.expected.as_slice());
});
}
/// This version reserves all needed capacity ahead of time and uses unsafe
/// code to append bytes without bounds checks.
fn append_slice_raw(c: &mut Criterion) {
c.bench_function("append_slice_raw", |b| {
let mut ctx = WriteCtx::new();
// Make sure we have enough memory in arr. This should be taken care of by WriteCtx::new()
// but it doesn't hurt to be explicit given we'll be memcpy-ing.
ctx.arr.reserve_exact(ctx.times * ctx.data.len());
b.iter(|| {
ctx.clear();
// Reserve the total required capacity.
let data_len = ctx.data.len();
// SAFETY: We reserved enough capacity. We set the new length and then copy bytes
// from ctx.data into the new, uninitialized region.
unsafe {
ctx.arr.set_len(data_len * ctx.times);
for pos in (0..data_len * ctx.times).step_by(data_len) {
std::ptr::copy_nonoverlapping(
ctx.data.as_ptr(),
ctx.arr.as_mut_ptr().add(pos),
data_len,
);
}
}
});
assert_eq!(ctx.arr.as_slice(), ctx.expected.as_slice());
});
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment