More elaborate

Author: alexqhj

module/source/hooks/use-unity-loader.ts

@@ -3,11 +3,15 @@ import { isBrowserEnvironment } from "../constants/is-browser-environment";
 import { UnityLoaderStatus } from "../enums/unity-loader-status";
 import { UnityConfig } from "../exports";
 
-// Map to track script references and their instance count
-const scriptReferenceMap = new Map<
-  string,
-  { count: number; status: UnityLoaderStatus }
->();
+interface ScriptData {
+  count: number;
+  status: UnityLoaderStatus;
+  handleLoad?: () => void;
+  handleError?: () => void;
+}
+
+// Map to track script references, status, and event handlers.
+const scriptReferenceMap = new Map<string, ScriptData>();
 
 /**
  * Hook to embed a Unity Loader script.
@@ -25,86 +29,111 @@ const useUnityLoader = (unityConfig: UnityConfig): UnityLoaderStatus => {
     // It is possible for the application being rendered server side. In
     // this scenario, the window is not available. We can't create a Unity
     // Loader in this case.
-    if (isBrowserEnvironment === false) {
+    if (!isBrowserEnvironment) {
       return undefined;
     }
+
     // If the script's source is null, we'll reset the status to idle.
-    if (unityConfig.loaderUrl === null) {
+    if (!loaderUrl) {
       setStatus(UnityLoaderStatus.Idle);
       return undefined;
     }
 
-    /**
-     * Find existing script element by source. It may have been added by
-     * another instance of this hook.
-     */
+    // Find existing script element by source.
+    // It may have been added by another instance of this hook.
     let script = document.querySelector(
       `script[src="${loaderUrl}"]`
     ) as HTMLScriptElement | null;
 
-    // If script exists, increase reference count, else we'll create a new script.
+    // Get existing data from the reference map.
+    const existingData = scriptReferenceMap.get(loaderUrl);
+
     if (script) {
-      const refData = scriptReferenceMap.get(loaderUrl) || {
-        count: 0,
-        status: UnityLoaderStatus.Loading,
-      };
-      scriptReferenceMap.set(loaderUrl, {
-        count: refData.count + 1,
-        status: refData.status,
-      });
-      setStatus(refData.status);
+      // If the script is already in the DOM, increment the ref count.
+      if (existingData) {
+        existingData.count += 1;
+        scriptReferenceMap.set(loaderUrl, existingData);
+        setStatus(existingData.status);
+      } else {
+        // Edge case: script is in the DOM, but not in the map.
+        // (Unlikely if all usage is through this hook)
+        scriptReferenceMap.set(loaderUrl, {
+          count: 1,
+          status: UnityLoaderStatus.Loaded,
+        });
+        setStatus(UnityLoaderStatus.Loaded);
+      }
     } else {
+      // Create a new script element.
       script = document.createElement("script");
       script.type = "text/javascript";
       script.src = loaderUrl;
       script.async = true;
       script.setAttribute("data-status", "loading");
       document.body.appendChild(script);
 
-      // Initialize reference map entry
-      scriptReferenceMap.set(loaderUrl, {
-        count: 1,
-        status: UnityLoaderStatus.Loading,
-      });
-
-      script.addEventListener("load", () => {
+      // Define load handler.
+      const handleLoad = () => {
         script?.setAttribute("data-status", "loaded");
-        scriptReferenceMap.set(loaderUrl, {
-          count: 1,
-          status: UnityLoaderStatus.Loaded,
-        });
+        const refData = scriptReferenceMap.get(loaderUrl);
+        if (refData) {
+          refData.status = UnityLoaderStatus.Loaded;
+          scriptReferenceMap.set(loaderUrl, refData);
+        }
         setStatus(UnityLoaderStatus.Loaded);
-      });
+      };
 
-      script.addEventListener("error", () => {
+      // Define error handler.
+      const handleError = () => {
         script?.setAttribute("data-status", "error");
-        scriptReferenceMap.set(loaderUrl, {
-          count: 1,
-          status: UnityLoaderStatus.Error,
-        });
+        const refData = scriptReferenceMap.get(loaderUrl);
+        if (refData) {
+          refData.status = UnityLoaderStatus.Error;
+          scriptReferenceMap.set(loaderUrl, refData);
+        }
         setStatus(UnityLoaderStatus.Error);
+      };
+
+      // Attach listeners.
+      script.addEventListener("load", handleLoad);
+      script.addEventListener("error", handleError);
+
+      // Initialize the reference map.
+      scriptReferenceMap.set(loaderUrl, {
+        count: 1,
+        status: UnityLoaderStatus.Loading,
+        handleLoad,
+        handleError,
       });
+
+      setStatus(UnityLoaderStatus.Loading);
     }
 
-    // Remove event listeners on cleanup.
     return () => {
       const refData = scriptReferenceMap.get(loaderUrl);
+      if (!refData) return;
+
+      // Decrement the ref count.
+      refData.count -= 1;
 
-      if (refData) {
-        if (refData.count > 1) {
-          // Decrease reference count when an instance unmounts
-          scriptReferenceMap.set(loaderUrl, {
-            ...refData,
-            count: refData.count - 1,
-          });
-        } else {
-          // Remove script only when the last instance unmounts
-          scriptReferenceMap.delete(loaderUrl);
-          script?.remove();
+      // If there are no more consumers of the script, remove it from the DOM.
+      if (refData.count <= 0) {
+        scriptReferenceMap.delete(loaderUrl);
+
+        // Also remove listeners before removing the script.
+        if (script && refData.handleLoad && refData.handleError) {
+          script.removeEventListener("load", refData.handleLoad);
+          script.removeEventListener("error", refData.handleError);
         }
+
+        script?.remove();
+      } else {
+        // If there's still at least one consumer of the script, update the map
+        // but don't remove the script from the DOM.
+        scriptReferenceMap.set(loaderUrl, refData);
       }
     };
-  }, [unityConfig.loaderUrl]);
+  }, [loaderUrl]);
 
   return status;
 };