simplify_cfg.rs

use std::collections::{HashMap, HashSet};

use hercules_ir::*;

use crate::*;

/*
 * Top level function to simplify control flow in a Hercules function.
 */
pub fn simplify_cfg(
    editor: &mut FunctionEditor,
    fork_join_map: &HashMap<NodeID, NodeID>,
    reduce_cycles: &HashMap<NodeID, HashSet<NodeID>>,
) {
    // Collapse region chains.
    collapse_region_chains(editor);

    // Get rid of unnecessary fork-joins.
    remove_useless_fork_joins(editor, fork_join_map, reduce_cycles);
}

/*
 * Function to collapse region chains. A chain is a list of at least one region
 * node that takes only one control input. Region chains can be deleted. The use
 * of the head of the chain can turn into the use by the user of the tail of the
 * chain.
 */
fn collapse_region_chains(editor: &mut FunctionEditor) {
    // Loop over all region nodes.
    for id in editor.node_ids() {
        if let Node::Region { preds } = &editor.func().nodes[id.idx()] {
            let has_call_user = editor
                .get_users(id)
                .any(|x| editor.func().nodes[x.idx()].is_call());

            if preds.len() == 1 && !has_call_user {
                // Step 1: bridge gap between use and user.
                let predecessor = preds[0];
                let successor = editor
                    .get_users(id)
                    .filter(|x| !editor.func().nodes[x.idx()].is_phi())
                    .next()
                    .expect("Region node doesn't have a non-phi user.");

                editor.edit(|edit| {
                    // Set successor's use of this region to use the region's use.
                    edit.replace_all_uses_where(id, predecessor, |n| *n == successor)
                });

                // Step 2: bridge gap between uses and users of corresponding
                // phi nodes.
                let phis: Vec<NodeID> = editor
                    .get_users(id)
                    .filter(|x| editor.func().nodes[x.idx()].is_phi())
                    .collect();
                for phi_id in phis {
                    let data_uses =
                        if let Node::Phi { control, data } = &editor.func().nodes[phi_id.idx()] {
                            assert!(*control == id);
                            data
                        } else {
                            panic!()
                        };
                    assert!(data_uses.len() == 1, "Phi node doesn't have exactly one data use, while corresponding region had exactly one control use.");
                    let predecessor = data_uses[0];

                    editor.edit(|mut edit| {
                        // Set successors' use of this phi to use the phi's use.
                        edit = edit.replace_all_uses(phi_id, predecessor)?;
                        // Delete this phi.
                        edit.delete_node(phi_id)
                    });
                }

                // Delete this region.
                editor.edit(|edit| edit.delete_node(id));
            }
        }
    }
}

/*
 * Function to remove unused fork-joins. A fork-join is unused if there are no
 * reduce users of the join node. In such situations, it is asserted there are
 * no thread ID users of the fork as well. Also look for reduces that weren't
 * eliminated by DCE because they have a user, but the only users are in the
 * corresponding reduce cycle, so the reduce has no user outside the fork-join.
 */
fn remove_useless_fork_joins(
    editor: &mut FunctionEditor,
    fork_join_map: &HashMap<NodeID, NodeID>,
    reduce_cycles: &HashMap<NodeID, HashSet<NodeID>>,
) {
    // First, try to get rid of reduces where possible. Look for reduces with no
    // users outside its reduce cycle, and its reduce cycle contains no other
    // reduce nodes.
    for (_, join) in fork_join_map {
        let reduces: Vec<_> = editor
            .get_users(*join)
            .filter(|id| editor.func().nodes[id.idx()].is_reduce())
            .collect();

        for reduce in reduces {
            // If the reduce has users only in the reduce cycle, and none of
            // the nodes in the cycle are reduce nodes, then the reduce and its
            // whole cycle can be deleted.
            if editor
                .get_users(reduce)
                .all(|user| reduce_cycles[&reduce].contains(&user))
                && reduce_cycles[&reduce]
                    .iter()
                    .all(|id| !editor.func().nodes[id.idx()].is_reduce())
            {
                editor.edit(|mut edit| {
                    for id in reduce_cycles[&reduce].iter() {
                        edit = edit.delete_node(*id)?;
                    }
                    edit.delete_node(reduce)
                });
            }
        }
    }

    // Second, run DCE to get rid of thread IDs.
    dce(editor);

    // Third, get rid of fork-joins.
    for (fork, join) in fork_join_map {
        if editor.get_users(*join).len() == 1 {
            assert_eq!(editor.get_users(*fork).len(), 1);

            let fork_use = get_uses(&editor.func().nodes[fork.idx()]).as_ref()[0];
            let join_use = get_uses(&editor.func().nodes[join.idx()]).as_ref()[0];

            editor.edit(|mut edit| {
                edit = edit.replace_all_uses(*join, join_use)?;
                edit = edit.replace_all_uses(*fork, fork_use)?;
                edit = edit.delete_node(*fork)?;
                edit.delete_node(*join)
            });
        }
    }
}