Tail calls not respected after inlining

[ospencer]

I ran into a bit of an issue with return_call instructions not being honored after inlining. This can cause code that otherwise works without blowing the stack to cause a stack overflow after being optimized.

A fairly small reproduction is available here (and my apologies for the OCaml 😶): https://github.com/ospencer/binaryen_tailcall_bug_repro/blob/main/lib/Util.re

If you have esy installed, you can run the example with esy && esy test, and you'll get this output:

Unoptimized Binaryen IR:
(module
 (type $i32_=>_i32 (func (param i32) (result i32)))
 (type $none_=>_i32 (func (result i32)))
 (export "wrapper" (func $wrapper))
 (func $is_odd (param $0 i32) (result i32)
  (if (result i32)
   (i32.le_u
    (local.get $0)
    (i32.const 1)
   )
   (i32.eq
    (local.get $0)
    (i32.const 1)
   )
   (return_call $is_even
    (i32.sub
     (local.get $0)
     (i32.const 1)
    )
   )
  )
 )
 (func $is_even (param $0 i32) (result i32)
  (if (result i32)
   (i32.le_u
    (local.get $0)
    (i32.const 1)
   )
   (i32.eq
    (local.get $0)
    (i32.const 0)
   )
   (return_call $is_odd
    (i32.sub
     (local.get $0)
     (i32.const 1)
    )
   )
  )
 )
 (func $wrapper (result i32)
  (call $is_odd
   (i32.const 99999)
  )
 )
)

Optimized Binaryen IR:
(module
 (type $i32_=>_i32 (func (param i32) (result i32)))
 (type $none_=>_i32 (func (result i32)))
 (export "wrapper" (func $wrapper))
 (func $is_odd (; has Stack IR ;) (param $0 i32) (result i32)
  (if (result i32)
   (i32.le_u
    (local.get $0)
    (i32.const 1)
   )
   (i32.eq
    (local.get $0)
    (i32.const 1)
   )
   (if (result i32)
    (i32.le_u
     (local.tee $0
      (i32.sub
       (local.get $0)
       (i32.const 1)
      )
     )
     (i32.const 1)
    )
    (i32.eqz
     (local.get $0)
    )
    (call $is_odd
     (i32.sub
      (local.get $0)
      (i32.const 1)
     )
    )
   )
  )
 )
 (func $wrapper (; has Stack IR ;) (result i32)
  (call $is_odd
   (i32.const 99999)
  )
 )
)

I haven't looked at the code just yet, but my intuition is that this happens because it wouldn't be safe to copy over a return_call in the body of a function, as it likely would change the semantics of the function it's being inlined into. I can't really think of something more clever than just not inlining functions with a return_call instruction, though in "simple" cases like this one where the result of the inlined function is returned, the return_call could be kept, but that analysis seems a bit annoying to write 😅

Let me know if there's any additional information I can provide or any way I can help.

[kripken]

This does look like a bug. To see it it's enough to run wasm-opt -all --inlining on the input testcase. The output has a normal call and not a call_return

The pass does have logic that should handle return_call safely, but it's not complete I guess,

binaryen/src/passes/Inlining.cpp

Lines 254 to 282 in 094deb0

	// Return calls in inlined functions should only break out of the scope of
	// the inlined code, not the entire function they are being inlined into. To
	// achieve this, make the call a non-return call and add a break. This does
	// 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) {
	curr->isReturn = false;
	curr->type = targetType.getSignature().results;
	if (curr->type.isConcrete()) {
	replaceCurrent(builder->makeBreak(returnName, curr));
	} else {
	replaceCurrent(builder->blockify(curr, builder->makeBreak(returnName)));
	}
	}
	void visitCall(Call* curr) {
	if (curr->isReturn) {
	handleReturnCall(curr, module->getFunction(curr->target)->type);
	}
	}
	void visitCallIndirect(CallIndirect* curr) {
	if (curr->isReturn) {
	handleReturnCall(curr, curr->heapType);
	}
	}
	void visitCallRef(CallRef* curr) {
	if (curr->isReturn) {
	handleReturnCall(curr, curr->target->type.getHeapType());
	}
	}

cc @tlively

[tlively]

That code failed to anticipate the case where there is a "loop" of (mutually) recursive tail calls. I think it would work to only downgrade return calls in the callee to normal calls if and only if the inlined call site is not a return call. If the inlined call site is not a return call, then this recursion could have blown the stack anyway, and if it is a return call, then the types must be compatible such that we don't need to downgrade the return calls in the callee.