use std::{
cell::{RefCell, UnsafeCell},
collections::HashMap,
fmt::Debug,
rc::Rc,
};
use bytesize::ByteSize;
use thousands::Separable;
#[derive(Clone, Hash, PartialEq, Eq)]
pub enum MemUsage {
Simple(usize),
Struct {
name: &'static str,
fields: Vec<(&'static str, MemUsage)>,
},
Collection {
name: &'static str,
overhead: usize,
items: Vec<MemUsage>,
},
Enum {
name: &'static str,
discriminant: usize,
min_size: usize,
variant: Box<(&'static str, MemUsage)>,
},
}
impl MemUsage {
pub fn effective_size(&self) -> usize {
match self {
MemUsage::Simple(size) => *size,
MemUsage::Struct { fields, .. } => fields
.iter()
.map(|(_, val)| val.effective_size())
.sum::<usize>(),
MemUsage::Collection {
overhead, items, ..
} => {
overhead
+ items
.iter()
.map(|item| item.effective_size())
.sum::<usize>()
}
MemUsage::Enum {
min_size, variant, ..
} => (*min_size).max(variant.1.effective_size().saturating_sub(*min_size)),
}
}
}
impl Debug for MemUsage {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::Simple(size) => write!(f, "{}", &ByteSize(*size as u64)),
Self::Struct { name, fields } => {
let mut struct_f = f.debug_struct(&format!(
"{}({})",
name,
ByteSize(self.effective_size() as u64)
));
let mut fields_sorted = fields.clone();
fields_sorted.sort_by_key(|(_, usage)| -(usage.effective_size() as isize));
for (name, usage) in fields_sorted.iter() {
struct_f.field(name, usage);
}
struct_f.finish()
}
Self::Collection {
name,
overhead,
items,
} => {
let mut exemplars = HashMap::new();
for item in items {
*exemplars.entry(item.clone()).or_insert(0) += 1;
}
let mut exemplars_with_count = exemplars
.into_iter()
.map(|(shape, count)| (count, count * shape.effective_size(), shape))
.collect::<Vec<_>>();
exemplars_with_count.sort_by_key(|(_, total_size, _)| -(*total_size as isize));
let mut d_struct = f.debug_struct(&format!(
"{}({})",
name,
ByteSize(self.effective_size() as u64)
));
d_struct.field("overhead", overhead);
if items.is_empty() {
d_struct.field("items", &"(empty)");
return d_struct.finish();
}
if items.len() > 1 {
d_struct.field("len", &items.len().separate_with_underscores());
}
if exemplars_with_count.len() == 1 {
d_struct.field(
&format!(
"{}x",
&exemplars_with_count[0].0.separate_with_underscores()
),
&exemplars_with_count[0].2,
);
} else {
d_struct.field(
"exemplars",
&exemplars_with_count
.into_iter()
.map(|(count, total_size, usage)| {
(
format!(
"{}x = {}",
&count.separate_with_underscores(),
ByteSize(total_size as u64)
),
usage,
)
})
.collect::<Vec<_>>(),
);
}
d_struct.finish()
}
Self::Enum {
name,
min_size,
variant,
..
} => f
.debug_struct(&format!(
"{}({})",
name,
ByteSize(self.effective_size() as u64)
))
.field("min_size", min_size)
.field(variant.0, &variant.1)
.finish(),
}
}
}
pub trait MemUser {
fn mem_use(&self) -> MemUsage;
}
impl<T: MemUser> MemUser for Vec<T> {
fn mem_use(&self) -> MemUsage {
MemUsage::Collection {
name: "Vec",
overhead: std::mem::size_of::<Vec<T>>(),
items: self.iter().map(|item| item.mem_use()).collect(),
}
}
}
impl MemUser for u8 {
fn mem_use(&self) -> MemUsage {
MemUsage::Simple(std::mem::size_of::<u8>())
}
}
impl MemUser for u32 {
fn mem_use(&self) -> MemUsage {
MemUsage::Simple(std::mem::size_of::<u32>())
}
}
impl MemUser for i32 {
fn mem_use(&self) -> MemUsage {
MemUsage::Simple(std::mem::size_of::<i32>())
}
}
impl MemUser for u64 {
fn mem_use(&self) -> MemUsage {
MemUsage::Simple(std::mem::size_of::<u64>())
}
}
impl MemUser for i64 {
fn mem_use(&self) -> MemUsage {
MemUsage::Simple(std::mem::size_of::<i64>())
}
}
impl MemUser for usize {
fn mem_use(&self) -> MemUsage {
MemUsage::Simple(std::mem::size_of::<usize>())
}
}
impl MemUser for String {
fn mem_use(&self) -> MemUsage {
MemUsage::Collection {
name: "String",
overhead: std::mem::size_of::<String>(),
items: vec![MemUsage::Simple(self.capacity())],
}
}
}
impl<T: MemUser> MemUser for UnsafeCell<T> {
fn mem_use(&self) -> MemUsage {
unsafe { (*self.get()).mem_use() }
}
}
impl<K: MemUser, V: MemUser> MemUser for im::HashMap<K, V> {
fn mem_use(&self) -> MemUsage {
MemUsage::Collection {
name: "im::HashMap",
overhead: std::mem::size_of::<im::HashMap<K, V>>(), items: self
.iter()
.map(|(k, v)| MemUsage::Struct {
name: "Entry",
fields: vec![("key", k.mem_use()), ("val", v.mem_use())],
})
.collect(),
}
}
}
impl<K: MemUser + Ord, V: MemUser> MemUser for im::OrdMap<K, V> {
fn mem_use(&self) -> MemUsage {
MemUsage::Collection {
name: "im::OrdMap",
overhead: std::mem::size_of::<im::OrdMap<K, V>>(), items: self
.iter()
.map(|(k, v)| MemUsage::Struct {
name: "Entry",
fields: vec![("key", k.mem_use()), ("val", v.mem_use())],
})
.collect(),
}
}
}
impl<K: MemUser, V: MemUser> MemUser for HashMap<K, V> {
fn mem_use(&self) -> MemUsage {
MemUsage::Collection {
name: "HashMap",
overhead: std::mem::size_of::<HashMap<K, V>>(), items: self
.iter()
.map(|(k, v)| MemUsage::Struct {
name: "Entry",
fields: vec![("key", k.mem_use()), ("val", v.mem_use())],
})
.collect(),
}
}
}
impl<T: MemUser> MemUser for im::HashSet<T> {
fn mem_use(&self) -> MemUsage {
MemUsage::Collection {
name: "im::HashSet",
overhead: std::mem::size_of::<im::HashSet<T>>(), items: self.iter().map(|item| item.mem_use()).collect(),
}
}
}
impl<T: MemUser + Ord> MemUser for im::OrdSet<T> {
fn mem_use(&self) -> MemUsage {
MemUsage::Collection {
name: "im::OrdSet",
overhead: std::mem::size_of::<im::OrdSet<T>>(), items: self.iter().map(|item| item.mem_use()).collect(),
}
}
}
impl<T: MemUser + Clone> MemUser for im::Vector<T> {
fn mem_use(&self) -> MemUsage {
MemUsage::Collection {
name: "im::Vector",
overhead: std::mem::size_of::<im::Vector<T>>(), items: self.iter().map(|item| item.mem_use()).collect(),
}
}
}
impl<T: MemUser> MemUser for Option<T> {
fn mem_use(&self) -> MemUsage {
MemUsage::Collection {
name: "Option",
overhead: 1,
items: match self {
None => vec![MemUsage::Simple(std::mem::size_of::<Option<T>>())],
Some(val) => {
vec![val.mem_use()]
}
},
}
}
}
impl<T: MemUser> MemUser for RefCell<T> {
fn mem_use(&self) -> MemUsage {
self.borrow().mem_use()
}
}
impl<T: MemUser> MemUser for Rc<T> {
fn mem_use(&self) -> MemUsage {
self.as_ref().mem_use()
}
}
impl MemUser for Box<str> {
fn mem_use(&self) -> MemUsage {
MemUsage::Collection {
name: "Box<str>",
overhead: std::mem::size_of::<Box<str>>(),
items: vec![MemUsage::Simple(self.len())],
}
}
}
impl MemUser for Box<Box<str>> {
fn mem_use(&self) -> MemUsage {
MemUsage::Collection {
name: "Box<Box<str>>",
overhead: std::mem::size_of::<Box<Box<str>>>(),
items: vec![self.as_ref().mem_use()],
}
}
}
impl MemUser for Box<[u8]> {
fn mem_use(&self) -> MemUsage {
MemUsage::Collection {
name: "Box<[u8]>",
overhead: std::mem::size_of::<Box<u8>>(),
items: vec![MemUsage::Simple(self.len())],
}
}
}
impl MemUser for Box<Box<[u8]>> {
fn mem_use(&self) -> MemUsage {
MemUsage::Collection {
name: "Box<Box<[u8]>>",
overhead: std::mem::size_of::<Box<Box<u8>>>(),
items: vec![self.as_ref().mem_use()],
}
}
}
impl<A: MemUser, B: MemUser> MemUser for (A, B) {
fn mem_use(&self) -> MemUsage {
MemUsage::Struct {
name: "Tuple",
fields: vec![("0", self.0.mem_use()), ("1", self.1.mem_use())],
}
}
}
impl MemUser for litl::Value {
fn mem_use(&self) -> MemUsage {
todo!()
}
}
impl MemUser for litl::Val {
fn mem_use(&self) -> MemUsage {
todo!()
}
}