// Copyright (c) Microsoft Corporation.
// Licensed under the MIT License.

use crate::rir::{Callable, Instruction, Literal, Operand, Prim, Program, Ty, Variable};

#[cfg(test)]
mod tests;

pub fn check_types(program: &Program) {
    for (_, block) in program.blocks.iter() {
        for instr in &block.0 {
            check_instr_types(program, instr);
        }
    }
}

fn check_instr_types(program: &Program, instr: &Instruction) {
    match instr {
        Instruction::Call(id, args, var, _) => {
            check_call_types(program.get_callable(*id), args, *var);
        }

        Instruction::Branch(var, _, _, _) => assert_eq!(var.ty, Ty::Prim(Prim::Boolean)),

        Instruction::Add(opr1, opr2, var)
        | Instruction::Sub(opr1, opr2, var)
        | Instruction::Mul(opr1, opr2, var)
        | Instruction::Sdiv(opr1, opr2, var)
        | Instruction::Srem(opr1, opr2, var)
        | Instruction::Shl(opr1, opr2, var)
        | Instruction::Ashr(opr1, opr2, var)
        | Instruction::Fadd(opr1, opr2, var)
        | Instruction::Fsub(opr1, opr2, var)
        | Instruction::Fmul(opr1, opr2, var)
        | Instruction::Fdiv(opr1, opr2, var)
        | Instruction::LogicalAnd(opr1, opr2, var)
        | Instruction::LogicalOr(opr1, opr2, var)
        | Instruction::BitwiseAnd(opr1, opr2, var)
        | Instruction::BitwiseOr(opr1, opr2, var)
        | Instruction::BitwiseXor(opr1, opr2, var) => {
            assert_eq!(opr1.get_type(), opr2.get_type());
            assert_eq!(opr1.get_type(), var.ty);
        }

        Instruction::Fcmp(_, opr1, opr2, var) | Instruction::Icmp(_, opr1, opr2, var) => {
            assert_eq!(opr1.get_type(), opr2.get_type());
            assert_eq!(Ty::Prim(Prim::Boolean), var.ty);
        }

        Instruction::Store(opr, var)
        | Instruction::LogicalNot(opr, var)
        | Instruction::BitwiseNot(opr, var) => {
            assert_eq!(opr.get_type(), var.ty);
        }

        Instruction::Phi(args, var) => {
            for (opr, _) in args {
                assert_eq!(opr.get_type(), var.ty);
            }
        }

        Instruction::Index(array, index, var) => {
            match array {
                Operand::Literal(Literal::Array(global_idx)) => {
                    let Some(array_literal) = program.array_literals.get(*global_idx) else {
                        panic!("array literal index {global_idx} is out of bounds");
                    };
                    assert_eq!(Ty::Prim(array_literal.ty), var.ty);
                }
                Operand::Variable(Variable {
                    ty: Ty::Array(_, elem_ty),
                    ..
                }) => {
                    assert_eq!(Ty::Prim(*elem_ty), var.ty);
                }
                _ => panic!(
                    "expected array operand to be either a global array literal or an array variable"
                ),
            }
            assert_eq!(index.get_type(), Ty::Prim(Prim::Integer));
        }

        Instruction::Convert(_, _)
        | Instruction::Jump(_)
        | Instruction::Alloca(..)
        | Instruction::Load(..)
        | Instruction::Return => {}
    }
}

fn check_call_types(callable: &Callable, args: &[Operand], var: Option<Variable>) {
    assert_eq!(
        callable.input_type.len(),
        args.len(),
        "incorrect number of arguments"
    );
    for (arg, ty) in args.iter().zip(callable.input_type.iter()) {
        assert_eq!(arg.get_type(), *ty);
    }

    match (var, callable.output_type) {
        (Some(var), Some(ty)) => assert_eq!(ty, var.ty),
        (None, None) => {}
        _ => panic!("expected return type to be present in both the instruction and the callable"),
    }
}
