From 519b1e907dc2f49be3f33fe8e55e05cea52ff97d Mon Sep 17 00:00:00 2001
From: Arnold Schwaighofer <aschwaighofer@apple.com>
Date: Sat, 15 May 2021 07:44:27 -0700
Subject: [PATCH] [5.5] Make sure we tail call optimize a call in concurrency
 runtime's switch_task_impl.

Without this hack the call will leave a stack frame around (not tail
call optimized) and blow the stack if we call switch_task often enough.

Ideally, clang would emit this call as `musttail` but currently it does
not.

rdar://76652421
---
 stdlib/public/Concurrency/Actor.cpp | 11 +++++-
 test/Interpreter/async_fib.swift    | 60 +++++++++++++++++++++++++++++
 2 files changed, 70 insertions(+), 1 deletion(-)
 create mode 100644 test/Interpreter/async_fib.swift

diff --git a/stdlib/public/Concurrency/Actor.cpp b/stdlib/public/Concurrency/Actor.cpp
index 7b4bae5c8733d..c3cba1e414517 100644
--- a/stdlib/public/Concurrency/Actor.cpp
+++ b/stdlib/public/Concurrency/Actor.cpp
@@ -1802,6 +1802,13 @@ static bool tryAssumeThreadForSwitch(ExecutorRef newExecutor,
   return false;
 }
 
+__attribute__((noinline))
+SWIFT_CC(swiftasync)
+static void force_tail_call_hack(AsyncTask *task) {
+  // This *should* be executed as a tail call.
+  return task->runInFullyEstablishedContext();
+}
+
 /// Given that we've assumed control of an executor on this thread,
 /// continue to run the given task on it.
 SWIFT_CC(swiftasync)
@@ -1815,7 +1822,9 @@ static void runOnAssumedThread(AsyncTask *task, ExecutorRef executor,
     oldTracking->setActiveExecutor(executor);
 
     // FIXME: force tail call
-    return task->runInFullyEstablishedContext();
+    // return task->runInFullyEstablishedContext();
+    // This hack "ensures" that this call gets executed as a tail call.
+    return force_tail_call_hack(task);
   }
 
   // Otherwise, set up tracking info.
diff --git a/test/Interpreter/async_fib.swift b/test/Interpreter/async_fib.swift
new file mode 100644
index 0000000000000..a447cc24dd10d
--- /dev/null
+++ b/test/Interpreter/async_fib.swift
@@ -0,0 +1,60 @@
+// RUN: %empty-directory(%t)
+// RUN: %target-build-swift -Xfrontend -enable-experimental-concurrency %s -parse-as-library -module-name main -o %t/main
+// RUN: %target-codesign %t/main
+// RUN: %target-run %t/main | %FileCheck %s
+
+// REQUIRES: concurrency
+// REQUIRES: executable_test
+// UNSUPPORTED: use_os_stdlib
+// UNSUPPORTED: back_deployment_runtime
+// UNSUPPORTED: OS=windows-msvc
+
+var gg = 0
+
+@inline(never)
+public func identity<T>(_ x: T) -> T {
+  gg += 1
+  return x
+}
+
+actor Actor {
+  var x: Int = 0
+  init(x: Int) { self.x = x }
+
+  @inline(never)
+  func get(_ i: Int ) async -> Int {
+    return i + x
+  }
+}
+
+// Used to crash with an stack overflow with m >= 18
+let m = 22
+
+@inline(never)
+func asyncFib(_ n: Int, _ a1: Actor, _ a2: Actor) async -> Int {
+  if n == 0 {
+    return await a1.get(n)
+  }
+  if n == 1 {
+    return await a2.get(n)
+  }
+
+  let first = await asyncFib(n-2, a1, a2)
+  let second = await asyncFib(n-1, a1, a2)
+
+  let result = first + second
+
+  return result
+}
+
+@main struct Main {
+  static func main() async {
+    let a1 = Actor(x: 0)
+    let a2 = Actor(x: 0)
+    _ = await asyncFib(identity(m), a1, a2)
+
+    // CHECK: result: 0
+    await print("result: \(a1.x)");
+    await print(a2.x)
+  }
+}