Use a call-graph in the inliner (FuelLabs#3186)

This reduces the time spent in the inliner, in compiling the [`amm`
library ](https://github.com/sway-libs/amm) from ~120s to ~5s.

Code sizes before/after on
[sway-applications](https://github.com/FuelLabs/sway-applications):
-----------------------------------
benchmark | before | after
----------------|----------|--------
dao-voting | 20568 | 20496
escrow | 72740 | 72668
fundraiser | 31344 | 31440
multisig-wallet | 6592 | 7000
NFT | 23200 | 23200
oracle | 1924 | 1924

While this change doesn't touch the inlining heuristics, but only the
order in which inlining of functions is undertaken, we have a heuristic
that checks the size of the function being inlined. **That** depends on
the order in which we inline, and hence the final set of inlining
decisions aren't identical (though largely similar). Overall, I think
the compile time benefits outweigh any differences in the inlining
decisions.

sway-core/src/lib.rs

@@ -27,7 +27,7 @@ use std::path::{Path, PathBuf};
 use std::sync::Arc;
 use sway_ast::Dependency;
 use sway_error::handler::{ErrorEmitted, Handler};
-use sway_ir::{Context, Function, Instruction, Kind, Module, Value};
+use sway_ir::{call_graph, Context, Function, Instruction, Kind, Module, Value};
 
 pub use semantic_analysis::namespace::{self, Namespace};
 pub mod types;
@@ -506,13 +506,16 @@ fn inline_function_calls(
         false
     };
 
+    let cg = call_graph::build_call_graph(ir, functions);
+    let functions = call_graph::callee_first_order(ir, &cg);
+
     for function in functions {
         if let Err(ir_error) = match tree_type {
             parsed::TreeType::Predicate => {
                 // Inline everything for predicates
-                sway_ir::optimize::inline_all_function_calls(ir, function)
+                sway_ir::optimize::inline_all_function_calls(ir, &function)
             }
-            _ => sway_ir::optimize::inline_some_function_calls(ir, function, inline_heuristic),
+            _ => sway_ir::optimize::inline_some_function_calls(ir, &function, inline_heuristic),
         } {
             return err(
                 Vec::new(),

sway-ir/src/analysis.rs

@@ -1,2 +1,4 @@
+pub mod call_graph;
+pub use call_graph::*;
 pub mod dominator;
 pub use dominator::*;

sway-ir/src/analysis/call_graph.rs

@@ -0,0 +1,54 @@
+/// Build call graphs for the program being compiled.
+/// If a function F1 calls function F2, then the call
+/// graph has an edge F1->F2.
+use crate::{Context, Function, Instruction, ValueDatum};
+
+use rustc_hash::{FxHashMap, FxHashSet};
+
+pub type CallGraph = FxHashMap<Function, FxHashSet<Function>>;
+
+/// Build call graph considering all providing functions.
+pub fn build_call_graph(ctx: &Context, functions: &[Function]) -> CallGraph {
+    let mut res = CallGraph::default();
+    for function in functions {
+        let entry = res.entry(*function);
+        let entry = entry.or_insert_with(FxHashSet::default);
+        for (_, inst) in function.instruction_iter(ctx) {
+            if let ValueDatum::Instruction(Instruction::Call(callee, _)) = ctx.values[inst.0].value
+            {
+                entry.insert(callee);
+            }
+        }
+    }
+    res
+}
+
+/// Given a call graph, return reverse topological sort
+/// (post order traversal), i.e., If A calls B, then B
+/// occurs before A in the returned Vec.
+pub fn callee_first_order(ctx: &Context, cg: &CallGraph) -> Vec<Function> {
+    let mut res = Vec::new();
+
+    let mut visited = FxHashSet::<Function>::default();
+    fn post_order_visitor(
+        ctx: &Context,
+        cg: &CallGraph,
+        visited: &mut FxHashSet<Function>,
+        res: &mut Vec<Function>,
+        node: Function,
+    ) {
+        if visited.contains(&node) {
+            return;
+        }
+        visited.insert(node);
+        for callee in &cg[&node] {
+            post_order_visitor(ctx, cg, visited, res, *callee);
+        }
+        res.push(node);
+    }
+    for node in cg.keys() {
+        post_order_visitor(ctx, cg, &mut visited, &mut res, *node);
+    }
+
+    res
+}

sway-ir/src/optimize/inline.rs

@@ -2,7 +2,9 @@
 //!
 //! Function inlining is pretty hairy so these passes must be maintained with care.
 
-use std::collections::HashMap;
+use std::{cell::RefCell, collections::HashMap};
+
+use rustc_hash::FxHashMap;
 
 use crate::{
     asm::AsmArg,
@@ -44,30 +46,33 @@ pub fn inline_some_function_calls<F: Fn(&Context, &Function, &Value) -> bool>(
     function: &Function,
     predicate: F,
 ) -> Result<bool, IrError> {
-    let mut modified = false;
-    loop {
-        // Find the next call site which passes the predicate.
-        let call_data = function
-            .instruction_iter(context)
-            .find_map(|(block, call_val)| match context.values[call_val.0].value {
-                ValueDatum::Instruction(Instruction::Call(inlined_function, _)) => predicate(
-                    context,
-                    &inlined_function,
-                    &call_val,
-                )
-                .then_some((block, call_val, inlined_function)),
-                _ => None,
-            });
-
-        match call_data {
-            Some((block, call_val, inlined_function)) => {
-                inline_function_call(context, *function, block, call_val, inlined_function)?;
-                modified = true;
+    // Find call sites which passes the predicate.
+    // We use a RefCell so that the inliner can modify the value
+    // when it moves other instructions (which could be in call_date) after an inline.
+    let call_data: FxHashMap<Value, RefCell<(Block, Function)>> = function
+        .instruction_iter(context)
+        .filter_map(|(block, call_val)| match context.values[call_val.0].value {
+            ValueDatum::Instruction(Instruction::Call(inlined_function, _)) => {
+                predicate(context, &inlined_function, &call_val)
+                    .then_some((call_val, RefCell::new((block, inlined_function))))
             }
-            None => break,
-        }
+            _ => None,
+        })
+        .collect();
+
+    for (call_site, call_site_in) in &call_data {
+        let (block, inlined_function) = *call_site_in.borrow();
+        inline_function_call(
+            context,
+            *function,
+            block,
+            *call_site,
+            inlined_function,
+            &call_data,
+        )?;
     }
-    Ok(modified)
+
+    Ok(!call_data.is_empty())
 }
 
 /// A utility to get a predicate which can be passed to inline_some_function_calls() based on
@@ -139,6 +144,7 @@ pub fn inline_function_call(
     block: Block,
     call_site: Value,
     inlined_function: Function,
+    call_data: &FxHashMap<Value, RefCell<(Block, Function)>>,
 ) -> Result<(), IrError> {
     // Split the block at right after the call site.
     let call_site_idx = context.blocks[block.0]
@@ -147,6 +153,19 @@ pub fn inline_function_call(
         .position(|&v| v == call_site)
         .unwrap();
     let (pre_block, post_block) = block.split_at(context, call_site_idx + 1);
+    if post_block != block {
+        // We need to update call_data for every call_site that was in block.
+        for inst in post_block.instruction_iter(context).filter(|inst| {
+            matches!(
+                context.values[inst.0].value,
+                ValueDatum::Instruction(Instruction::Call(..))
+            )
+        }) {
+            if let Some(call_info) = call_data.get(&inst) {
+                call_info.borrow_mut().0 = post_block;
+            }
+        }
+    }
 
     // Remove the call from the pre_block instructions.  It's still in the context.values[] though.
     context.blocks[pre_block.0].instructions.pop();

test/src/e2e_vm_tests/test_programs/should_pass/require_contract_deployment/call_basic_storage/src/main.sw

@@ -3,7 +3,7 @@ use basic_storage_abi::{StoreU64, Quad};
 use std::assert::assert;
 
 fn main() -> u64 {
-    let addr = abi(StoreU64, 0x5c452aa39ace5925513ac6a078eee827fc78245258c71ed9851ffe2251719b9d);
+    let addr = abi(StoreU64, 0x38cc2a8e447e57c335ffea24407527070d0b9f63c51358aca726fdd47b45a592);
     let key = 0xffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff;
     let value = 4242;