From 8e181d6ae2dff82875c60be789dcd63bedaba883 Mon Sep 17 00:00:00 2001
From: Hernan Torrisi <hernantorrisi@gmail.com>
Date: Mon, 13 May 2024 16:16:40 -0700
Subject: [PATCH] change hooks lifecycle to account for component reloading

---
 src/hooks/useResizeCanvas.ts | 28 +++++++++------
 src/hooks/useRive.tsx        | 66 ++++++++++++++++++++----------------
 2 files changed, 55 insertions(+), 39 deletions(-)

diff --git a/src/hooks/useResizeCanvas.ts b/src/hooks/useResizeCanvas.ts
index 1f6bd06..1151577 100644
--- a/src/hooks/useResizeCanvas.ts
+++ b/src/hooks/useResizeCanvas.ts
@@ -13,7 +13,7 @@ interface UseResizeCanvasProps {
   /**
    * Ref to the canvas element
    */
-  canvasRef: MutableRefObject<HTMLCanvasElement | null>;
+  canvasElem: HTMLCanvasElement | null;
   /**
    * Ref to the container element of the canvas
    */
@@ -55,7 +55,7 @@ interface UseResizeCanvasProps {
  */
 export default function useResizeCanvas({
   riveLoaded = false,
-  canvasRef,
+  canvasElem,
   containerRef,
   options = {},
   onCanvasHasResized,
@@ -120,7 +120,7 @@ export default function useResizeCanvas({
 
     const { width, height } = getContainerDimensions();
     let hasResized = false;
-    if (canvasRef.current) {
+    if (canvasElem) {
       // Check if the canvas parent container bounds have changed and set
       // new values accordingly
       const boundsChanged =
@@ -138,10 +138,10 @@ export default function useResizeCanvas({
         if (boundsChanged || canvasSizeChanged) {
           const newCanvasWidthProp = currentDevicePixelRatio * width;
           const newCanvasHeightProp = currentDevicePixelRatio * height;
-          canvasRef.current.width = newCanvasWidthProp;
-          canvasRef.current.height = newCanvasHeightProp;
-          canvasRef.current.style.width = width + 'px';
-          canvasRef.current.style.height = height + 'px';
+          canvasElem.width = newCanvasWidthProp;
+          canvasElem.height = newCanvasHeightProp;
+          canvasElem.style.width = width + 'px';
+          canvasElem.style.height = height + 'px';
           setLastCanvasSize({
             width: newCanvasWidthProp,
             height: newCanvasHeightProp,
@@ -149,8 +149,8 @@ export default function useResizeCanvas({
           hasResized = true;
         }
       } else if (boundsChanged) {
-        canvasRef.current.width = width;
-        canvasRef.current.height = height;
+        canvasElem.width = width;
+        canvasElem.height = height;
         setLastCanvasSize({
           width: width,
           height: height,
@@ -167,7 +167,7 @@ export default function useResizeCanvas({
     }
     isFirstSizing && setIsFirstSizing(false);
   }, [
-    canvasRef,
+    canvasElem,
     containerRef,
     containerSize,
     currentDevicePixelRatio,
@@ -184,4 +184,12 @@ export default function useResizeCanvas({
     shouldUseDevicePixelRatio,
     riveLoaded,
   ]);
+
+  // Reset width and height values when the canvas changes
+  useEffect(() => {
+    setLastCanvasSize({
+      width: 0,
+      height: 0,
+    });
+  }, [canvasElem]);
 }
diff --git a/src/hooks/useRive.tsx b/src/hooks/useRive.tsx
index 3fed470..354d977 100644
--- a/src/hooks/useRive.tsx
+++ b/src/hooks/useRive.tsx
@@ -65,7 +65,7 @@ export default function useRive(
   riveParams?: UseRiveParameters,
   opts: Partial<UseRiveOptions> = {}
 ): RiveState {
-  const canvasRef = useRef<HTMLCanvasElement | null>(null);
+  const [canvasElem, setCanvasElem] = useState<HTMLCanvasElement | null>(null);
   const containerRef = useRef<HTMLElement | null>(null);
 
   const [rive, setRive] = useState<Rive | null>(null);
@@ -87,7 +87,7 @@ export default function useRive(
   // Watch the canvas parent container resize and size the canvas to match
   useResizeCanvas({
     riveLoaded: !!rive,
-    canvasRef,
+    canvasElem,
     containerRef,
     options,
     onCanvasHasResized,
@@ -99,32 +99,40 @@ export default function useRive(
    */
   const setCanvasRef: RefCallback<HTMLCanvasElement> = useCallback(
     (canvas: HTMLCanvasElement | null) => {
-      if (canvas && riveParams && isParamsLoaded) {
-        const { useOffscreenRenderer } = options;
-        const r = new Rive({
-          useOffscreenRenderer,
-          ...riveParams,
-          canvas,
-        });
-        r.on(EventType.Load, () => {
-          // Check if the component/canvas is mounted before setting state to avoid setState
-          // on an unmounted component in some rare cases
-          if (canvasRef.current) {
-            setRive(r);
-          } else {
-            // If unmounted, cleanup the rive object immediately
-            r.cleanup();
-          }
-        });
-      } else if (canvas === null && canvasRef.current) {
-        canvasRef.current.height = 0;
-        canvasRef.current.width = 0;
+
+      if (canvas === null && canvasElem) {
+        canvasElem.height = 0;
+        canvasElem.width = 0;
       }
 
-      canvasRef.current = canvas;
+      setCanvasElem(canvas);
     },
-    [isParamsLoaded]
+    []
   );
+
+  useEffect(() => {
+    if(!canvasElem || !riveParams) {
+      return;
+    }
+    if (rive == null) {
+      const { useOffscreenRenderer } = options;
+      const r = new Rive({
+        useOffscreenRenderer,
+        ...riveParams,
+        canvas: canvasElem,
+      });
+      r.on(EventType.Load, () => {
+        // Check if the component/canvas is mounted before setting state to avoid setState
+        // on an unmounted component in some rare cases
+        if (canvasElem) {
+          setRive(r);
+        } else {
+          // If unmounted, cleanup the rive object immediately
+          r.cleanup();
+        }
+      });
+    }
+  },[canvasElem, isParamsLoaded, rive]);
   /**
    * Ref callback called when the container element mounts
    */
@@ -146,14 +154,14 @@ export default function useRive(
         : rive && rive.stopRendering();
     });
 
-    if (canvasRef.current) {
-      observer.observe(canvasRef.current);
+    if (canvasElem) {
+      observer.observe(canvasElem);
     }
 
     return () => {
       observer.disconnect();
     };
-  }, [rive]);
+  }, [rive, canvasElem]);
 
   /**
    * On unmount, call cleanup to cleanup any WASM generated objects that need
@@ -166,7 +174,7 @@ export default function useRive(
         setRive(null);
       }
     };
-  }, [rive]);
+  }, [rive, canvasElem]);
 
   /**
    * Listen for changes in the animations params
@@ -198,7 +206,7 @@ export default function useRive(
   );
 
   return {
-    canvas: canvasRef.current,
+    canvas: canvasElem,
     container: containerRef.current,
     setCanvasRef,
     setContainerRef,