extern crate hercules_ir;

use self::hercules_ir::ir::*;

use crate::*;

/*
 * Substitute all uses of a dynamic constant A with dynamic constant B in a
 * type. Return the substituted version of the type, once memozied.
 */
pub(crate) fn substitute_dynamic_constants_in_type(
    dc_a: DynamicConstantID,
    dc_b: DynamicConstantID,
    ty: TypeID,
    edit: &mut FunctionEdit,
) -> TypeID {
    // Look inside the type for references to dynamic constants.
    let ty_clone = edit.get_type(ty).clone();
    match ty_clone {
        Type::Product(ref fields) => {
            let new_fields = fields
                .into_iter()
                .map(|field_id| substitute_dynamic_constants_in_type(dc_a, dc_b, *field_id, edit))
                .collect();
            if new_fields != *fields {
                edit.add_type(Type::Product(new_fields))
            } else {
                ty
            }
        }
        Type::Summation(ref variants) => {
            let new_variants = variants
                .into_iter()
                .map(|variant_id| {
                    substitute_dynamic_constants_in_type(dc_a, dc_b, *variant_id, edit)
                })
                .collect();
            if new_variants != *variants {
                edit.add_type(Type::Summation(new_variants))
            } else {
                ty
            }
        }
        Type::Array(elem_ty, ref dims) => {
            let new_elem_ty = substitute_dynamic_constants_in_type(dc_a, dc_b, elem_ty, edit);
            let new_dims = dims
                .into_iter()
                .map(|dim_id| substitute_dynamic_constants(dc_a, dc_b, *dim_id, edit))
                .collect();
            if new_elem_ty != elem_ty || new_dims != *dims {
                edit.add_type(Type::Array(new_elem_ty, new_dims))
            } else {
                ty
            }
        }
        _ => ty,
    }
}

/*
 * Substitute all uses of a dynamic constant A with dynamic constant B in a
 * dynamic constant C. Return the substituted version of C, once memoized. Takes
 * a mutable edit instead of an editor since this may create new dynamic
 * constants, which can only be done inside an edit.
 */
pub(crate) fn substitute_dynamic_constants(
    dc_a: DynamicConstantID,
    dc_b: DynamicConstantID,
    dc_c: DynamicConstantID,
    edit: &mut FunctionEdit,
) -> DynamicConstantID {
    // If C is just A, then just replace all of C with B.
    if dc_a == dc_c {
        return dc_b;
    }

    // Since we substitute non-sense dynamic constant IDs earlier, we explicitly
    // check that the provided ID to replace inside of is valid. Otherwise,
    // ignore.
    if dc_c.idx() >= edit.num_dynamic_constants() {
        return dc_c;
    }

    // If C is not just A, look inside of it to possibly substitute a child DC.
    let dc_clone = edit.get_dynamic_constant(dc_c).clone();
    match dc_clone {
        DynamicConstant::Constant(_) | DynamicConstant::Parameter(_) => dc_c,
        // This is a certified Rust moment.
        DynamicConstant::Add(left, right) => {
            let new_left = substitute_dynamic_constants(dc_a, dc_b, left, edit);
            let new_right = substitute_dynamic_constants(dc_a, dc_b, right, edit);
            if new_left != left || new_right != right {
                edit.add_dynamic_constant(DynamicConstant::Add(new_left, new_right))
            } else {
                dc_c
            }
        }
        DynamicConstant::Sub(left, right) => {
            let new_left = substitute_dynamic_constants(dc_a, dc_b, left, edit);
            let new_right = substitute_dynamic_constants(dc_a, dc_b, right, edit);
            if new_left != left || new_right != right {
                edit.add_dynamic_constant(DynamicConstant::Sub(new_left, new_right))
            } else {
                dc_c
            }
        }
        DynamicConstant::Mul(left, right) => {
            let new_left = substitute_dynamic_constants(dc_a, dc_b, left, edit);
            let new_right = substitute_dynamic_constants(dc_a, dc_b, right, edit);
            if new_left != left || new_right != right {
                edit.add_dynamic_constant(DynamicConstant::Mul(new_left, new_right))
            } else {
                dc_c
            }
        }
        DynamicConstant::Div(left, right) => {
            let new_left = substitute_dynamic_constants(dc_a, dc_b, left, edit);
            let new_right = substitute_dynamic_constants(dc_a, dc_b, right, edit);
            if new_left != left || new_right != right {
                edit.add_dynamic_constant(DynamicConstant::Div(new_left, new_right))
            } else {
                dc_c
            }
        }
        DynamicConstant::Rem(left, right) => {
            let new_left = substitute_dynamic_constants(dc_a, dc_b, left, edit);
            let new_right = substitute_dynamic_constants(dc_a, dc_b, right, edit);
            if new_left != left || new_right != right {
                edit.add_dynamic_constant(DynamicConstant::Rem(new_left, new_right))
            } else {
                dc_c
            }
        }
    }
}

