cubecl-opt 0.11.0-pre.1

use core::fmt::Display;

use alloc::{
    format,
    rc::Rc,
    string::{String, ToString},
    vec::Vec,
};
use petgraph::{
    dot::{Config, Dot},
    prelude::StableDiGraph,
    visit::EdgeRef,
};

use crate::{
    BasicBlock, ControlFlow, Function, Optimizer,
    analyses::{
        liveness::{
            Liveness,
            shared::{SharedLiveness, SmemAllocation},
        },
        uniformity::Uniformity,
    },
    gvn::{BlockSets, Expression, GlobalValues, Instruction, ValueTable},
};

const DEBUG_GVN: bool = option_env!("CUBECL_DEBUG_GVN").is_some();

impl Display for Optimizer {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        write!(f, "main: {{\n{}\n}}", self.main)?;
        for (id, extra_func) in self.global_state.extra_functions.iter() {
            write!(
                f,
                "\n\nfunc_{id}[{}]({}): {{\n{}\n}}",
                extra_func
                    .implicit_params
                    .iter()
                    .map(|it| it.to_string())
                    .collect::<Vec<_>>()
                    .join(", "),
                extra_func
                    .explicit_params
                    .iter()
                    .map(|it| it.to_string())
                    .collect::<Vec<_>>()
                    .join(", "),
                extra_func
            )?;
        }
        Ok(())
    }
}

/// Debug display for the program state.
impl Display for Function {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        let global_nums = self
            .analysis_cache
            .try_get::<GlobalValues>()
            .unwrap_or_default();
        let liveness = self
            .analysis_cache
            .try_get::<Liveness>()
            .unwrap_or_else(|| Rc::new(Liveness::empty(self)));
        let shared_liveness = self
            .analysis_cache
            .try_get::<SharedLiveness>()
            .unwrap_or_else(|| Rc::new(SharedLiveness::empty(self)));
        let uniformity = self
            .analysis_cache
            .try_get::<Uniformity>()
            .unwrap_or_default();

        if DEBUG_GVN {
            writeln!(f, "# Value Table:")?;
            writeln!(f, "{}", global_nums.borrow().values)?;
        }

        let smems = shared_liveness
            .allocations
            .values()
            .map(|it| format!("    {it}"));
        let smems = smems.collect::<Vec<_>>().join(",\n");
        writeln!(f, "Shared memories: [\n{smems}\n]\n")?;

        for node in self.node_indices() {
            let id = node.index();
            let bb = &self[node];
            let uniform = match uniformity.is_block_uniform(node) {
                true => "uniform ",
                false => "",
            };
            writeln!(f, "{uniform}bb{id} {{")?;
            if DEBUG_GVN {
                let block_sets = &global_nums
                    .borrow()
                    .block_sets
                    .get(&node)
                    .cloned()
                    .unwrap_or_default();
                writeln!(f, "{block_sets}")?;
            }

            if !bb.block_use.is_empty() {
                writeln!(f, "    Uses: {:?}", bb.block_use)?;
            }
            let live_vars = liveness.at_block(node).iter();
            let live_vars = live_vars.map(|it| format!("%{it}"));
            let live_vars = live_vars.collect::<Vec<_>>();
            writeln!(f, "    Live variables: [{}]\n", live_vars.join(", "))?;
            let live_shared = shared_liveness.at_block(node).iter();
            let live_shared = live_shared.map(|it| format!("shared({it})"));
            let live_shared = live_shared.collect::<Vec<_>>();
            writeln!(
                f,
                "    Live shared memories: [{}]\n",
                live_shared.join(", ")
            )?;

            for phi in bb.phi_nodes.borrow().iter() {
                write!(f, "    {} = phi ", phi.out)?;
                for entry in &phi.entries {
                    write!(f, "[bb{}: ", entry.block.index())?;
                    write!(f, "{}]", entry.value)?;
                }
                let is_uniform = match uniformity.is_val_uniform(phi.out) {
                    true => " @ uniform",
                    false => "",
                };
                writeln!(f, ";{is_uniform}")?;
            }
            if !bb.phi_nodes.borrow().is_empty() {
                writeln!(f)?;
            }

            for op in bb.ops.borrow_mut().values_mut() {
                let op_fmt = op.to_string();
                if op_fmt.is_empty() {
                    continue;
                }

                let is_uniform = match op.out.is_some_and(|out| uniformity.is_val_uniform(out)) {
                    true => " @ uniform",
                    false => "",
                };
                writeln!(f, "    {op_fmt};{is_uniform}")?;
            }
            match &*bb.control_flow.borrow() {
                ControlFlow::IfElse {
                    cond,
                    then,
                    or_else,
                    merge,
                } => {
                    writeln!(
                        f,
                        "    {cond} ? bb{} : bb{}; merge: {}",
                        then.index(),
                        or_else.index(),
                        merge
                            .as_ref()
                            .map(|it| format!("bb{}", it.index()))
                            .unwrap_or("None".to_string())
                    )?;
                }
                super::ControlFlow::Switch {
                    value,
                    default,
                    branches,
                    ..
                } => {
                    write!(f, "    switch({value}) ")?;
                    for (val, block) in branches {
                        write!(f, "[{val}: bb{}] ", block.index())?;
                    }
                    writeln!(f, "[default: bb{}];", default.index())?;
                }
                super::ControlFlow::Loop {
                    body,
                    continue_target,
                    merge,
                } => {
                    writeln!(
                        f,
                        "    loop(continue: bb{}, merge: bb{})",
                        continue_target.index(),
                        merge.index()
                    )?;
                    writeln!(f, "    branch bb{};", body.index())?
                }
                super::ControlFlow::LoopBreak {
                    break_cond,
                    body,
                    continue_target,
                    merge,
                } => {
                    writeln!(
                        f,
                        "    loop(cond: {}, body: bb{} continue: bb{}, break: bb{})",
                        break_cond,
                        body.index(),
                        continue_target.index(),
                        merge.index()
                    )?;
                }
                super::ControlFlow::Return { value } => writeln!(
                    f,
                    "    return{};",
                    value.map(|it| format!(" {it}")).unwrap_or_default()
                )?,
                super::ControlFlow::Unreachable => writeln!(f, "    unreachable;")?,
                super::ControlFlow::None => {
                    let edge = self.edges(node).next();
                    let target = edge.map(|it| it.target().index()).unwrap_or(255);
                    writeln!(f, "    branch bb{target};")?;
                }
            }
            f.write_str("}\n\n")?;
        }

