bit simpler

Author: majocha

src/Compiler/Utilities/Cancellable.fs

@@ -85,16 +85,17 @@ and CancellableResumptionDynamicInfo<'TOverall> = ResumptionDynamicInfo<Cancella
 and CancellableCode<'TOverall, 'T> = ResumableCode<CancellableStateMachineData<'TOverall>, 'T>
 
 [<Sealed>]
-type Trampoline(cancellationToken: CancellationToken) =
+type Trampoline() =
 
     let mutable bindDepth = 0
+    let cancellationToken = FSharp.Compiler.Cancellable.Token
 
     [<Literal>]
     static let bindDepthLimit = 100
 
-    static let current = new ThreadLocal<Trampoline>()
+    static let current = new AsyncLocal<Trampoline voption>()
 
-    let delayed = System.Collections.Generic.Stack<ITrampolineInvocation>()
+    let pending = System.Collections.Generic.Stack<ITrampolineInvocation>()
 
     member this.IsCancelled = cancellationToken.IsCancellationRequested
 
@@ -103,15 +104,13 @@ type Trampoline(cancellationToken: CancellationToken) =
 
     member this.ShoudBounce = bindDepth % bindDepthLimit = 0
 
-    static member Install ct = current.Value <- Trampoline ct
-
     member val LastError: ExceptionDispatchInfo voption = ValueNone with get, set
 
     member this.RunDelayed(continuation, invocation) =
         // The calling state machine is now suspended. We need to resume it next.
-        delayed.Push continuation
+        pending.Push continuation
         // Schedule the delayed invocation to run.
-        delayed.Push invocation
+        pending.Push invocation
 
     member this.RunImmediate(invocation: ITrampolineInvocation) =
         bindDepth <- bindDepth + 1
@@ -121,16 +120,19 @@ type Trampoline(cancellationToken: CancellationToken) =
             invocation.MoveNext()
 
             while not invocation.IsCompleted do
-                if delayed.Peek().IsCompleted then
-                    delayed.Pop() |> ignore
+                if pending.Peek().IsCompleted then
+                    pending.Pop() |> ignore
                 else
-                    delayed.Peek().MoveNext()
+                    pending.Peek().MoveNext()
             // In case this was a delayed invocation, which captures the exception, we need to replay it.
             invocation.ReplayExceptionIfStored()
         finally
             bindDepth <- bindDepth - 1
 
-    static member Current = current.Value
+    static member Current = current.Value.Value
+
+    static member EnsureInstalled() =
+        if current.Value.IsNone then current.Value <- ValueSome <| Trampoline()
 
 type ITrampolineInvocation<'T> =
     inherit ITrampolineInvocation
@@ -354,22 +356,20 @@ module CancellableAutoOpens =
 module Cancellable =
     open Internal.Utilities.Library.CancellableImplementation
 
-    let run ct (code: Cancellable<_>) =
-        use _ = FSharp.Compiler.Cancellable.UsingToken ct
-
+    let run (code: Cancellable<_>) =
         let invocation = code.GetInvocation(false)
-        Trampoline.Install ct
+        Trampoline.EnsureInstalled()
         Trampoline.Current.RunImmediate invocation
-        invocation
+        invocation.Result
 
     let runWithoutCancellation code =
-        code |> run CancellationToken.None |> _.Result
+        use _ = FSharp.Compiler.Cancellable.UsingToken CancellationToken.None
+        run code
 
     let toAsync code =
         async {
-            let! ct = Async.CancellationToken
-
-            return run ct code |> _.Result
+            use! _holder = FSharp.Compiler.Cancellable.UseToken()
+            return run code
         }
 
     let token () =

src/Compiler/Utilities/Cancellable.fsi

@@ -115,7 +115,6 @@ module internal CancellableAutoOpens =
     val cancellable: CancellableImplementation.CancellableBuilder
 
 module internal Cancellable =
-    val run: ct: CancellationToken -> code: Cancellable<'a> -> ITrampolineInvocation<'a>
     val runWithoutCancellation: code: Cancellable<'a> -> 'a
     val toAsync: code: Cancellable<'a> -> Async<'a>
     val token: unit -> Cancellable<CancellationToken>