mem_dbg/

lib.rs

/*
 * SPDX-FileCopyrightText: 2023 Tommaso Fontana
 * SPDX-FileCopyrightText: 2023 Inria
 * SPDX-FileCopyrightText: 2023 Sebastiano Vigna
 *
 * SPDX-License-Identifier: Apache-2.0 OR LGPL-2.1-or-later
 */
#![doc = include_str!("../README.md")]
#![deny(unconditional_recursion)]
#![cfg_attr(not(feature = "std"), no_std)]
#[cfg(not(feature = "std"))]
extern crate alloc;

#[cfg(not(feature = "std"))]
use alloc::string::String;

// HashMap for pointer deduplication in mem_size (re-exported for derive macro)
pub use hashbrown::HashMap;
// HashSet for pointer deduplication in mem_dbg
pub use hashbrown::HashSet;

#[cfg(feature = "derive")]
pub use mem_dbg_derive::{MemDbg, MemSize};

mod impl_mem_dbg;
mod impl_mem_size;

mod utils;
pub use utils::*;

/// Internal trait used within [`FlatType`] to implement [`MemSize`] depending
/// on whether a type is flat or not.
///
/// It has only two implementations, [`True`] and [`False`].
///
/// The `And` associated type computes the logical AND with another [`Boolean`]:
/// - `True::And<B> = B` (true AND x = x)
/// - `False::And<B> = False` (false AND x = false)
///
/// This is used to determine [`FlatType`] for composite types like tuples:
/// a tuple is `FlatType<Flat=True>` only if all its components are
/// `FlatType<Flat=True>`.
pub trait Boolean {
    type And<B: Boolean>: Boolean;
    const VALUE: bool;
}

/// One of the two possible implementations of [`Boolean`].
pub struct True {}
impl Boolean for True {
    type And<B: Boolean> = B;
    const VALUE: bool = true;
}

/// One of the two possible implementations of [`Boolean`].
pub struct False {}
impl Boolean for False {
    type And<B: Boolean> = False;
    const VALUE: bool = false;
}

/// How to display a reference address in [`MemDbgImpl::_mem_dbg_depth_on_impl`].
#[derive(Clone, Copy)]
pub enum RefDisplay {
    /// No address display.
    None,
    /// First encounter of a reference: display `@ 0x...` after type name.
    FirstEncounter(usize),
    /// Back-reference to already-seen pointer: display `→ 0x...`, use pointer size, skip recursion.
    BackReference(usize),
}

/// Marker trait for flat types.
///
/// "Flat" means that the type has no heap indirection to account for, so its
/// size can be computed using [`size_of`]. The scope of the trait is slightly
/// wider than that of [`Copy`] because, for example, atomic types are not
/// [`Copy`] but they are considered to be flat. In a non-flat type the
/// computation of size must be done by iterating recursively on the size of the
/// fields.
///
/// The trait comes in two flavors: `FlatType<Flat=True>` and
/// `FlatType<Flat=False>`. In the first case, [`MemSize::mem_size`] can be
/// computed on arrays, vectors, slices, and supported containers by multiplying
/// the length or capacity by the size of the element type; in the second case,
/// it is necessary to iterate on each element.
///
/// Since we cannot use negative trait bounds, every type that is used as a
/// parameter of an array, vector, slice, or container type, must implement
/// either `FlatType<Flat=True>` or `FlatType<Flat=False>`. If you do not
/// implement either of these traits, you will not be able to compute the size
/// of arrays, vectors, slices, and containers, but error messages will be very
/// unhelpful due to the contrived way we have to implement mutually exclusive
/// types [working around the bug that prevents the compiler from understanding
/// that implementations for the two flavors of `FlatType` are mutually
/// exclusive](https://github.com/rust-lang/rfcs/pull/1672#issuecomment-1405377983).
///
/// If you use the provided derive macros all this logic will be hidden from
/// you. You'll just have to add the attribute `#[mem_size(flat)]` to your
/// structures if they are flat types that do not contain non-`'static`
/// references (typically [`Copy`] + `'static` types, but this is not enforced).
///
/// When all fields of a struct or enum implement `FlatType<Flat=True>` but the
/// type itself is not annotated with `#[mem_size(flat)]`, a compile-time error
/// will suggest adding `#[mem_size(flat)]` (if the type is flat) or
/// `#[mem_size(rec)]` (to explicitly opt out of the optimization and silence
/// the error).
///
/// For example, the following will not compile because `usize` implements
/// `FlatType<Flat=True>`, but the struct is not annotated with
/// `#[mem_size(flat)]` or `#[mem_size(rec)]`:
///
/// ```compile_fail
/// #[derive(mem_dbg::MemSize)]
/// struct MyStruct(usize);
/// ```
///
/// Note that this approach forces us to compute the size of non-flat types that
/// contain references by iteration _even if you do not specify_
/// [`SizeFlags::FOLLOW_REFS`].
pub trait FlatType {
    type Flat: Boolean;
}

