From 013a45fa5a7e5e87eb8ffc86130aa8b7747bba84 Mon Sep 17 00:00:00 2001
From: Alfonso Garcia-Caro <alfonso.garcia-caro@outlook.com>
Date: Sat, 6 Nov 2021 18:13:11 +0900
Subject: [PATCH] Make UseElmish compatible with Fast Refresh

---
 Feliz.UseElmish/UseElmish.fs | 69 +++++++++++++++++++++++++-----------
 1 file changed, 49 insertions(+), 20 deletions(-)

diff --git a/Feliz.UseElmish/UseElmish.fs b/Feliz.UseElmish/UseElmish.fs
index e0c569e1..cb988f11 100644
--- a/Feliz.UseElmish/UseElmish.fs
+++ b/Feliz.UseElmish/UseElmish.fs
@@ -2,10 +2,11 @@ namespace Feliz.UseElmish
 
 open Feliz
 open Elmish
+open Fable.React
 
 [<AutoOpen>]
 module UseElmishExtensions =
-    type private ElmishObservable<'Model, 'Msg>() =
+    type ElmishObservable<'Model, 'Msg>() =
         let mutable hasDisposedOnce = false
         let mutable state: 'Model option = None
         let mutable listener: ('Model -> unit) option = None
@@ -39,11 +40,11 @@ module UseElmishExtensions =
                 | :? System.IDisposable as disp -> disp.Dispose()
                 | _ -> ()
             | _ -> ()
-            dispatcher <- None
-            state <- None
+            // dispatcher <- None
+            // state <- None
             hasDisposedOnce <- true
 
-    let private runProgram (program: unit -> Program<'Arg, 'Model, 'Msg, unit>) (arg: 'Arg) (obs: ElmishObservable<'Model, 'Msg>) () =
+    let runProgram (program: unit -> Program<'Arg, 'Model, 'Msg, unit>) (arg: 'Arg) (obs: ElmishObservable<'Model, 'Msg>) () =
         program()
         |> Program.withSetState obs.SetState
         |> Program.runWith arg
@@ -58,34 +59,62 @@ module UseElmishExtensions =
         | _ -> ()
 
     type React with
-        [<Hook>]
-        static member useElmish(program: unit -> Program<'Arg, 'Model, 'Msg, unit>, arg: 'Arg, ?dependencies: obj array) =
+        static member
+#if DEBUG
+            inline
+#endif
+            useElmish(program: unit -> Program<'Arg, 'Model, 'Msg, unit>, arg: 'Arg, ?dependencies: obj array) =
+
             // Don't use useMemo here because React doesn't guarantee it won't recreate it again
-            let obs, _ = React.useState(fun () -> ElmishObservable<'Model, 'Msg>())
+            let obs = Hooks.useStateLazy(fun () -> ElmishObservable<'Model, 'Msg>())
+            let obs = obs.current
 
-            let state, setState = React.useState(runProgram program arg obs)
+            let state = Hooks.useStateLazy(runProgram program arg obs)
 
-            React.useEffect((fun () ->
+            Hooks.useEffectDisposable((fun () ->
                 if obs.HasDisposedOnce then
-                    runProgram program arg obs () |> setState
+                    // In Debug mode, a component may be disconnected and reconnected during HMR
+                    // TODO: We need a way to tell if the effect has been triggered because of a change in dependencies,
+                    // in that case state does need to be reset.
+                    let newValue =
+#if DEBUG
+                        let curValue = obs.Value
+                        let newValue = runProgram program arg obs ()
+                        defaultArg curValue newValue
+#else
+                        runProgram program arg obs ()
+#endif
+                    state.update(fun _ -> newValue)
                 React.createDisposable(obs.DisposeState)
             ), defaultArg dependencies [||])
 
-            obs.Subscribe(setState)
-            state, obs.Dispatch
+            obs.Subscribe(fun v -> state.update(fun _ -> v))
+            state.current, obs.Dispatch
 
-        [<Hook>]
-        static member useElmish(program: unit -> Program<unit, 'Model, 'Msg, unit>, ?dependencies: obj array) =
+        static member 
+#if DEBUG
+            inline
+#endif        
+            useElmish(program: unit -> Program<unit, 'Model, 'Msg, unit>, ?dependencies: obj array) =
             React.useElmish(program, (), ?dependencies=dependencies)
 
-        [<Hook>]
-        static member useElmish(init: 'Arg -> 'Model * Cmd<'Msg>, update: 'Msg -> 'Model -> 'Model * Cmd<'Msg>, arg: 'Arg, ?dependencies: obj array) =
+        static member 
+#if DEBUG
+            inline
+#endif        
+            useElmish(init: 'Arg -> 'Model * Cmd<'Msg>, update: 'Msg -> 'Model -> 'Model * Cmd<'Msg>, arg: 'Arg, ?dependencies: obj array) =
             React.useElmish((fun () -> Program.mkProgram init update (fun _ _ -> ())), arg, ?dependencies=dependencies)
 
-        [<Hook>]
-        static member useElmish(init: unit -> 'Model * Cmd<'Msg>, update: 'Msg -> 'Model -> 'Model * Cmd<'Msg>, ?dependencies: obj array) =
+        static member 
+#if DEBUG
+            inline
+#endif        
+            useElmish(init: unit -> 'Model * Cmd<'Msg>, update: 'Msg -> 'Model -> 'Model * Cmd<'Msg>, ?dependencies: obj array) =
             React.useElmish((fun () -> Program.mkProgram init update (fun _ _ -> ())), ?dependencies=dependencies)
 
-        [<Hook>]
-        static member useElmish(init: 'Model * Cmd<'Msg>, update: 'Msg -> 'Model -> 'Model * Cmd<'Msg>, ?dependencies: obj array) =
+        static member 
+#if DEBUG
+            inline
+#endif        
+            useElmish(init: 'Model * Cmd<'Msg>, update: 'Msg -> 'Model -> 'Model * Cmd<'Msg>, ?dependencies: obj array) =
             React.useElmish((fun () -> Program.mkProgram (fun () -> init) update (fun _ _ -> ())), ?dependencies=dependencies)