        Ok(())
    }
}

impl Display for BlockSets {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        let mut exp_gen = self.exp_gen.iter().collect::<Vec<_>>();
        exp_gen.sort_by_key(|it| it.0);
        let exp_gen = exp_gen
            .into_iter()
            .map(|(val, expr)| format!("{val}: {expr}"))
            .collect::<Vec<_>>();
        let mut phi_gen = self.phi_gen.iter().collect::<Vec<_>>();
        phi_gen.sort_by_key(|it| it.0);
        let phi_gen = phi_gen
            .into_iter()
            .map(|(val, expr)| format!("{val}: {expr}"))
            .collect::<Vec<_>>();
        let tmp_gen = self
            .tmp_gen
            .iter()
            .map(|it| format!("{it}"))
            .collect::<Vec<_>>();
        let mut leaders = self.leaders.iter().collect::<Vec<_>>();
        leaders.sort_by_key(|it| it.0);
        let leaders = leaders
            .into_iter()
            .map(|(val, expr)| format!("{val}: {expr}"))
            .collect::<Vec<_>>();
        let mut antic_out = self.antic_out.iter().collect::<Vec<_>>();
        antic_out.sort_by_key(|it| it.0);
        let antic_out = antic_out
            .into_iter()
            .map(|(val, expr)| format!("{val}: {expr}"))
            .collect::<Vec<_>>();
        let mut antic_in = self.antic_in.iter().collect::<Vec<_>>();
        antic_in.sort_by_key(|it| it.0);
        let antic_in = antic_in
            .into_iter()
            .map(|(val, expr)| format!("{val}: {expr}"))
            .collect::<Vec<_>>();

        writeln!(f, "    exp_gen: [{}]", exp_gen.join(", "))?;
        writeln!(f, "    phi_gen: [{}]", phi_gen.join(", "))?;
        writeln!(f, "    tmp_gen: [{}]", tmp_gen.join(", "))?;
        writeln!(f, "    leaders: [{}]", leaders.join(", "))?;
        writeln!(f, "    antic_in: [{}]", antic_in.join(", "))?;
        writeln!(f, "    antic_out: [{}]", antic_out.join(", "))
    }
}

impl Display for ValueTable {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        let mut values = self.value_numbers.iter().collect::<Vec<_>>();
        values.sort_by_key(|it| it.1);
        writeln!(f, "values: [")?;
        for (val, num) in values {
            writeln!(f, "    {num}: {val},")?;
        }
        writeln!(f, "]")?;
        writeln!(f, "expressions: [")?;
        let mut expressions = self.expression_numbers.iter().collect::<Vec<_>>();
        expressions.sort_by_key(|it| it.1);
        for (expr, val) in expressions {
            writeln!(f, "    {val}: {expr},")?;
        }
        writeln!(f, "]")
    }
}

impl Display for Expression {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        match self {
            Expression::Instruction(instruction) => write!(f, "{instruction}"),
            Expression::Copy(val, _) => write!(f, "copy({val})"),
            Expression::Value(value) => write!(f, "{value}"),
            Expression::Volatile(value) => write!(f, "volatile({value})"),
            Expression::Phi(entries) => write!(
                f,
                "phi({})",
                entries
                    .iter()
                    .map(|(val, b)| format!("{val}: bb{}", b.index()))
                    .collect::<Vec<_>>()
                    .join(", ")
            ),
            Expression::Builtin(builtin, _) => write!(f, "builtin({builtin:?})"),
        }
    }
}

impl Display for Instruction {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        write!(f, "{:?}: [{:?}]", self.op, self.args)
    }
}