bitflags::bitflags! {
    /// Flags for [`MemSize`].
    #[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
    pub struct SizeFlags: u32 {
        /// Follow references.
        ///
        /// By default [`MemSize::mem_size`] does not follow references and
        /// computes only the size of the reference itself. Note that the size
        /// of every reference will be added once (i.e., if you have two
        /// identical references to the same memory region, the size of that
        /// region will be added only once).
        const FOLLOW_REFS = 1 << 0;
        /// Return capacity instead of size.
        ///
        /// Size does not include memory allocated but not used: for example, in
        /// the case of a vector [`MemSize::mem_size`] calls [`Vec::len`] rather
        /// than [`Vec::capacity`].
        ///
        /// However, when this flag is specified [`MemSize::mem_size`] will
        /// return the size of all memory allocated, even if it is not used: for
        /// example, in the case of a vector this option makes
        /// [`MemSize::mem_size`] call [`Vec::capacity`] rather than
        /// [`Vec::len`].
        const CAPACITY = 1 << 1;
        /// Follow counted references (i.e., [`Rc`](std::rc::Rc) and
        /// [`Arc`](std::sync::Arc)).
        ///
        /// By default [`MemSize::mem_size`] does not follow counted references
        /// and computes only the size of the reference itself. Note that the
        /// size of every counted reference will be added once (i.e., if you
        /// have two identical counted references to the same memory region,
        /// the size of that region will be added only once).
        const FOLLOW_RCS = 1 << 2;
    }
}

impl Default for SizeFlags {
    /// The default set of flags is the empty set.
    fn default() -> Self {
        Self::empty()
    }
}

/// A trait to compute recursively the overall size or capacity of a structure,
/// as opposed to the stack size returned by [`core::mem::size_of()`].
///
/// You can derive this trait with `#[derive(MemSize)]` if all the fields of
/// your type implement [`MemSize`].
///
/// When implementing this trait manually, you should implement
/// [`mem_size_rec`](MemSize::mem_size_rec).
pub trait MemSize {
    /// Returns the (recursively computed) overall memory size of the structure
    /// in bytes.
    fn mem_size(&self, flags: SizeFlags) -> usize {
        let mut refs = HashMap::new();
        let base = self.mem_size_rec(flags, &mut refs);
        base + refs.into_values().sum::<usize>()
    }

    /// Recursive implementation that tracks visited references for deduplication.
    ///
    /// The parameter `refs` is a map from pointer addresses (coming from
    /// references) to their computed size that is used to count the space
    /// occupied by references only once in case any of the flags
    /// [`SizeFlags::FOLLOW_REFS`] or [`SizeFlags::FOLLOW_RCS`] is set.
    ///
    /// In case of custom (non-derive) implementations: types that do not
    /// contain references can simply ignore the `refs` parameter; otherwise,
    /// please have a look at the [implementation for
    /// references](#impl-MemSize-for-%26T).
    fn mem_size_rec(&self, flags: SizeFlags, refs: &mut HashMap<usize, usize>) -> usize;
}

bitflags::bitflags! {
    /// Flags for [`MemDbg`].
    #[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
    pub struct DbgFlags: u32 {
        /// Follow references. See [`SizeFlags::FOLLOW_REFS`].
        const FOLLOW_REFS = 1 << 0;
        /// Print memory usage in human readable format.
        const HUMANIZE = 1 << 1;
        /// Print memory usage as a percentage.
        const PERCENTAGE = 1 << 2;
        /// Print the type name.
        const TYPE_NAME = 1 << 3;
        /// Display capacity instead of size. See [`SizeFlags::CAPACITY`].
        const CAPACITY = 1 << 4;
        /// Add an underscore every 3 digits, when `HUMANIZE` is not set.
        const SEPARATOR = 1 << 5;
        /// Print fields in memory order (i.e., using the layout chosen by the
        /// compiler), rather than in declaration order.
        const RUST_LAYOUT = 1 << 6;
        /// Use colors to distinguish sizes.
        const COLOR = 1 << 7;
        /// Follow counted references. See [`SizeFlags::FOLLOW_RCS`].
        const FOLLOW_RCS = 1 << 8;
    }
}

