From acded67b3ec77ce1753356ad46f7ae17290f2ee0 Mon Sep 17 00:00:00 2001
From: Vaivaswatha N <vaivaswatha.nagaraj@fuel.sh>
Date: Thu, 4 Jul 2024 19:20:37 +0530
Subject: [PATCH] Inliner shouldn't inline a function into itself (#6218)

Also update the testsuite inliner predicate to respect inline
annotations.

Closes #6149
---
 sway-ir/src/optimize/inline.rs      | 10 +++-
 sway-ir/tests/inline/recursive-1.ir | 71 +++++++++++++++++++++++++++++
 sway-ir/tests/tests.rs              | 26 +++++++----
 3 files changed, 95 insertions(+), 12 deletions(-)
 create mode 100644 sway-ir/tests/inline/recursive-1.ir

diff --git a/sway-ir/src/optimize/inline.rs b/sway-ir/src/optimize/inline.rs
index b4ef08d3155..4a3e97302f6 100644
--- a/sway-ir/src/optimize/inline.rs
+++ b/sway-ir/src/optimize/inline.rs
@@ -35,12 +35,12 @@ pub fn create_fn_inline_pass() -> Pass {
 /// This is a copy of sway_core::inline::Inline.
 /// TODO: Reuse: Depend on sway_core? Move it to sway_types?
 #[derive(Debug)]
-enum Inline {
+pub enum Inline {
     Always,
     Never,
 }
 
-fn metadata_to_inline(context: &Context, md_idx: Option<MetadataIndex>) -> Option<Inline> {
+pub fn metadata_to_inline(context: &Context, md_idx: Option<MetadataIndex>) -> Option<Inline> {
     fn for_each_md_idx<T, F: FnMut(MetadataIndex) -> Option<T>>(
         context: &Context,
         md_idx: Option<MetadataIndex>,
@@ -189,6 +189,12 @@ pub fn inline_some_function_calls<F: Fn(&Context, &Function, &Value) -> bool>(
     for call_site in &call_sites {
         let call_site_in = call_data.get(call_site).unwrap();
         let (block, inlined_function) = *call_site_in.borrow();
+
+        if function == &inlined_function {
+            // We can't inline a function into itself.
+            continue;
+        }
+
         inline_function_call(
             context,
             *function,
diff --git a/sway-ir/tests/inline/recursive-1.ir b/sway-ir/tests/inline/recursive-1.ir
new file mode 100644
index 00000000000..b7dc13d62f3
--- /dev/null
+++ b/sway-ir/tests/inline/recursive-1.ir
@@ -0,0 +1,71 @@
+//
+// regex: VAR=v\d+
+// regex: LABEL=[[:alpha:]0-9_]+
+// regex: PING=(ping|id_from_ping)
+// regex: PONG=(pong|id_from_pong)
+script {
+
+    fn id_from_foo(b: u64) -> u64, !1 {
+       entry(b: u64):
+       ret u64 b
+    }
+
+    // check: fn foo
+    fn foo(b: u64) -> u64 {
+        entry(b: u64):
+
+        // check: call id_from_foo
+        v1 = call id_from_foo(b)
+        // check: call foo
+        v0 = call foo(v1)
+        ret u64 v0
+    }
+
+    fn id_from_ping(b: u64) -> u64, !1 {
+       entry(b: u64):
+       ret u64 b
+    }
+
+    fn id_from_pong(b: u64) -> u64, !1 {
+       entry(b: u64):
+       ret u64 b
+    }
+
+    // check: fn main
+    fn main() -> u64 {
+        entry():
+
+        v0 = const u64 11
+        // check: call foo
+        v1 = call foo(v0)
+
+        // check: $PING
+        v2 = call ping(v1)
+        v3 = add v1, v2
+
+        ret u64 v3
+    }
+
+    // check: fn ping
+    fn ping(b: u64) -> u64 {
+        entry(b: u64):
+
+        // check: id_from_ping
+        v1 = call id_from_ping(b)
+        v0 = call pong(v1)
+        ret u64 v0
+    }
+
+    // check: fn pong
+    fn pong(b: u64) -> u64 {
+        entry(b: u64):
+
+        // check: $PONG
+        v1 = call id_from_pong(b)
+        v0 = call ping(v1)
+        ret u64 v0
+    }
+
+}
+
+!1 = inline "never"
diff --git a/sway-ir/tests/tests.rs b/sway-ir/tests/tests.rs
index ed7a745ad73..54bc88176c4 100644
--- a/sway-ir/tests/tests.rs
+++ b/sway-ir/tests/tests.rs
@@ -5,9 +5,10 @@ use sway_ir::{
     create_arg_demotion_pass, create_const_demotion_pass, create_const_folding_pass,
     create_dce_pass, create_dom_fronts_pass, create_dominators_pass, create_escaped_symbols_pass,
     create_mem2reg_pass, create_memcpyopt_pass, create_misc_demotion_pass, create_postorder_pass,
-    create_ret_demotion_pass, create_simplify_cfg_pass, optimize as opt, register_known_passes,
-    Context, ExperimentalFlags, IrError, PassGroup, PassManager, DCE_NAME, FN_DCE_NAME,
-    FN_DEDUP_DEBUG_PROFILE_NAME, FN_DEDUP_RELEASE_PROFILE_NAME, MEM2REG_NAME, SROA_NAME,
+    create_ret_demotion_pass, create_simplify_cfg_pass, metadata_to_inline, optimize as opt,
+    register_known_passes, Context, ExperimentalFlags, Function, IrError, PassGroup, PassManager,
+    Value, DCE_NAME, FN_DCE_NAME, FN_DEDUP_DEBUG_PROFILE_NAME, FN_DEDUP_RELEASE_PROFILE_NAME,
+    MEM2REG_NAME, SROA_NAME,
 };
 use sway_types::SourceEngine;
 
@@ -201,13 +202,18 @@ fn inline() {
                     );
 
             funcs.into_iter().fold(false, |acc, func| {
-                opt::inline_some_function_calls(
-                    ir,
-                    &func,
-                    opt::is_small_fn(max_blocks, max_instrs, max_stack),
-                )
-                .unwrap()
-                    || acc
+                let predicate = |context: &Context, function: &Function, call_site: &Value| {
+                    let attributed_inline =
+                        metadata_to_inline(context, function.get_metadata(context));
+                    match attributed_inline {
+                        Some(opt::Inline::Never) => false,
+                        Some(opt::Inline::Always) => true,
+                        None => (opt::is_small_fn(max_blocks, max_instrs, max_stack))(
+                            context, function, call_site,
+                        ),
+                    }
+                };
+                opt::inline_some_function_calls(ir, &func, predicate).unwrap() || acc
             })
         }
     })