cubecl-opt 0.11.0-pre.1

use alloc::vec::Vec;
use cubecl_ir::{Builtin, OpCode, Type, Value};
use hashbrown::HashMap;
use petgraph::graph::NodeIndex;
use smallvec::SmallVec;

use crate::{AtomicCounter, Function, GlobalState, PhiInstruction, passes::OptimizerPass};

use super::GlobalValues;

#[derive(Debug, Clone, Default)]
pub struct GvnPass;

impl OptimizerPass for GvnPass {
    fn apply_post_ssa(&mut self, func: &mut Function, state: &GlobalState, changes: AtomicCounter) {
        self.run(func, state, &changes);
    }
}

impl GvnPass {
    /// Run the GVN-PRE algorithm
    /// 1. Build forward and backward dominator trees
    /// 2. Run `build_sets` step to annotate the tree with value information
    /// 3. Insert expressions where they're needed to make partially redundant expressions fully
    ///    redundant
    /// 4. Replace fully redundant expressions with simple assignments from the leader of that
    ///    expression to `out`
    pub fn run(&mut self, func: &mut Function, state: &GlobalState, changes: &AtomicCounter) {
        let analysis = func.analysis::<GlobalValues>(state);

        analysis.0.borrow_mut().insert(func, state, changes);
        analysis.0.borrow_mut().eliminate(func, changes);
    }
}

/// A global value table that maps expressions and locals to the values they represent.
#[derive(Debug, Clone)]
pub struct ValueTable {
    pub(crate) value_numbers: HashMap<Value, u32>,
    pub(crate) expression_numbers: HashMap<Expression, u32>,

    pub(crate) next_expr_num: u32,
    pub(crate) next_value_num: u32,
}
impl ValueTable {
    pub(crate) fn insert_phi(&mut self, func: &Function, phi: &PhiInstruction, val: u32) {
        let expr = Expression::Phi(
            phi.entries
                .iter()
                .map(|it| (func.value_of_var(&it.value).unwrap(), it.block))
                .collect(),
        );
        let out = func.value_of_var(&phi.out).unwrap();
        self.expression_numbers.insert(expr, val);
        self.value_numbers.insert(out, val);
    }
}

impl Default for ValueTable {
    fn default() -> Self {
        Self {
            value_numbers: Default::default(),
            expression_numbers: Default::default(),
            next_expr_num: 0,
            next_value_num: 1,
        }
    }
}

#[derive(PartialEq, Eq, PartialOrd, Ord, Hash, Clone, Debug)]
pub enum Expression {
    Instruction(Instruction),
    Copy(u32, Type),
    Value(Value),
    Volatile(Value),
    Phi(Vec<(Value, NodeIndex)>),
    Builtin(Builtin, Type),
}

impl Expression {
    pub fn depends_on(&self) -> SmallVec<[u32; 4]> {
        match self {
            Expression::Instruction(instruction) => instruction.args.clone(),
            Expression::Copy(val, _) => SmallVec::from_slice(&[*val]),
            Expression::Phi(_)
            | Expression::Volatile(_)
            | Expression::Value(_)
            | Expression::Builtin(..) => SmallVec::new(),
        }
    }

    /// Whether the expression is a trivial copy (which does not need to be hoisted since it's free)
    pub fn is_simple(&self) -> bool {
        matches!(self, Expression::Copy(_, _))
    }

    pub fn item(&self) -> Type {
        match self {
            Expression::Instruction(instruction) => instruction.item,
            Expression::Copy(_, item) => *item,
            Expression::Value(value) => value.ty,
            Expression::Volatile(value) => value.ty,
            Expression::Phi(entries) => entries[0].0.ty,
            Expression::Builtin(_, ty) => *ty,
        }
    }
}

impl From<Instruction> for Expression {
    fn from(value: Instruction) -> Self {
        Expression::Instruction(value)
    }
}

#[derive(PartialEq, Eq, PartialOrd, Ord, Hash, Clone, Debug)]
pub struct Instruction {
    pub(crate) op: OpCode,
    pub(crate) commutative: bool,
    pub(crate) args: SmallVec<[u32; 4]>,
    pub(crate) item: Type,
}

impl Instruction {
    pub fn new(op: impl Into<OpCode>, args: &[u32], item: Type) -> Self {
        Self {
            op: op.into(),
            commutative: false,
            args: SmallVec::from_slice(args),
            item,
        }
    }

    pub fn commutative(op: impl Into<OpCode>, args: &[u32], item: Type) -> Self {
        Self {
            op: op.into(),
            commutative: true,
            args: SmallVec::from_slice(args),
            item,
        }
    }
}