impl DbgFlags {
    /// Translates flags that are in common with [`MemSize`] into [`SizeFlags`].
    pub const fn to_size_flags(&self) -> SizeFlags {
        let mut flags = SizeFlags::empty();
        if self.contains(DbgFlags::FOLLOW_REFS) {
            flags = flags.union(SizeFlags::FOLLOW_REFS);
        }
        if self.contains(DbgFlags::CAPACITY) {
            flags = flags.union(SizeFlags::CAPACITY);
        }
        if self.contains(DbgFlags::FOLLOW_RCS) {
            flags = flags.union(SizeFlags::FOLLOW_RCS);
        }
        flags
    }
}

impl Default for DbgFlags {
    /// The default set of flags contains [`DbgFlags::TYPE_NAME`],
    /// [`DbgFlags::SEPARATOR`], and [`DbgFlags::PERCENTAGE`].
    fn default() -> Self {
        Self::TYPE_NAME | Self::SEPARATOR | Self::PERCENTAGE
    }
}

/// A trait providing methods to display recursively the content and size of a
/// structure.
///
/// You can derive this trait with `#[derive(MemDbg)]` if all the fields of your
/// type implement [`MemDbg`]. Note that you will also need to derive
/// [`MemSize`].
pub trait MemDbg: MemDbgImpl {
    /// Writes to stderr debug info about the structure memory usage, expanding
    /// all levels of nested structures.
    #[cfg(feature = "std")]
    fn mem_dbg(&self, flags: DbgFlags) -> core::fmt::Result {
        self._mem_dbg_depth(
            <Self as MemSize>::mem_size(self, flags.to_size_flags()),
            usize::MAX,
            core::mem::size_of_val(self),
            flags,
        )
    }

    /// Writes to a [`core::fmt::Write`] debug info about the structure memory
    /// usage, expanding all levels of nested structures.
    fn mem_dbg_on(&self, writer: &mut impl core::fmt::Write, flags: DbgFlags) -> core::fmt::Result {
        let mut dbg_refs = HashSet::new();
        self._mem_dbg_depth_on(
            writer,
            <Self as MemSize>::mem_size(self, flags.to_size_flags()),
            usize::MAX,
            &mut String::new(),
            Some("⏺"),
            true,
            core::mem::size_of_val(self),
            flags,
            &mut dbg_refs,
        )
    }

    /// Writes to stderr debug info about the structure memory usage as
    /// [`mem_dbg`](MemDbg::mem_dbg), but expanding only up to `max_depth`
    /// levels of nested structures.
    #[cfg(feature = "std")]
    fn mem_dbg_depth(&self, max_depth: usize, flags: DbgFlags) -> core::fmt::Result {
        self._mem_dbg_depth(
            <Self as MemSize>::mem_size(self, flags.to_size_flags()),
            max_depth,
            core::mem::size_of_val(self),
            flags,
        )
    }

    /// Writes to a [`core::fmt::Write`] debug info about the structure memory
    /// usage as [`mem_dbg_on`](MemDbg::mem_dbg_on), but expanding only up to
    /// `max_depth` levels of nested structures.
    fn mem_dbg_depth_on(
        &self,
        writer: &mut impl core::fmt::Write,
        max_depth: usize,
        flags: DbgFlags,
    ) -> core::fmt::Result {
        let mut dbg_refs = HashSet::new();
        self._mem_dbg_depth_on(
            writer,
            <Self as MemSize>::mem_size(self, flags.to_size_flags()),
            max_depth,
            &mut String::new(),
            Some("⏺"),
            true,
            core::mem::size_of_val(self),
            flags,
            &mut dbg_refs,
        )
    }
}

/// Implements [`MemDbg`] for all types that implement [`MemDbgImpl`].
///
/// This is done so that no one can change the implementation of [`MemDbg`],
/// which ensures consistency in printing.
impl<T: MemDbgImpl> MemDbg for T {}

