gvn.rs

use std::collections::HashMap;

use hercules_ir::ir::*;

use crate::*;

/*
 * Top level function to run global value numbering. In the sea of nodes, GVN is
 * fairly simple compared to in a normal CFG. Needs access to constants for
 * identity function simplification.
 */
pub fn gvn(editor: &mut FunctionEditor, gvn_constants_and_clones: bool) {
    // Create worklist (starts as all nodes) and value number hashmap.
    let mut worklist: Vec<NodeID> = (0..editor.func().nodes.len()).map(NodeID::new).collect();
    let mut value_numbers: HashMap<Node, NodeID> = HashMap::new();

    while let Some(work) = worklist.pop() {
        // Deleted nodes (and technically the original start) should be skipped.
        if editor.func().nodes[work.idx()].is_start() {
            continue;
        }

        // First, simplify the work node by unwrapping identity functions.
        let value = crawl_identities(work, editor);

        // Next, check if there is a value number for this simplified value yet.
        if let Some(number) = value_numbers.get(&editor.func().nodes[value.idx()]) {
            // If the number is this worklist item, there's nothing to be done.
            // Also, don't GVN constants and clones if indicated to not do so.
            if *number == work
                || (!gvn_constants_and_clones
                    && (editor.func().nodes[work.idx()]
                        .try_constant()
                        .map(|cons_id| !editor.get_constant(cons_id).is_scalar())
                        .unwrap_or(false)
                        || editor.func().nodes[work.idx()]
                            .try_write()
                            .map(|(_, _, indices)| indices.is_empty())
                            .unwrap_or(false)))
            {
                continue;
            }

            // Record the users of `work` before making any edits.
            let work_users: Vec<NodeID> = editor.get_users(work).collect();

            // At this point, we know the number of the node IDed `work` is
            // `number`. We want to replace `work` with `number`, which means
            // 1. replacing all uses of `work` with `number`
            // 2. deleting `work`
            let success = editor.edit(|mut edit| {
                edit.sub_edit(work, *number);
                let edit = edit.replace_all_uses(work, *number)?;
                edit.delete_node(work)
            });

            // If the edit was performed, then the old users of `work` may now
            // be congruent to other nodes.
            if success {
                worklist.extend(work_users);
            }
        } else {
            // If there isn't a number yet for this value, assign the value the
            // ID of the node as its number.
            value_numbers.insert(editor.func().nodes[value.idx()].clone(), value);
            // The crawled identity `value` may be different than `work` - add
            // `work` back on to the worklist so that in a subsequent iteration,
            // we can simplify it.
            worklist.push(work);
        }
    }
}

/*
 * Helper function for unwrapping identity functions.
 */
fn crawl_identities(mut work: NodeID, editor: &FunctionEditor) -> NodeID {
    loop {
        let func = editor.func();
        // TODO: replace with API for saner pattern matching on IR. Also,
        // actually add the rest of the identity funcs.
        if let Node::Binary {
            left,
            right,
            op: BinaryOperator::Add,
        } = func.nodes[work.idx()]
        {
            if let Node::Constant { id } = func.nodes[left.idx()] {
                if editor.get_constant(id).is_zero() {
                    work = right;
                    continue;
                }
            }
        }

        if let Node::Binary {
            left,
            right,
            op: BinaryOperator::Add,
        } = func.nodes[work.idx()]
        {
            if let Node::Constant { id } = func.nodes[right.idx()] {
                if editor.get_constant(id).is_zero() {
                    work = left;
                    continue;
                }
            }
        }

        return work;
    }
}