#![allow(dead_code)]
use alloc::vec::Vec;
use core::ops::{Deref, DerefMut};
#[cfg(feature = "__alloc-instrument")]
pub type ProfiledVec<T> = InstrumentedVec<T>;
#[cfg(not(feature = "__alloc-instrument"))]
pub type ProfiledVec<T> = Vec<T>;
pub trait ProfiledVecExt<T> {
fn with_capacity_profiled(capacity: usize, context: &'static str) -> Self;
fn new_profiled(context: &'static str) -> Self;
}
#[cfg(feature = "__alloc-instrument")]
impl<T> ProfiledVecExt<T> for ProfiledVec<T> {
fn with_capacity_profiled(capacity: usize, context: &'static str) -> Self {
InstrumentedVec::with_capacity(capacity, context)
}
fn new_profiled(context: &'static str) -> Self {
InstrumentedVec::new(context)
}
}
#[cfg(not(feature = "__alloc-instrument"))]
impl<T> ProfiledVecExt<T> for ProfiledVec<T> {
#[inline]
fn with_capacity_profiled(capacity: usize, _context: &'static str) -> Self {
Vec::with_capacity(capacity)
}
#[inline]
fn new_profiled(_context: &'static str) -> Self {
Vec::new()
}
}
#[derive(Debug, Clone)]
pub struct VecStats {
pub context: &'static str,
pub final_len: usize,
pub final_capacity: usize,
pub initial_capacity: usize,
pub realloc_count: u32,
pub peak_capacity: usize,
pub element_size: usize,
}
impl VecStats {
pub fn utilization_pct(&self) -> f32 {
if self.final_capacity == 0 {
100.0
} else {
(self.final_len as f32 / self.final_capacity as f32) * 100.0
}
}
pub fn wasted_bytes(&self) -> usize {
(self.final_capacity.saturating_sub(self.final_len)) * self.element_size
}
}
#[cfg(feature = "__alloc-instrument")]
static STATS_ENABLED: core::sync::atomic::AtomicBool = core::sync::atomic::AtomicBool::new(true);
#[cfg(feature = "__alloc-instrument")]
static MIN_WASTE_REPORT: core::sync::atomic::AtomicUsize =
core::sync::atomic::AtomicUsize::new(1024);
#[cfg(feature = "__alloc-instrument")]
pub fn set_stats_enabled(enabled: bool) {
STATS_ENABLED.store(enabled, core::sync::atomic::Ordering::Relaxed);
}
#[cfg(feature = "__alloc-instrument")]
pub fn set_min_waste_report(bytes: usize) {
MIN_WASTE_REPORT.store(bytes, core::sync::atomic::Ordering::Relaxed);
}
#[derive(Debug)]
pub struct InstrumentedVec<T> {
inner: Vec<T>,
context: &'static str,
initial_capacity: usize,
realloc_count: u32,
peak_capacity: usize,
last_capacity: usize,
}
impl<T> InstrumentedVec<T> {
pub fn new(context: &'static str) -> Self {
Self {
inner: Vec::new(),
context,
initial_capacity: 0,
realloc_count: 0,
peak_capacity: 0,
last_capacity: 0,
}
}
pub fn with_capacity(capacity: usize, context: &'static str) -> Self {
Self {
inner: Vec::with_capacity(capacity),
context,
initial_capacity: capacity,
realloc_count: 0,
peak_capacity: capacity,
last_capacity: capacity,
}
}
pub fn from_vec(vec: Vec<T>, context: &'static str) -> Self {
let cap = vec.capacity();
Self {
inner: vec,
context,
initial_capacity: cap,
realloc_count: 0,
peak_capacity: cap,
last_capacity: cap,
}
}
#[inline]
fn check_realloc(&mut self) {
let current = self.inner.capacity();
if current != self.last_capacity {
if self.last_capacity > 0 {
self.realloc_count += 1;
}
self.last_capacity = current;
if current > self.peak_capacity {
self.peak_capacity = current;
}
}
}
pub fn push(&mut self, value: T) {
self.inner.push(value);
self.check_realloc();
}
pub fn extend<I: IntoIterator<Item = T>>(&mut self, iter: I) {
self.inner.extend(iter);
self.check_realloc();
}
pub fn reserve(&mut self, additional: usize) {
self.inner.reserve(additional);
self.check_realloc();
}
pub fn resize(&mut self, new_len: usize, value: T)
where
T: Clone,
{
self.inner.resize(new_len, value);
self.check_realloc();
}
pub fn stats(&self) -> VecStats {
VecStats {
context: self.context,
final_len: self.inner.len(),
final_capacity: self.inner.capacity(),
initial_capacity: self.initial_capacity,
realloc_count: self.realloc_count,
peak_capacity: self.peak_capacity,
element_size: core::mem::size_of::<T>(),
}
}
pub fn into_inner(self) -> Vec<T> {
let mut this = core::mem::ManuallyDrop::new(self);
core::mem::take(&mut this.inner)
}
pub fn inner(&self) -> &Vec<T> {
&self.inner
}
}
impl<T> Deref for InstrumentedVec<T> {
type Target = Vec<T>;
#[inline]
fn deref(&self) -> &Self::Target {
&self.inner
}
}
impl<T> DerefMut for InstrumentedVec<T> {
#[inline]
fn deref_mut(&mut self) -> &mut Self::Target {
&mut self.inner
}
}
impl<T> From<InstrumentedVec<T>> for Vec<T> {
fn from(v: InstrumentedVec<T>) -> Vec<T> {
v.into_inner()
}
}
#[cfg(feature = "__alloc-instrument")]
impl<T> Drop for InstrumentedVec<T> {
fn drop(&mut self) {
if !STATS_ENABLED.load(core::sync::atomic::Ordering::Relaxed) {
return;
}
let stats = self.stats();
let wasted = stats.wasted_bytes();
let min_report = MIN_WASTE_REPORT.load(core::sync::atomic::Ordering::Relaxed);
if wasted >= min_report || stats.realloc_count > 0 {
let util = stats.utilization_pct();
eprintln!(
"[alloc] {}: len={} cap={} ({:.0}% util, {}B wasted) init={} reallocs={}",
stats.context,
stats.final_len,
stats.final_capacity,
util,
wasted,
stats.initial_capacity,
stats.realloc_count,
);
}
}
}
#[cfg(not(feature = "__alloc-instrument"))]
impl<T> Drop for InstrumentedVec<T> {
fn drop(&mut self) {
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_profiled_vec_basic() {
let mut v: ProfiledVec<u8> = ProfiledVec::with_capacity_profiled(100, "test");
for i in 0..50u8 {
v.push(i);
}
assert_eq!(v.len(), 50);
}
#[test]
fn test_into_inner() {
let mut v = InstrumentedVec::with_capacity(10, "test");
v.push(1u8);
v.push(2);
let inner: Vec<u8> = v.into_inner();
assert_eq!(inner, vec![1, 2]);
}
}