/// Inner trait used to implement [`MemDbg`].
///
/// This trait should not be implemented by users, which should use the
/// [`MemDbg`](mem_dbg_derive::MemDbg) derive macro instead.
///
/// The default no-op implementation is used by all types in which it does not
/// make sense, or it is impossible, to recurse.
#[allow(clippy::too_many_arguments)]
pub trait MemDbgImpl: MemSize {
    fn _mem_dbg_rec_on(
        &self,
        _writer: &mut impl core::fmt::Write,
        _total_size: usize,
        _max_depth: usize,
        _prefix: &mut String,
        _is_last: bool,
        _flags: DbgFlags,
        _dbg_refs: &mut HashSet<usize>,
    ) -> core::fmt::Result {
        Ok(())
    }

    #[cfg(feature = "std")]
    #[doc(hidden)]
    fn _mem_dbg_depth(
        &self,
        total_size: usize,
        max_depth: usize,
        padded_size: usize,
        flags: DbgFlags,
    ) -> core::fmt::Result {
        struct Wrapper(std::io::Stderr);
        impl core::fmt::Write for Wrapper {
            #[inline(always)]
            fn write_str(&mut self, s: &str) -> core::fmt::Result {
                use std::io::Write;
                self.0
                    .lock()
                    .write(s.as_bytes())
                    .map_err(|_| core::fmt::Error)
                    .map(|_| ())
            }
        }
        let mut dbg_refs = HashSet::new();
        self._mem_dbg_depth_on(
            &mut Wrapper(std::io::stderr()),
            total_size,
            max_depth,
            &mut String::new(),
            Some("⏺"),
            true,
            padded_size,
            flags,
            &mut dbg_refs,
        )
    }

    fn _mem_dbg_depth_on(
        &self,
        writer: &mut impl core::fmt::Write,
        total_size: usize,
        max_depth: usize,
        prefix: &mut String,
        field_name: Option<&str>,
        is_last: bool,
        padded_size: usize,
        flags: DbgFlags,
        dbg_refs: &mut HashSet<usize>,
    ) -> core::fmt::Result {
        self._mem_dbg_depth_on_impl(
            writer,
            total_size,
            max_depth,
            prefix,
            field_name,
            is_last,
            padded_size,
            flags,
            dbg_refs,
            RefDisplay::None,
        )
    }

