From 1865594231fc55b3b7a3228a74f16e05b1694a29 Mon Sep 17 00:00:00 2001
From: Martin <mautner@vrvis.at>
Date: Wed, 5 Jul 2023 15:33:14 +0200
Subject: [PATCH] More robust host invalid invocation detection

---
 src/Libs/FShade.Imperative/Utilities.fs  | 42 +++++++++++++++++++-----
 src/Tests/FShade.Core.Tests/Optimizer.fs | 35 +++++++++++++++++++-
 2 files changed, 68 insertions(+), 9 deletions(-)

diff --git a/src/Libs/FShade.Imperative/Utilities.fs b/src/Libs/FShade.Imperative/Utilities.fs
index 0d45f51..5033081 100644
--- a/src/Libs/FShade.Imperative/Utilities.fs
+++ b/src/Libs/FShade.Imperative/Utilities.fs
@@ -16,6 +16,7 @@ exception FShadeOnlyInShaderCodeException of string
 
 [<AutoOpen>]
 module DynamicHostInvocation =
+    open System.Collections.Generic
 
     let rec (|ShaderOnlyExn|_|) (e : Exception) =
         match e with
@@ -35,22 +36,47 @@ module DynamicHostInvocation =
         | _ ->
             None
 
-    let private invalidHostInvocation = new ThreadLocal<bool>(fun _ -> false)
+    let private invalidHostInvocation = new ThreadLocal<_>(fun _ -> Stack<bool>())
 
-    /// Returns a default value and sets a flag to indicate that it must not be used by the optimizer.
-    let onlyInShaderCode<'T> (_name : string) : 'T =
-        invalidHostInvocation.Value <- true
-        Unchecked.defaultof<'T>
+    /// <summary>
+    /// Returns a default value and signals that it must not be used by the optimizer.
+    /// </summary>
+    /// <exception cref="FShadeOnlyInShaderCodeException">Thrown when invoked directly.</exception>
+    let onlyInShaderCode<'T> (name : string) : 'T =
+        let invalid = invalidHostInvocation.Value
+
+        if invalid.Count = 0 then
+            raise <| FShadeOnlyInShaderCodeException name
+        else
+            // Replace the flag to indicate to the optimizer that the result is invalid.
+            // We use a stack rather than a simple flag to avoid potential issues with APC.
+            // See: https://stackoverflow.com/questions/8431221/why-did-entering-a-lock-on-a-ui-thread-trigger-an-onpaint-event
+            invalid.Pop() |> ignore
+            invalid.Push true
+            Unchecked.defaultof<'T>
 
     let inline private tryInvoke (f : unit -> 'T) =
+        let invalid = invalidHostInvocation.Value
+
         try
-            invalidHostInvocation.Value <- false
+            invalid.Push false
             let result = f()
-            if invalidHostInvocation.Value then None
+
+            if invalid.Count > 1 then
+                Log.warn "[FShade] Stack for invalid host invocations has %d values (should be 1)" invalid.Count
+
+            if invalid.Pop() then None
             else Some result
-        with ShaderOnlyExn _ ->
+
+        with
+        | ShaderOnlyExn _ ->
+            invalid.Pop() |> ignore
             None
 
+        | _ ->
+            invalid.Pop() |> ignore
+            reraise()
+
     type MethodInfo with
         member x.TryInvoke(obj : obj, parameters : obj[]) =
             tryInvoke (fun _ -> x.Invoke(obj, parameters))
diff --git a/src/Tests/FShade.Core.Tests/Optimizer.fs b/src/Tests/FShade.Core.Tests/Optimizer.fs
index 2aa8218..d8682e0 100644
--- a/src/Tests/FShade.Core.Tests/Optimizer.fs
+++ b/src/Tests/FShade.Core.Tests/Optimizer.fs
@@ -519,4 +519,37 @@ let ``[Constant] enum to float``() =
 let ``[Constant] enum to float32``() =
     let input    = <@ keep (float32 MyEnum.C) @>
     let expected = <@ keep 4.0f @>
-    input |> Opt.run |> should exprEqual expected
\ No newline at end of file
+    input |> Opt.run |> should exprEqual expected
+
+let private myCoolIntrinsic (_value : float) =
+    onlyInShaderCode<float> "myCoolIntrinsic"
+
+let private myCoolIntrinsic2 (_value : float) =
+    raise <| FShadeOnlyInShaderCodeException "myCoolIntrinsic"
+
+let private myCoolIntrinsic3 (value : float) =
+    value + 1.0
+
+let private myNotSoCoolIntrinsic (_value : float) =
+    raise <| System.ArgumentException()
+
+[<Test>]
+let ``[Constant] respect onlyInShaderCode``() =
+    let input    = <@ keep (myCoolIntrinsic 1.0) @>
+    input |> Opt.run |> should exprEqual input
+
+[<Test>]
+let ``[Constant] respect FShadeOnlyInShaderCodeException``() =
+    let input    = <@ keep (myCoolIntrinsic2 1.0) @>
+    input |> Opt.run |> should exprEqual input
+
+[<Test>]
+let ``[Constant] evaluate method``() =
+    let input    = <@ keep (myCoolIntrinsic3 1.0) @>
+    let expected = <@ keep 2.0 @>
+    input |> Opt.run |> should exprEqual expected
+
+[<Test>]
+let ``[Constant] evaluate throwing method``() =
+    let input    = <@ keep (myNotSoCoolIntrinsic 1.0) @>
+    input |> Opt.run |> should exprEqual input
\ No newline at end of file