mem_dbg 0.4.3

/*
 * 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;

/// String type used in `MemDbgImpl` signatures (re-exported for the derive
/// macro, which must name it in generated code without relying on the
/// caller's prelude).
#[doc(hidden)]
#[cfg(not(feature = "std"))]
pub use alloc::string::String;
/// String type used in `MemDbgImpl` signatures (re-exported for the derive
/// macro, which must name it in generated code without relying on the
/// caller's prelude).
#[doc(hidden)]
#[cfg(feature = "std")]
pub use std::string::String;

/// Hash map for pointer deduplication in mem_size (re-exported for derive macro).
#[doc(hidden)]
pub use hashbrown::HashMap;
/// Hash set for pointer deduplication in mem_dbg.
#[doc(hidden)]
pub use hashbrown::HashSet;

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

mod impl_mem_dbg;
mod impl_mem_size;

/// Layout-equivalent mirror of the private internal node that `LinkedList<T>`
/// heap-allocates per element. Re-exported so integration tests can compute
/// expected sizes against the same definition. Not part of the stable
/// public API. See the type's documentation for details.
#[doc(hidden)]
pub use impl_mem_size::LinkedListNode;

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 {
    /// Logical AND with another [`Boolean`].
    type And<B: Boolean>: Boolean;
    /// The Boolean value of the type.
    const VALUE: bool;
    /// Per-field flatness check used by the derive macro to warn when a
    /// `#[mem_size(flat)]` type has a non-flat field: `()` for [`True`] (flat),
    /// and the `#[must_use]` [`NonFlatField`] for [`False`] (non-flat). Not part
    /// of the public API.
    #[doc(hidden)]
    type FlatCheck: Copy;
    /// The value of [`Boolean::FlatCheck`].
    #[doc(hidden)]
    const FLAT_CHECK: Self::FlatCheck;
}

/// Marker discarded, one per field, by `#[derive(MemSize)]` on a
/// `#[mem_size(flat)]` type. It is `#[must_use]`, so a non-flat field (whose
/// [`FlatType::Flat`] is [`False`], mapping to this type) raises an
/// `unused_must_use` warning at the field. Not part of the public API.
#[must_use = "this field is not flat, so the enclosing `#[mem_size(flat)]` type under-counts its size: use `#[mem_size(rec)]`, or implement `FlatType` by hand. This will become a hard error in a future release."]
#[doc(hidden)]
#[derive(Clone, Copy)]
pub struct NonFlatField;

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

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

/// 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>`; otherwise, you
/// will not be able to compute the size of arrays, vectors, slices, and
/// containers.
///
/// 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, i.e., types all of whose fields are
/// themselves `FlatType<Flat = True>` (typically [`Copy`] + `'static` types
/// with no heap allocation and no references). `#[mem_size(flat)]` on a type
/// with a non-flat field (e.g., a [`Vec`], [`String`], [`Box`], or a reference,
/// which are all `Flat = False` because their size is not constant) currently
/// produces a compile-time `unused_must_use` warning at the offending field
/// (for non-generic types) and will become a hard error in a future release. If
/// you need to assert flatness the macro cannot verify, implement [`FlatType`]
/// by hand.
///
/// 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)]`:
///
/// ```ignore
/// #[derive(mem_dbg::MemSize)]
/// struct MyStruct(usize);
/// ```
pub trait FlatType {
    /// Whether the type is flat ([`True`]) or not ([`False`]).
    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). Reference cycles are handled by
        /// this same once-per-address accounting, so following a cyclic value
        /// terminates instead of recursing forever.
        ///
        /// A limitation of the current approach is that each reference is just
        /// a thin pointer (i.e., a memory address). Thus, we cannot tell, for
        /// example, that an inner reference to a referenced structure is being
        /// counted twice, as the referenced structure has no associated size.
        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). Counted-reference
        /// cycles are handled by this same once-per-address accounting, so
        /// following a cyclic value terminates instead of recursing forever.
        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.
    ///
    /// Implementations must return at least `core::mem::size_of_val(self)`.
    /// Wrapper and derive implementations subtract inline field storage before
    /// adding recursive contributions, so returning a smaller value can
    /// underflow in debug builds or wrap in optimized builds.
    /// 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 {
    /// Recursively displays the fields of `self` below the current line.
    ///
    /// The default implementation is a no-op, used by types in which it does
    /// not make sense, or it is impossible, to recurse.
    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<'a>(std::io::StderrLock<'a>);
        impl core::fmt::Write for Wrapper<'_> {
            #[inline(always)]
            fn write_str(&mut self, s: &str) -> core::fmt::Result {
                use std::io::Write;
                self.0.write_all(s.as_bytes()).map_err(|_| core::fmt::Error)
            }
        }
        let stderr = std::io::stderr();
        let mut dbg_refs = HashSet::new();
        self._mem_dbg_depth_on(
            &mut Wrapper(stderr.lock()),
            total_size,
            max_depth,
            &mut String::new(),
            Some("⏺"),
            true,
            padded_size,
            flags,
            &mut dbg_refs,
        )
    }

    /// Displays `self` at the given depth without any reference-address
    /// annotation; forwards to [`_mem_dbg_depth_on_impl`](Self::_mem_dbg_depth_on_impl)
    /// with [`RefDisplay::None`].
    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 {
            // Saturate so a misbehaving manual `MemDbgImpl` cannot trigger
            // a debug-build underflow on the padding line.
            let padding = padded_size.saturating_sub(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(())
    }
}