ling-mir

use super::Transform;
use crate::ir::*;
use crate::loop_utils;
use std::collections::HashMap;

pub struct ConstantPropagation;

impl Transform for ConstantPropagation {
    fn run(&self, func: &mut MirFunction) -> bool {
        let mut assign_counts: HashMap<Local, usize> = HashMap::new();
        let mut constant_assignments: HashMap<Local, Constant> = HashMap::new();

        let loops = loop_utils::find_loops(func);
        let mut loop_blocks = std::collections::HashSet::new();
        for lp in &loops {
            loop_blocks.insert(lp.header);
            for &bb in &lp.body {
                loop_blocks.insert(bb);
            }
        }

        for (bb_idx, bb) in func.basic_blocks.iter().enumerate() {
            let in_loop = loop_blocks.contains(&BasicBlockId(bb_idx));
            for stmt in &bb.statements {
                if let StatementKind::Assign(dest, rval) = &stmt.kind {
                    *assign_counts.entry(*dest).or_insert(0) += if in_loop { 2 } else { 1 };
                    if let Rvalue::Use(Operand::Constant(c)) = rval {
                        constant_assignments.insert(*dest, c.clone());
                    }
                } else if let StatementKind::SetAttr(Operand::Copy(l) | Operand::Move(l), _, _) =
                    &stmt.kind
                {
                    *assign_counts.entry(*l).or_insert(0) += if in_loop { 2 } else { 1 };
                } else if let StatementKind::SetIndex(Operand::Copy(l) | Operand::Move(l), _, _) =
                    &stmt.kind
                {
                    *assign_counts.entry(*l).or_insert(0) += if in_loop { 2 } else { 1 };
                } else if let StatementKind::VectorStore(
                    Operand::Copy(l) | Operand::Move(l),
                    _,
                    _,
                ) = &stmt.kind
                {
                    *assign_counts.entry(*l).or_insert(0) += if in_loop { 2 } else { 1 };
                }
            }
        }

        let mut safe_constants: HashMap<Local, Constant> = HashMap::new();
        for (local, count) in &assign_counts {
            if *count == 1 {
                if let Some(c) = constant_assignments.get(local) {
                    safe_constants.insert(*local, c.clone());
                }
            }
        }

        let mut changed = false;

        if !safe_constants.is_empty() {
            for bb in &mut func.basic_blocks {
                for stmt in &mut bb.statements {
                    if let StatementKind::Assign(_, rval) = &mut stmt.kind {
                        changed |= self.propagate_constants_in_rvalue(rval, &safe_constants);
                    } else if let StatementKind::SetIndex(obj, idx, val) = &mut stmt.kind {
                        changed |= self.propagate_constants_in_operand(obj, &safe_constants);
                        changed |= self.propagate_constants_in_operand(idx, &safe_constants);
                        changed |= self.propagate_constants_in_operand(val, &safe_constants);
                    } else if let StatementKind::SetAttr(obj, _, val) = &mut stmt.kind {
                        changed |= self.propagate_constants_in_operand(obj, &safe_constants);
                        changed |= self.propagate_constants_in_operand(val, &safe_constants);
                    } else if let StatementKind::VectorStore(obj, idx, val) = &mut stmt.kind {
                        changed |= self.propagate_constants_in_operand(obj, &safe_constants);
                        changed |= self.propagate_constants_in_operand(idx, &safe_constants);
                        changed |= self.propagate_constants_in_operand(val, &safe_constants);
                    }
                }
                if let Some(term) = &mut bb.terminator {
                    if let TerminatorKind::SwitchInt { discr, .. } = &mut term.kind {
                        changed |= self.propagate_constants_in_operand(discr, &safe_constants);
                    }
                }
            }
        }

        for bb in &mut func.basic_blocks {
            let mut local_consts: HashMap<Local, Constant> = HashMap::new();
            for stmt in &mut bb.statements {
                if let StatementKind::Assign(dest, rval) = &mut stmt.kind {
                    changed |= self.propagate_constants_in_rvalue(rval, &local_consts);

                    if let Rvalue::Use(Operand::Constant(c)) = rval {
                        local_consts.insert(*dest, c.clone());
                    } else {
                        local_consts.remove(dest);
                    }
                }
            }
        }

        changed
    }
}

impl ConstantPropagation {
    fn propagate_constants_in_rvalue(
        &self,
        rval: &mut Rvalue,
        map: &HashMap<Local, Constant>,
    ) -> bool {
        match rval {
            Rvalue::Use(op) | Rvalue::UnaryOp(_, op) | Rvalue::GetAttr(op, _) => {
                self.propagate_constants_in_operand(op, map)
            },
            Rvalue::BinaryOp(_, l, r) | Rvalue::GetIndex(l, r) => {
                let mut changed = self.propagate_constants_in_operand(l, map);
                changed |= self.propagate_constants_in_operand(r, map);
                changed
            },
            Rvalue::Call { func, args } => {
                let mut changed = self.propagate_constants_in_operand(func, map);
                for arg in args {
                    changed |= self.propagate_constants_in_operand(arg, map);
                }
                changed
            },
            Rvalue::Aggregate(_, ops) => {
                let mut changed = false;
                for op in ops {
                    changed |= self.propagate_constants_in_operand(op, map);
                }
                changed
            },
            Rvalue::VectorSplat(op, _) => self.propagate_constants_in_operand(op, map),
            Rvalue::VectorLoad(obj, idx, _) => {
                let mut changed = self.propagate_constants_in_operand(obj, map);
                changed |= self.propagate_constants_in_operand(idx, map);
                changed
            },
            Rvalue::VectorFMA(a, b, c) => {
                let mut changed = self.propagate_constants_in_operand(a, map);
                changed |= self.propagate_constants_in_operand(b, map);
                changed |= self.propagate_constants_in_operand(c, map);
                changed
            },
            _ => false,
        }
    }

    fn propagate_constants_in_operand(
        &self,
        op: &mut Operand,
        map: &HashMap<Local, Constant>,
    ) -> bool {
        if let Operand::Copy(l) | Operand::Move(l) = op {
            if let Some(c) = map.get(l) {
                *op = Operand::Constant(c.clone());
                return true;
            }
        }
        false
    }
}