delete_uncalled.rs

use bitvec::prelude::*;

use hercules_ir::callgraph::*;
use hercules_ir::ir::*;

use crate::*;

/**
 * First renumber functions by deleting all functions unreachable from any entry function.
 * Then, updates all functions to use the new function numbering.
 * Returns a vec where the element at index i is the new idx of the ith function (if it was not
 * deleted).
 */
pub fn delete_uncalled(
    editors: &mut Vec<FunctionEditor>,
    callgraph: &CallGraph,
) -> Vec<Option<usize>> {
    // Step 1. Identify which functions are not reachable from entry nodes using BFS.
    let mut reachable = bitvec![u8, Lsb0; 0; editors.len()];
    let mut worklist = vec![];
    for (idx, editor) in editors.iter().enumerate() {
        if editor.func().entry {
            worklist.push(idx);
            reachable.set(idx, true);
        }
    }

    while let Some(func_idx) = worklist.pop() {
        for callee in callgraph.get_callees(FunctionID::new(func_idx)) {
            if !reachable[callee.idx()] {
                reachable.set(callee.idx(), true);
                worklist.push(callee.idx());
            }
        }
    }

    // Step 2. Compute the new index of each function, which is obtained by
    // deleteting all unreachable non-entry functions before it
    let mut new_idx = vec![None; editors.len()];
    let mut deleted_count = 0;
    for idx in 0..editors.len() {
        if !reachable[idx] {
            deleted_count += 1;
        } else {
            new_idx[idx] = Some(idx - deleted_count);
        }
    }

    // Step 3. Update all function callsites to use new indices. We assume
    // that all nodes in all functions will be mutable, so panic if any
    // edit fails.
    for editor in editors.iter_mut() {
        // Don't make edits to dead functions as they may include calls to dead functions
        if new_idx[editor.func_id().idx()].is_none() {
            continue;
        }

        let callsites: Vec<_> = editor
            .node_ids()
            .filter(|id| editor.func().nodes[id.idx()].is_call())
            .collect();

        for callsite in callsites {
            let success = editor.edit(|mut edit| {
                let (control, function, dynamic_constants, args) =
                    edit.get_node(callsite).try_call().unwrap();
                let new_node = edit.add_node(Node::Call {
                    control,
                    function: FunctionID::new(new_idx[function.idx()].unwrap()),
                    dynamic_constants: dynamic_constants.clone(),
                    args: args.clone(),
                });
                edit.sub_edit(callsite, new_node);
                let edit = edit.replace_all_uses(callsite, new_node)?;
                edit.delete_node(callsite)
            });
            assert!(
                success,
                "Expected all delete_uncalled callsite edits to succeed!"
            );
        }
    }

    new_idx
}