Tail Recursion

[TheAngryByrd]

Is your feature request related to a problem? Please describe.

Tasks currently do not support tail call optimization. If we want to support more features like async does, then we should.

Describe the solution you'd like

I'm not smart enough to know what this looks like yet.

Describe alternatives you've considered

Don't do it. There's always an option not to.

Additional context
Add any other context or screenshots about the feature request here.

[TheAngryByrd]

I looked into this over Christmas and it seems to be the case that is isn't simple because ReturnFrom doesn't know if we're coming from our own state machine builder or another Task/TaskLike and can't use hijack to call MoveNext on the statemachine.. I think we have to wait for F# compiler to implement fsharp/fslang-suggestions#1200

[majocha]

The difference with legacy async builder is that async handles safely all recursive and mutually recursive binds, not just tail recursion. It would be great to have parity in this regard.

As a simple example this works with async but overflows stack with cancellableTask:

        let rec foo n =
            cancellableTask {
                if n = 0 then
                    return 0
                else
                    let! result = foo (n - 1)
                    return result
            }
        let! actual = foo 10_000
        Expect.equal actual 0 ""

I've been wondering where exactly the SO arises in resumable code execution? How is it different from async trampolined execution? It would seem calling MoveNext() on a state machine is conceptually not that different from trampolining?

I admit I have a very vague idea how this all works, but my suspicion is that binds cause SO here. It seems in a recursive code like above, each bind synchronously starts another state machine, executes it's MoveNext and so on:

Stack overflow.
   at IcedTasks.Tests.CancellableTaskTests.foo@1271-18(Int32)
   at IcedTasks.Tests.CancellableTaskTests+foo@1272-19.MoveNext()
   at System.Runtime.CompilerServices.AsyncMethodBuilderCore.Start[[IcedTasks.Tests.CancellableTaskTests+foo@1272-19, IcedTasks.Tests, Version=0.0.0.0, Culture=neutral, PublicKeyToken=null]](foo@1272-19 ByRef)
   at System.Runtime.CompilerServices.AsyncTaskMethodBuilder`1[[System.Int32, System.Private.CoreLib, Version=8.0.0.0, Culture=neutral, PublicKeyToken=7cec85d7bea7798e]].Start[[IcedTasks.Tests.CancellableTaskTests+foo@1272-19, IcedTasks.Tests, Version=0.0.0.0, Culture=neutral, PublicKeyToken=null]](foo@1272-19 ByRef)
   at IcedTasks.Tests.CancellableTaskTests+foo@1272-21.Invoke(System.Threading.CancellationToken)
   at IcedTasks.Tests.CancellableTaskTests+foo@1276-20.Invoke(System.Threading.CancellationToken)
   at IcedTasks.Tests.CancellableTaskTests+foo@1272-19.MoveNext()
   at System.Runtime.CompilerServices.AsyncMethodBuilderCore.Start[[IcedTasks.Tests.CancellableTaskTests+foo@1272-19, IcedTasks.Tests, Version=0.0.0.0, Culture=neutral, PublicKeyToken=null]](foo@1272-19 ByRef)
   at System.Runtime.CompilerServices.AsyncTaskMethodBuilder`1[[System.Int32, System.Private.CoreLib, Version=8.0.0.0, Culture=neutral, PublicKeyToken=7cec85d7bea7798e]].Start[[IcedTasks.Tests.CancellableTaskTests+foo@1272-19, IcedTasks.Tests, Version=0.0.0.0, Culture=neutral, PublicKeyToken=null]](foo@1272-19 ByRef)
   at IcedTasks.Tests.CancellableTaskTests+foo@1272-21.Invoke(System.Threading.CancellationToken)
   at IcedTasks.Tests.CancellableTaskTests+foo@1276-20.Invoke(System.Threading.CancellationToken)
   at IcedTasks.Tests.CancellableTaskTests+foo@1272-19.MoveNext()
etc. etc.

Interestingly, the stack overflow in the example above goes away when I wrap the builder starting code in a Task.Run like this:

                    (AfterCode<_, _>(fun sm ->
                        let sm = sm

                        fun (ct) ->
                            if ct.IsCancellationRequested then
                                Task.FromCanceled<_>(ct)
                            else
                                Task.Run<'T>((fun () ->
                                    let mutable sm = sm
                                    sm.Data.CancellationToken <- ct
                                    sm.Data.MethodBuilder <- AsyncTaskMethodBuilder<'T>.Create()
                                    sm.Data.MethodBuilder.Start(&sm)
                                    sm.Data.MethodBuilder.Task), ct) : Task<'T>
                    ))

So that each Start is not synchronous but unconditionally queued on the threadpool.

[TheAngryByrd]

Yeah Task.Run in the return position essentially a stack guard technique that F# compiler uses for synchronous recursion.

[majocha]

I'm looking at the coroutine example for inspiration.

It seems, one way to make all binds safe from SO would be to:

1. Delay starting of bound computation.
2. Store the reference to the bound computation in current state machine data just like the tailcall target in example.
3. Start it from there and make sure the continuation resumes on completed.

That way instead of growing stack, we would grow a linked list of recursively bound computations. (?)

[majocha]

A less invasive alternative would be to implement the stackguard:

1. Just keep a bind count in the state machine data
2. In Bind method conditionally wrap getAwaiter in a Task.Run

[TheAngryByrd]

