blart 0.5.0

use std::{
    fmt,
    io::{self, Write},
    ptr,
};

use crate::{
    allocator::Allocator,
    raw::{InnerNode, NodePtr, NodeType, OpaqueNodePtr},
    visitor::{Visitable, Visitor},
    AsBytes, TreeMap,
};

/// Settings which customize the output of the [`DotPrinter`] visitor.
#[derive(Debug, Default, Clone, PartialEq, Eq)]
#[non_exhaustive]
pub struct DotPrinterSettings {
    /// Add node address to output in graphs
    pub display_node_address: bool,
}

/// A visitor of the radix trie that will print the tree in "dot" notation.
///
/// See ['DOT Language | Graphviz'](https://graphviz.org/doc/info/lang.html) for
/// information about syntax and example of the language.
#[derive(Debug)]
pub struct DotPrinter<O: Write, K, V> {
    output: O,
    next_id: usize,
    settings: DotPrinterSettings,
    key_fmt: fn(&K, &mut fmt::Formatter<'_>) -> fmt::Result,
    value_fmt: fn(&V, &mut fmt::Formatter<'_>) -> fmt::Result,
}

impl<O: Write, K, V> DotPrinter<O, K, V> {
    /// Write the dot-format of the given tree to the given output.
    pub fn print<const PREFIX_LEN: usize, A: Allocator>(
        output: O,
        tree: &TreeMap<K, V, PREFIX_LEN, A>,
        settings: DotPrinterSettings,
    ) -> Option<io::Result<()>>
    where
        K: fmt::Display,
        V: fmt::Display,
    {
        Self::print_with_fmt(
            output,
            tree,
            settings,
            <K as fmt::Display>::fmt,
            <V as fmt::Display>::fmt,
        )
    }

    /// Write the dot-format of the given tree to the given output.
    pub fn print_with_fmt<const PREFIX_LEN: usize, A: Allocator>(
        output: O,
        tree: &TreeMap<K, V, PREFIX_LEN, A>,
        settings: DotPrinterSettings,
        key_fmt: fn(&K, &mut fmt::Formatter<'_>) -> fmt::Result,
        value_fmt: fn(&V, &mut fmt::Formatter<'_>) -> fmt::Result,
    ) -> Option<io::Result<()>> {
        tree.state.as_ref().map(|state| {
            // SAFETY: Since we get a reference to the `TreeMap`, we know the
            // node and all descendants will not be mutated
            unsafe { Self::print_tree(output, &state.root, settings, key_fmt, value_fmt) }
        })
    }

    /// Write the dot-format of the given tree to the given output.
    ///
    /// # Safety
    ///  - For the duration of this function, the given node and all its
    ///    children nodes must not get mutated.
    pub(crate) unsafe fn print_tree<const PREFIX_LEN: usize>(
        output: O,
        tree: &OpaqueNodePtr<K, V, PREFIX_LEN>,
        settings: DotPrinterSettings,
        key_fmt: fn(&K, &mut fmt::Formatter<'_>) -> fmt::Result,
        value_fmt: fn(&V, &mut fmt::Formatter<'_>) -> fmt::Result,
    ) -> io::Result<()> {
        let mut visitor = DotPrinter {
            output,
            next_id: 0,
            settings,
            key_fmt,
            value_fmt,
        };

        visitor.output_prelude()?;
        let _ = tree.visit_with(&mut visitor)?;
        visitor.output_epilogue()
    }

    fn output_prelude(&mut self) -> io::Result<()> {
        writeln!(self.output, "strict digraph G {{")?;
        writeln!(self.output, "node [shape=record]")
    }

    fn output_epilogue(&mut self) -> io::Result<()> {
        writeln!(self.output, "}}")
    }

    fn get_id(&mut self) -> usize {
        let new_id = self.next_id;
        self.next_id += 1;
        new_id
    }

    fn write_inner_node<N, const PREFIX_LEN: usize>(&mut self, inner_node: &N) -> io::Result<usize>
    where
        N: InnerNode<PREFIX_LEN, Key = K, Value = V>,
    {
        let header = inner_node.header();
        let node_id = self.get_id();
        write!(self.output, "n{node_id} ")?;
        write!(self.output, "[label=\"{{")?;
        // write header line
        if self.settings.display_node_address {
            write!(
                self.output,
                "{{<h0> {:p}}}  | {{{:?} | {:?} | {:?}}} | {{",
                ptr::from_ref(inner_node),
                N::TYPE,
                header.prefix_len(),
                header.read_prefix()
            )?;
        } else {
            write!(
                self.output,
                "{{<h0> {:?} | {:?} | {:?}}} | {{",
                N::TYPE,
                header.prefix_len(),
                header.read_prefix()
            )?;
        }

        // SAFETY: The `child_it` does not live beyond the following loop and will not
        // overlap with any mutating access or operation, which is guaranteed by the
        // `print_tree` caller requirements.
        let child_it = inner_node.iter();
        for (idx, (key_fragment, _)) in child_it.enumerate() {
            if idx == 0 {
                write!(self.output, "<c{idx}> {key_fragment}")?;
            } else {
                write!(self.output, "| <c{idx}> {key_fragment}")?;
            }
        }
        writeln!(self.output, "}}}}\"]")?;

        // SAFETY: The `child_it` does not live beyond the following loop and will not
        // overlap with any mutating access or operation, which is guaranteed by the
        // `print_tree` caller requirements.
        let child_it = inner_node.iter();
        for (key_frag_id, (_, child)) in child_it.enumerate() {
            let child_id = child.visit_with(self)?;

            writeln!(self.output, "n{node_id}:c{key_frag_id} -> n{child_id}:h0")?;
        }

        Ok(node_id)
    }
}

impl<K, V, O, const PREFIX_LEN: usize> Visitor<K, V, PREFIX_LEN> for DotPrinter<O, K, V>
where
    O: Write,
{
    type Output = io::Result<usize>;

    fn default_output(&self) -> Self::Output {
        unimplemented!("this visitor should never use the default output")
    }

    fn combine_output(&self, _: Self::Output, _: Self::Output) -> Self::Output {
        unimplemented!("this visitor should never combine outputs")
    }

    fn visit_inner_node<N>(&mut self, t: &N) -> Self::Output
    where
        N: InnerNode<PREFIX_LEN, Key = K, Value = V> + Visitable<K, V, PREFIX_LEN>,
    {
        self.write_inner_node(t)
    }

    fn visit_leaf(&mut self, t: &super::LeafNode<K, V, PREFIX_LEN>) -> Self::Output {
        let key = OverrideDisplay {
            value: t.key_ref(),
            fmt_fn: self.key_fmt,
        };
        let value = OverrideDisplay {
            value: t.value_ref(),
            fmt_fn: self.value_fmt,
        };

        let node_id = self.get_id();
        write!(self.output, "n{node_id} ")?;
        write!(self.output, "[label=\"{{")?;
        // write header line
        if self.settings.display_node_address {
            writeln!(
                self.output,
                "{{<h0> {:p}}} | {{{:?}}} | {{{}}} | {{{}}} | {{{:p} {:p}}} }}\"]",
                ptr::from_ref(t),
                NodeType::Leaf,
                key,
                value,
                t.previous.map_or(ptr::null_mut(), NodePtr::to_ptr),
                t.next.map_or(ptr::null_mut(), NodePtr::to_ptr),
            )?;
        } else {
            writeln!(
                self.output,
                "{{<h0> {:?}}} | {{{}}} | {{{}}}}}\"]",
                NodeType::Leaf,
                key,
                value
            )?;
        }

        Ok(node_id)
    }
}

#[derive(Debug)]
struct OverrideDisplay<'a, T> {
    value: &'a T,
    fmt_fn: fn(&T, &mut fmt::Formatter<'_>) -> fmt::Result,
}

impl<T> fmt::Display for OverrideDisplay<'_, T> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        (self.fmt_fn)(self.value, f)
    }
}

/// Print a given value using the [`Debug::fmt`] implementation.
///
/// This function can be used to override the formatting of key or value types
/// when printing with [`DotPrinter`].
///
/// # Errors
///
/// This function returns an error if the provided [`fmt::Formatter`] returns
/// `Err`.
///
/// [`Debug::fmt`]: std::fmt::Debug::fmt
pub fn debug_as_display_fmt(value: &impl fmt::Debug, f: &mut fmt::Formatter) -> fmt::Result {
    value.fmt(f)
}

/// Print the string "\[null\]".
///
/// This function can be used to override the formatting of key or value types
/// when printing with [`DotPrinter`].
///
/// # Errors
///
/// This function returns an error if the provided [`fmt::Formatter`] returns
/// `Err`.
pub fn null_display_fmt(_value: &impl fmt::Debug, f: &mut fmt::Formatter) -> fmt::Result {
    write!(f, "[null]")
}

/// Print the byte representation of the given value as returned by [`AsBytes`].
///
/// This function can be used to override the formatting of key or value types
/// when printing with [`DotPrinter`].
///
/// # Errors
///
/// This function returns an error if the provided [`fmt::Formatter`] returns
/// `Err`.
pub fn bytes_display_fmt(value: &impl AsBytes, f: &mut fmt::Formatter) -> fmt::Result {
    write!(f, "{:?}", value.as_bytes())
}

#[cfg(test)]
mod tests {
    use alloc::{string::String, vec::Vec};

    #[cfg(not(miri))]
    use expect_test::expect_file;

    use super::*;
    use crate::testing::swap;

    #[test]
    #[cfg(not(miri))]
    fn simple_tree_output_to_dot() {
        let tree: TreeMap<_, _> = crate::testing::generate_key_fixed_length([3, 3])
            .enumerate()
            .map(swap)
            .collect();

        let mut buffer = Vec::new();

        DotPrinter::print_with_fmt(
            &mut buffer,
            &tree,
            DotPrinterSettings {
                display_node_address: false,
            },
            debug_as_display_fmt,
            <usize as fmt::Display>::fmt,
        )
        .unwrap()
        .unwrap();

        let output = String::from_utf8(buffer).unwrap();
        expect_file!["./pretty_printer/simple_tree_output_to_dot.expect"].assert_eq(&output);
    }

    #[test]
    fn debug_display_fmt_works() {
        struct TestFmt;

        impl fmt::Debug for TestFmt {
            fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
                write!(f, "This is the Debug impl!")
            }
        }

        let output = format!(
            "{}",
            OverrideDisplay {
                value: &TestFmt,
                fmt_fn: debug_as_display_fmt,
            }
        );
        assert_eq!(output, "This is the Debug impl!");
    }

    #[test]
    fn null_display_fmt_works() {
        let output = format!(
            "{}",
            OverrideDisplay {
                value: &(),
                fmt_fn: null_display_fmt,
            }
        );
        assert_eq!(output, "[null]");
    }

    #[test]
    fn bytes_display_fmt_works() {
        #[derive(Copy, Clone)]
        struct MyBytes(&'static [u8]);

        impl AsBytes for MyBytes {
            fn as_bytes(&self) -> &[u8] {
                self.0
            }
        }

        let value = MyBytes(&[0x01, 0x02, 0x03, 0x04]);
        let output = format!(
            "{}",
            OverrideDisplay {
                value: &value,
                fmt_fn: bytes_display_fmt
            }
        );
        assert_eq!(output, "[1, 2, 3, 4]");
    }

    #[test]
    fn print_empty_tree_returns_none() {
        let tree: TreeMap<u8, u8> = TreeMap::new();
        let mut buffer = Vec::new();

        let result = DotPrinter::print_with_fmt(
            &mut buffer,
            &tree,
            DotPrinterSettings::default(),
            debug_as_display_fmt,
            debug_as_display_fmt,
        );

        assert!(result.is_none());
        assert!(buffer.is_empty());
    }
}