[Inlining] Preserve return_calls when possible (#4589)

We can preserve return_calls in inlined functions when the inlined call site is itself a return_call, since the call result types must transitively match in that case. This solves a problem where the previous inlining logic could introduce stack exhaustion by downgrading recursive return_calls to normal calls. Fixes #4587.
WebAssembly · Apr 11, 2022 · 5f88dcd · 5f88dcd
1 parent 5d5e465
commit 5f88dcd
Show file tree

Hide file tree

Showing 2 changed files with 72 additions and 0 deletions.
diff --git a/src/passes/Inlining.cpp b/src/passes/Inlining.cpp
@@ -247,6 +247,7 @@ struct Updater : public PostWalker<Updater> {
   Module* module;
   std::map<Index, Index> localMapping;
   Name returnName;
+  bool isReturn;
   Builder* builder;
   void visitReturn(Return* curr) {
     replaceCurrent(builder->makeBreak(returnName, curr->value));
@@ -257,6 +258,14 @@ struct Updater : public PostWalker<Updater> {
   // not cause unbounded stack growth because inlining and return calling both
   // avoid creating a new stack frame.
   template<typename T> void handleReturnCall(T* curr, HeapType targetType) {
+    if (isReturn) {
+      // If the inlined callsite was already a return_call, then we can keep
+      // return_calls in the inlined function rather than downgrading them.
+      // That is, if A->B and B->C and both those calls are return_calls
+      // then after inlining A->B we want to now have A->C be a
+      // return_call.
+      return;
+    }
     curr->isReturn = false;
     curr->type = targetType.getSignature().results;
     if (curr->type.isConcrete()) {
@@ -329,6 +338,7 @@ doInlining(Module* module, Function* into, const InliningAction& action) {
   Updater updater;
   updater.module = module;
   updater.returnName = block->name;
+  updater.isReturn = call->isReturn;
   updater.builder = &builder;
   // Set up a locals mapping
   for (Index i = 0; i < from->getNumLocals(); i++) {

diff --git a/test/lit/passes/inlining_enable-tail-call.wast b/test/lit/passes/inlining_enable-tail-call.wast
@@ -705,3 +705,65 @@
   (unreachable)
  )
 )
+
+(module
+ ;; CHECK:      (type $i32_=>_i32 (func (param i32) (result i32)))
+
+ ;; CHECK:      (export "is_even" (func $is_even))
+ (export "is_even" (func $is_even))
+ ;; CHECK:      (func $is_even (param $i i32) (result i32)
+ ;; CHECK-NEXT:  (local $1 i32)
+ ;; CHECK-NEXT:  (if (result i32)
+ ;; CHECK-NEXT:   (i32.eqz
+ ;; CHECK-NEXT:    (local.get $i)
+ ;; CHECK-NEXT:   )
+ ;; CHECK-NEXT:   (i32.const 1)
+ ;; CHECK-NEXT:   (return
+ ;; CHECK-NEXT:    (block $__inlined_func$is_odd (result i32)
+ ;; CHECK-NEXT:     (local.set $1
+ ;; CHECK-NEXT:      (i32.sub
+ ;; CHECK-NEXT:       (local.get $i)
+ ;; CHECK-NEXT:       (i32.const 1)
+ ;; CHECK-NEXT:      )
+ ;; CHECK-NEXT:     )
+ ;; CHECK-NEXT:     (if (result i32)
+ ;; CHECK-NEXT:      (i32.eqz
+ ;; CHECK-NEXT:       (local.get $1)
+ ;; CHECK-NEXT:      )
+ ;; CHECK-NEXT:      (i32.const 0)
+ ;; CHECK-NEXT:      (return_call $is_even
+ ;; CHECK-NEXT:       (i32.sub
+ ;; CHECK-NEXT:        (local.get $1)
+ ;; CHECK-NEXT:        (i32.const 1)
+ ;; CHECK-NEXT:       )
+ ;; CHECK-NEXT:      )
+ ;; CHECK-NEXT:     )
+ ;; CHECK-NEXT:    )
+ ;; CHECK-NEXT:   )
+ ;; CHECK-NEXT:  )
+ ;; CHECK-NEXT: )
+ (func $is_even (param $i i32) (result i32)
+  (if (result i32)
+   (i32.eqz (local.get $i))
+   (i32.const 1)
+   (return_call $is_odd
+    (i32.sub
+     (local.get $i)
+     (i32.const 1)
+    )
+   )
+  )
+ )
+ (func $is_odd (param $i i32) (result i32)
+  (if (result i32)
+   (i32.eqz (local.get $i))
+   (i32.const 0)
+   (return_call $is_even
+    (i32.sub
+     (local.get $i)
+     (i32.const 1)
+    )
+   )
+  )
+ )
+)