/*
 * Substitute all uses of a dynamic constant A with dynamic constant B in a
 * constant. Return the substituted version of the constant, once memozied.
 */
pub(crate) fn substitute_dynamic_constants_in_constant(
    dc_a: DynamicConstantID,
    dc_b: DynamicConstantID,
    cons: ConstantID,
    edit: &mut FunctionEdit,
) -> ConstantID {
    // Look inside the type for references to dynamic constants.
    let cons_clone = edit.get_constant(cons).clone();
    match cons_clone {
        Constant::Product(ty, fields) => {
            let new_ty = substitute_dynamic_constants_in_type(dc_a, dc_b, ty, edit);
            let new_fields = fields
                .iter()
                .map(|field_id| {
                    substitute_dynamic_constants_in_constant(dc_a, dc_b, *field_id, edit)
                })
                .collect();
            if new_ty != ty || new_fields != fields {
                edit.add_constant(Constant::Product(new_ty, new_fields))
            } else {
                cons
            }
        }
        Constant::Summation(ty, idx, variant) => {
            let new_ty = substitute_dynamic_constants_in_type(dc_a, dc_b, ty, edit);
            let new_variant = substitute_dynamic_constants_in_constant(dc_a, dc_b, variant, edit);
            if new_ty != ty || new_variant != variant {
                edit.add_constant(Constant::Summation(new_ty, idx, new_variant))
            } else {
                cons
            }
        }
        Constant::Array(ty) => {
            let new_ty = substitute_dynamic_constants_in_type(dc_a, dc_b, ty, edit);
            if new_ty != ty {
                edit.add_constant(Constant::Array(new_ty))
            } else {
                cons
            }
        }
        _ => cons,
    }
}

/*
 * Substitute all uses of a dynamic constant A with dynamic constant B in a
 * node.
 */
pub(crate) fn substitute_dynamic_constants_in_node(
    dc_a: DynamicConstantID,
    dc_b: DynamicConstantID,
    node: &mut Node,
    edit: &mut FunctionEdit,
) {
    match node {
        Node::Fork {
            control: _,
            factors,
        } => {
            for factor in factors.into_iter() {
                *factor = substitute_dynamic_constants(dc_a, dc_b, *factor, edit);
            }
        }
        Node::Constant { id } => {
            *id = substitute_dynamic_constants_in_constant(dc_a, dc_b, *id, edit);
        }
        Node::DynamicConstant { id } => {
            *id = substitute_dynamic_constants(dc_a, dc_b, *id, edit);
        }
        Node::Call {
            control: _,
            function: _,
            dynamic_constants,
            args: _,
        } => {
            for dc_arg in dynamic_constants.into_iter() {
                *dc_arg = substitute_dynamic_constants(dc_a, dc_b, *dc_arg, edit);
            }
        }
        _ => {}
    }
}