impl Display for BasicBlock {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        for phi in self.phi_nodes.borrow().iter() {
            write!(f, "    {} = phi ", phi.out)?;
            for entry in &phi.entries {
                write!(f, "[bb{}: ", entry.block.index())?;
                write!(f, "{}]", entry.value)?;
            }
            writeln!(f, ";\n")?;
        }
        if !self.phi_nodes.borrow().is_empty() {
            writeln!(f)?;
        }

        for op in self.ops.borrow_mut().values_mut() {
            let op_fmt = op.to_string();
            if op_fmt.is_empty() {
                continue;
            }

            writeln!(f, "    {op_fmt};")?;
        }
        match &*self.control_flow.borrow() {
            ControlFlow::IfElse {
                cond,
                then,
                or_else,
                merge,
            } => {
                writeln!(
                    f,
                    "    {cond} ? bb{} : bb{}; merge: {}",
                    then.index(),
                    or_else.index(),
                    merge
                        .as_ref()
                        .map(|it| format!("bb{}", it.index()))
                        .unwrap_or("None".to_string())
                )?;
            }
            super::ControlFlow::Switch {
                value,
                default,
                branches,
                ..
            } => {
                write!(f, "    switch({value}) ")?;
                for (val, block) in branches {
                    write!(f, "[{val}: bb{}] ", block.index())?;
                }
                writeln!(f, "[default: bb{}];", default.index())?;
            }
            super::ControlFlow::Loop {
                body,
                continue_target,
                merge,
            } => {
                writeln!(
                    f,
                    "    loop(continue: bb{}, merge: bb{})",
                    continue_target.index(),
                    merge.index()
                )?;
                writeln!(f, "    branch bb{};", body.index())?
            }
            super::ControlFlow::LoopBreak {
                break_cond,
                body,
                continue_target,
                merge,
            } => {
                writeln!(
                    f,
                    "    loop(cond: {}, body: bb{} continue: bb{}, break: bb{})",
                    break_cond,
                    body.index(),
                    continue_target.index(),
                    merge.index()
                )?;
            }
            super::ControlFlow::Return { value } => writeln!(
                f,
                "    return{};",
                value.map(|it| format!(" {it}")).unwrap_or_default()
            )?,
            super::ControlFlow::Unreachable => writeln!(f, "    unreachable;")?,
            super::ControlFlow::None => {
                writeln!(f, "    branch;")?;
            }
        }
        Ok(())
    }
}

impl Display for SmemAllocation {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        let crate::MemoryBlock {
            value_ty,
            alignment,
            ..
        } = self.smem;
        write!(
            f,
            "shared(id: {}, offset: {}, align: {alignment}, ty: {value_ty})",
            self.id, self.offset,
        )
    }
}

impl Function {
    pub fn dot_viz(&self) -> String {
        let uniformity = self.analysis_cache.try_get::<Uniformity>();

        let get_node_attributes = |_, (index, bb)| {
            let uniform = uniformity
                .as_ref()
                .map(|uniformity| uniformity.is_block_uniform(index))
                .unwrap_or(false);
            let title = match uniform {
                true => format!("uniform bb{}", index.index()),
                false => format!("bb{}", index.index()),
            };
            let bb = format!("{bb}");
            let lines = bb
                    .lines()
                    .map(|it| it.trim())
                    .map(escape_html)
                    .filter(|it| !it.is_empty())
                    .enumerate()
                    .map(|(i, it)| {
                        format!(r#"<TR><TD ALIGN="LEFT"><FONT COLOR="dimgray">{i} </FONT></TD><TD ALIGN="LEFT">{it}</TD></TR>"#)
                    })
                    .collect::<Vec<_>>();
            format!(
                r#"label = <
<TABLE ALIGN="LEFT" BORDER="0" CELLSPACING="0" CELLBORDER="1" CELLPADDING="0">
    <TR><TD BGCOLOR="lightgray" ALIGN="LEFT" CELLPADDING="3"><B>{title}</B></TD></TR>
    <TR><TD CELLPADDING="4">
    <TABLE ALIGN="LEFT" BORDER="0" CELLSPACING="0" CELLPADDING="0">
        {}
    </TABLE>
    </TD></TR>
</TABLE>>"#,
                lines.join("")
            )
        };

        //Dot::with_config(&self.program, &[Config::EdgeNoLabel])
        let content: Dot<'_, &StableDiGraph<BasicBlock, u32>> = Dot::with_attr_getters(
            &self.graph,
            &[
                Config::EdgeNoLabel,
                Config::NodeNoLabel,
                Config::GraphContentOnly,
            ],
            &|_, _| String::new(),
            &get_node_attributes,
        );
        format!(
            r#"
digraph {{
    node [ shape = box, fontname = "Consolas, 'Courier New', monospace", fontsize = "12", margin = 0 ]
{content}
}}
"#
        )
    }
}

fn escape_html(s: &str) -> String {
    s.replace("&", "&amp;")
        .replace("<", "&lt;")
        .replace(">", "&gt;")
}