utils.rs

use std::iter::zip;

use hercules_ir::def_use::*;
use 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
            }
        }
        DynamicConstant::Min(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::Min(new_left, new_right))
            } else {
                dc_c
            }
        }
        DynamicConstant::Max(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::Max(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);
            }
        }
        _ => {}
    }
}

/*
 * Top level function to make a function have only a single return.
 */
pub fn collapse_returns(editor: &mut FunctionEditor) -> Option<NodeID> {
    let returns: Vec<NodeID> = (0..editor.func().nodes.len())
        .filter(|idx| editor.func().nodes[*idx].is_return())
        .map(NodeID::new)
        .collect();
    assert!(!returns.is_empty());
    if returns.len() == 1 {
        return Some(returns[0]);
    }
    let preds_before_returns: Vec<NodeID> = returns
        .iter()
        .map(|ret_id| get_uses(&editor.func().nodes[ret_id.idx()]).as_ref()[0])
        .collect();
    let data_to_return: Vec<NodeID> = returns
        .iter()
        .map(|ret_id| get_uses(&editor.func().nodes[ret_id.idx()]).as_ref()[1])
        .collect();

    // All of the old returns get replaced in a single edit.
    let mut new_return = None;
    editor.edit(|mut edit| {
        let region = edit.add_node(Node::Region {
            preds: preds_before_returns.into_boxed_slice(),
        });
        let phi = edit.add_node(Node::Phi {
            control: region,
            data: data_to_return.into_boxed_slice(),
        });
        for ret in returns {
            edit = edit.delete_node(ret)?;
        }
        new_return = Some(edit.add_node(Node::Return {
            control: region,
            data: phi,
        }));
        Ok(edit)
    });
    new_return
}

pub(crate) fn contains_between_control_flow(func: &Function) -> bool {
    let num_control = func.nodes.iter().filter(|node| node.is_control()).count();
    assert!(num_control >= 2, "PANIC: A Hercules function must have at least two control nodes: a start node and at least one return node.");
    num_control > 2
}

/*
 * Top level function to ensure a Hercules function contains at least one
 * control node that isn't the start or return nodes.
 */
pub fn ensure_between_control_flow(editor: &mut FunctionEditor) -> Option<NodeID> {
    if !contains_between_control_flow(editor.func()) {
        let ret = editor
            .node_ids()
            .skip(1)
            .filter(|id| editor.func().nodes[id.idx()].is_control())
            .next()
            .unwrap();
        let Node::Return { control, data } = editor.func().nodes[ret.idx()] else {
            panic!("PANIC: A Hercules function with only two control nodes must have a return node be the other control node, other than the start node.")
        };
        assert_eq!(control, NodeID::new(0), "PANIC: The only other control node in a Hercules function, the return node, is not using the start node.");
        let mut region_id = None;
        editor.edit(|mut edit| {
            edit = edit.delete_node(ret)?;
            region_id = Some(edit.add_node(Node::Region {
                preds: Box::new([NodeID::new(0)]),
            }));
            edit.add_node(Node::Return {
                control: region_id.unwrap(),
                data,
            });
            Ok(edit)
        });
        region_id
    } else {
        Some(
            editor
                .get_users(NodeID::new(0))
                .filter(|id| editor.func().nodes[id.idx()].is_control())
                .next()
                .unwrap(),
        )
    }
}

/*
 * Helper function to tell if two lists of indices have the same structure.
 */
pub(crate) fn indices_structurally_equivalent(indices1: &[Index], indices2: &[Index]) -> bool {
    if indices1.len() == indices2.len() {
        let mut equiv = true;
        for pair in zip(indices1, indices2) {
            equiv = equiv
                && match pair {
                    (Index::Field(idx1), Index::Field(idx2)) => idx1 == idx2,
                    (Index::Variant(idx1), Index::Variant(idx2)) => idx1 == idx2,
                    (Index::Position(ref pos1), Index::Position(ref pos2)) => {
                        assert_eq!(pos1.len(), pos2.len());
                        true
                    }
                    _ => false,
                };
        }
        equiv
    } else {
        false
    }
}

/*
 * Helper function to determine if two lists of indices may overlap.
 */
pub(crate) fn indices_may_overlap(indices1: &[Index], indices2: &[Index]) -> bool {
    for pair in zip(indices1, indices2) {
        match pair {
            // Check that the field numbers are the same.
            (Index::Field(idx1), Index::Field(idx2)) => {
                if idx1 != idx2 {
                    return false;
                }
            }
            // Variant indices always may overlap, since it's the same
            // underlying memory. Position indices always may overlap, since the
            // indexing nodes may be the same at runtime.
            (Index::Variant(_), Index::Variant(_)) | (Index::Position(_), Index::Position(_)) => {}
            _ => panic!(),
        }
    }
    // `zip` will exit as soon as either iterator is done - two sets of indices
    // may overlap when one indexes a larger sub-value than the other.
    true
}