use crate::mir::{BasicBlockId, BinaryOp, ConstValue, MirFunction, MirInstruction, ValueId};
use std::collections::BTreeMap;

/// Apply `var += step` before continue so that header sees updated value.
/// Returns the new ValueId of the variable if updated, otherwise None.
pub fn apply_increment_before_continue(
    f: &mut MirFunction,
    cur_bb: BasicBlockId,
    vars: &mut BTreeMap<String, ValueId>,
    var_name: &str,
    step: i64,
) -> Option<ValueId> {
    let cur_val = match vars.get(var_name) {
        Some(v) => *v,
        None => return None,
    };
    // Emit const step
    let step_v = f.next_value_id();
    if let Some(bb) = f.get_block_mut(cur_bb) {
        bb.add_instruction(MirInstruction::Const {
            dst: step_v,
            value: ConstValue::Integer(step),
        });
    }
    // Emit add
    let new_v = f.next_value_id();
    if let Some(bb) = f.get_block_mut(cur_bb) {
        bb.add_instruction(MirInstruction::BinOp {
            dst: new_v,
            op: BinaryOp::Add,
            lhs: cur_val,
            rhs: step_v,
        });
    }
    vars.insert(var_name.to_string(), new_v);
    Some(new_v)
}