I'm looking at the coroutine example for inspiration

My earlier comment describes why #3 (comment) we don't. We'd have to define our Task type. Async2 and Coroutine do this.

You could implement the a Stackguard like the one in F# in your own project, doesn't have to change the CEs in this project. But unfortunately, Task.Run isn't really the answer here either. This eats up tons of memory for deeper loops. It might be acceptable for your project, but wouldn't be good for this library to have.

If you really want to implement Tail Recursion I'd recommend asking about it at the compiler level.

[TheAngryByrd]

I found the old stash I had:

Stackguard.fs:

namespace IcedTasks

open System.Threading
open System.Threading.Tasks
open IcedTasks.Polyfill.Task

type StackGuard(maxDepth: int, name: string) =

    let mutable depth = 1

    member _.Guard(f: CancellationToken -> Task<'t>, ?cancellationToken: CancellationToken) =
        let cancellationToken = defaultArg cancellationToken CancellationToken.None

        task {
            Interlocked.Increment(&depth)
            |> ignore

            try
                if depth % maxDepth = 0 then
                    // printfn "Running %s on a new thread %d" name depth
                    return! Task.Run<'t>((fun () -> f cancellationToken), cancellationToken)
                else
                    return! f cancellationToken
            finally
                Interlocked.Decrement(&depth)
                |> ignore
        }

    member _.Guard(f: CancellationToken -> ValueTask<'t>, ?cancellationToken: CancellationToken) =
        let cancellationToken = defaultArg cancellationToken CancellationToken.None

        valueTask {
            Interlocked.Increment(&depth)
            |> ignore

            try
                if depth % maxDepth = 0 then
                    return!
                        Task.Run<'t>((fun () -> (f cancellationToken).AsTask()), cancellationToken)
                else
                    return! f cancellationToken
            finally
                Interlocked.Decrement(&depth)
                |> ignore
        }

    member _.Guard(f: CancellationToken -> Task, ?cancellationToken: CancellationToken) =
        let cancellationToken = defaultArg cancellationToken CancellationToken.None

        taskUnit {
            Interlocked.Increment(&depth)
            |> ignore

            try
                if depth % maxDepth = 0 then
                    return! Task.Run((fun () -> f cancellationToken), cancellationToken)
                else
                    return! f cancellationToken
            finally
                Interlocked.Decrement(&depth)
                |> ignore
        }

    member _.Guard(f: CancellationToken -> ValueTask, ?cancellationToken: CancellationToken) =
        let cancellationToken = defaultArg cancellationToken CancellationToken.None

        valueTaskUnit {
            Interlocked.Increment(&depth)
            |> ignore

            try
                if depth % maxDepth = 0 then
                    return! Task.Run((fun () -> (f cancellationToken).AsTask()), cancellationToken)
                else
                    return! f cancellationToken
            finally
                Interlocked.Decrement(&depth)
                |> ignore
        }

Tests.fs:

                testCaseAsync "TailRecursion - Happy Path"
                <| async {
                    let stackGuard = StackGuard(1000, "TailRecursion - Happy Path")

                    let rec loop (size: int) =
                        task {
                            return!
                                stackGuard.Guard(
                                    (fun _ ->
                                        task {
                                            do! Task.Yield()
                                            if size = 0 then return 0 else return! loop (size - 1)
                                        }
                                    ),
                                    CancellationToken.None
                                )
                        }

                    let size = 1000000

                    let! _ =
                        loop size
                        |> Async.AwaitTask

                    ()

                }

                // this will eat up a ton of memory trying to join stack traces together, something hitting OutOfMemory
                testCaseAsync "TailRecursion - Sad Path"
                <| async {
                    let stackGuard = StackGuard(100, "TailRecursion - Sad Path")

                    let rec loop (size: int) =
                        task {
                            return!
                                stackGuard.Guard(
                                    (fun _ ->
                                        task {
                                            do! Task.Yield()

                                            if size = 42 then
                                                failwith $"can't find meaning of life at {size}"

                                            if size = 0 then return 0 else return! loop (size - 1)
                                        }
                                    ),
                                    CancellationToken.None
                                )
                        }

                    let size = 1000000

                    let! _ =
                        loop size
                        |> Async.AwaitTask

                    ()

                }

I might consider adding this as a helper but I wouldn't want to push this automatically into a CE.

[majocha]

I've been thinking about it some more. It turns out it's trivial to install a synchronous trampoline like this:

type Trampoline() =
    let mutable next: ValueOption<Action> = ValueNone
    let mutable executing = false

    let rec loop () =
        match next with
        | ValueSome action ->
            next <- ValueNone
            action.Invoke()
            loop ()
        | ValueNone -> executing <- false

    let start action =
        next <- ValueSome action
        executing <- true
        loop ()

    let set action =
        if executing then
            next <- ValueSome action
        else
            start action

    member val Awaiter = 
        { new ICriticalNotifyCompletion with
            member this.OnCompleted(continuation: Action) = set continuation
            member this.UnsafeOnCompleted(continuation: Action) = set continuation }
        |> ref

Instead of yielding to the threadpool which is super slow, we can just suspend and do

sm.Data.Builder.AwaitUnsafeOnCompleted(ThisThreadTrampoline.Awaiter, &sm)

whenever the bind limit is reached. This handles both tail call and non tail call scenarios, just like old async builder.

[TheAngryByrd]

Let’s give it a try!