    /// Internal implementation for depth display.
    ///
    /// The `ref_display` parameter controls how reference addresses are shown:
    /// - `RefDisplay::None`: no address display
    /// - `RefDisplay::FirstEncounter(ptr)`: show `@ 0x...` after type name
    /// - `RefDisplay::BackReference(ptr)`: show `→ 0x...`, use pointer size, skip recursion
    #[allow(clippy::too_many_arguments)]
    fn _mem_dbg_depth_on_impl(
        &self,
        writer: &mut impl core::fmt::Write,
        total_size: usize,
        max_depth: usize,
        prefix: &mut String,
        field_name: Option<&str>,
        is_last: bool,
        padded_size: usize,
        flags: DbgFlags,
        dbg_refs: &mut HashSet<usize>,
        ref_display: RefDisplay,
    ) -> core::fmt::Result {
        // Each depth level adds 2 characters to prefix ("│ " or "  ")
        // Use chars().count() since prefix contains multi-byte UTF-8 chars
        if prefix.chars().count() / 2 > max_depth {
            return Ok(());
        }

        // For back-references, use pointer size; otherwise compute full size
        let is_backref = matches!(ref_display, RefDisplay::BackReference(_));
        let display_size = if is_backref {
            core::mem::size_of_val(self)
        } else {
            <Self as MemSize>::mem_size(self, flags.to_size_flags())
        };

        if flags.contains(DbgFlags::COLOR) {
            let color = utils::color(display_size);
            writer.write_fmt(format_args!("{color}"))?;
        };

        if flags.contains(DbgFlags::HUMANIZE) {
            let (value, uom) = crate::utils::humanize_float(display_size);
            if uom == " B" {
                writer.write_fmt(format_args!("{:>5}  B ", display_size))?;
            } else {
                let precision = if value >= 100.0 {
                    1
                } else if value >= 10.0 {
                    2
                } else if value >= 1.0 {
                    3
                } else {
                    4
                };
                writer.write_fmt(format_args!("{0:>4.1$} {2} ", value, precision, uom))?;
            }
        } else if flags.contains(DbgFlags::SEPARATOR) {
            let mut align = crate::utils::n_of_digits(total_size);
            let mut size_for_sep = display_size;
            align += align / 3;
            let mut digits = crate::utils::n_of_digits(size_for_sep);
            let digit_align = digits + digits / 3;
            for _ in digit_align..align {
                writer.write_char(' ')?;
            }

            let first_digits = digits % 3;
            let mut multiplier = 10_usize.pow((digits - first_digits) as u32);
            if first_digits != 0 {
                writer.write_fmt(format_args!("{}", size_for_sep / multiplier))?;
            } else {
                multiplier /= 1000;
                digits -= 3;
                writer.write_fmt(format_args!(" {}", size_for_sep / multiplier))?;
            }

            while digits >= 3 {
                size_for_sep %= multiplier;
                multiplier /= 1000;
                writer.write_fmt(format_args!("_{:03}", size_for_sep / multiplier))?;
                digits -= 3;
            }

            writer.write_str(" B ")?;
        } else {
            let align = crate::utils::n_of_digits(total_size);
            writer.write_fmt(format_args!("{:>align$} B ", display_size))?;
        }

        if flags.contains(DbgFlags::PERCENTAGE) {
            writer.write_fmt(format_args!(
                "{:>6.2}% ",
                if total_size == 0 {
                    100.0
                } else {
                    100.0 * display_size as f64 / total_size as f64
                }
            ))?;
        }
        if flags.contains(DbgFlags::COLOR) {
            let reset_color = utils::reset_color();
            writer.write_fmt(format_args!("{reset_color}"))?;
        };
        if !prefix.is_empty() {
            // Find the byte index of the 3rd character
            let start_byte = prefix
                .char_indices()
                .nth(2) // Skip 2 characters to get to the 3rd
                .map(|(idx, _)| idx)
                .unwrap_or(prefix.len());
            writer.write_str(&prefix[start_byte..])?;
            if is_last {
                writer.write_char('╰')?;
            } else {
                writer.write_char('├')?;
            }
            writer.write_char('╴')?;
        }

        if let Some(field_name) = field_name {
            writer.write_fmt(format_args!("{}", field_name))?;
        }

        if flags.contains(DbgFlags::TYPE_NAME) {
            if flags.contains(DbgFlags::COLOR) {
                writer.write_fmt(format_args!("{}", utils::type_color()))?;
            }
            writer.write_fmt(format_args!(": {}", core::any::type_name::<Self>()))?;
            if flags.contains(DbgFlags::COLOR) {
                writer.write_fmt(format_args!("{}", utils::reset_color()))?;
            }
        }

        // Display reference address based on RefDisplay
        match ref_display {
            RefDisplay::FirstEncounter(ptr) => {
                if flags.contains(DbgFlags::COLOR) {
                    writer.write_fmt(format_args!("{}", utils::ref_color()))?;
                }
                writer.write_fmt(format_args!(
                    " @ 0x{:0width$x}",
                    ptr,
                    width = 2 * core::mem::size_of::<usize>()
                ))?;
                if flags.contains(DbgFlags::COLOR) {
                    writer.write_fmt(format_args!("{}", utils::reset_color()))?;
                }
            }
            RefDisplay::BackReference(ptr) => {
                if flags.contains(DbgFlags::COLOR) {
                    writer.write_fmt(format_args!("{}", utils::backref_color()))?;
                }
                writer.write_fmt(format_args!(
                    " → 0x{:0width$x}",
                    ptr,
                    width = 2 * core::mem::size_of::<usize>()
                ))?;
                if flags.contains(DbgFlags::COLOR) {
                    writer.write_fmt(format_args!("{}", utils::reset_color()))?;
                }
            }
            RefDisplay::None => {}
        }

        // Skip padding and recursion for back-references
        if !is_backref {
            let padding = padded_size - core::mem::size_of_val(self);

            if padding != 0 {
                writer.write_fmt(format_args!(" [{}B]", padding))?;
            }
        }

        writer.write_char('\n')?;

        // Skip recursion for back-references
        if !is_backref {
            if is_last {
                prefix.push_str("  ");
            } else {
                prefix.push_str("│ ");
            }

            self._mem_dbg_rec_on(
                writer, total_size, max_depth, prefix, is_last, flags, dbg_refs,
            )?;

            prefix.pop();
            prefix.pop();
        }

        Ok(())
    }
}