use crate::core::{
    compile_passes::compilation_pass::CompilationPass,
    ir::IntermediateRepresentation,
};

#[derive(Default)]
pub struct DeadCodeEliminator {
    used: Vec<bool>,
}

impl DeadCodeEliminator {
    pub fn eliminate(old_ir: IntermediateRepresentation) -> IntermediateRepresentation {
        Self::make_and_run(old_ir)
    }
}

impl CompilationPass for DeadCodeEliminator {
    fn run(&mut self, old_ir: IntermediateRepresentation) -> IntermediateRepresentation {
        let len = old_ir.get_exprs().len();
        self.used = vec![false; len];
        let mut to_handle = old_ir.get_outputs().to_vec();
        while let Some(expr_id) = to_handle.pop() {
            self.used[expr_id] = true;
            let deps = old_ir.get_exprs()[expr_id].get_deps();
            for dep in deps {
                if self.used[dep] {
                    continue;
                }
                self.used[dep] = true;
                to_handle.push(dep);
            }
        }
        let mut old_id_to_new_id = Vec::new();
        let mut exprs = Vec::new();
        let mut bounds = Vec::new();
        let mut is_plaintext = Vec::new();
        for old_expr_id in 0..len {
            let expr = &old_ir.get_exprs()[old_expr_id];
            if !self.used[old_expr_id] {
                if let Some(optimized_out) = expr.get_is_input_already_optimized_out() {
                    if !optimized_out.get() {
                        if !expr.get_input_name().starts_with("_") {
                            println!("Warning: arcis compiler detected that the value of input `{}` is unused or could be optimized out.", expr.get_input_name());
                        }
                        optimized_out.set(true);
                    }
                } else {
                    old_id_to_new_id.push(usize::MAX);
                    continue;
                }
            }
            let new_id = exprs.len();
            old_id_to_new_id.push(new_id);
            let new_expr = expr.clone().apply(|x| old_id_to_new_id[x]);
            exprs.push(new_expr);
            bounds.push(old_ir.get_bounds()[old_expr_id]);
            is_plaintext.push(old_ir.get_is_plaintext()[old_expr_id]);
        }
        let (_, old_ir_outputs, _, _, mut tracking) = old_ir.destructure();
        let new_outputs = old_ir_outputs
            .iter()
            .map(|x| old_id_to_new_id[*x])
            .collect();
        tracking.combine(old_id_to_new_id, exprs.len());
        IntermediateRepresentation::new_with_tracking(
            exprs,
            new_outputs,
            bounds,
            is_plaintext,
            tracking,
        )
    }
}
