From 76bede61fca7dbbbe7d15696e9bcad5251406266 Mon Sep 17 00:00:00 2001
From: Rohan <45748283+r100-stack@users.noreply.github.com>
Date: Tue, 21 Jan 2025 15:48:49 -0500
Subject: [PATCH] Cleanup

---
 .../itwinui-react/src/core/Toast/Toast.tsx    | 47 +++++++++++--------
 1 file changed, 27 insertions(+), 20 deletions(-)

diff --git a/packages/itwinui-react/src/core/Toast/Toast.tsx b/packages/itwinui-react/src/core/Toast/Toast.tsx
index 714d511d7af..e6a5d248b10 100644
--- a/packages/itwinui-react/src/core/Toast/Toast.tsx
+++ b/packages/itwinui-react/src/core/Toast/Toast.tsx
@@ -288,7 +288,6 @@ const useAnimateToastBasedOnVisibility = (
   const motionOk = useMediaQuery('(prefers-reduced-motion: no-preference)');
   const animateOutToRef = useLatestRef(animateOutTo);
   const onRemoveRef = useLatestRef(onRemove);
-  const setShouldBeMountedRef = useLatestRef(setShouldBeMounted);
 
   const [prevIsVisible, setPrevIsVisible] = React.useState<
     typeof isVisible | undefined
@@ -300,25 +299,9 @@ const useAnimateToastBasedOnVisibility = (
       setPrevIsVisible(isVisible);
 
       if (isVisible) {
-        setShouldBeMountedRef.current(true);
-
-        // Mount *before* handling dialog entry.
-        queueMicrotask(() => {
-          animateIn();
-        });
+        safeAnimateIn();
       } else {
-        if (!motionOk) {
-          setShouldBeMountedRef.current(false);
-          onRemoveRef.current?.();
-        } else {
-          const animation = animateOut();
-
-          // Unmount *after* handling dialog exit.
-          animation?.addEventListener('finish', () => {
-            setShouldBeMountedRef.current(false);
-            onRemoveRef.current?.();
-          });
-        }
+        safeAnimateOut();
       }
     }
 
@@ -335,6 +318,30 @@ const useAnimateToastBasedOnVisibility = (
       return { translateX, translateY };
     }
 
+    function safeAnimateIn() {
+      setShouldBeMounted(true);
+
+      // Mount *before* handling dialog entry.
+      queueMicrotask(() => {
+        animateIn();
+      });
+    }
+
+    function safeAnimateOut() {
+      if (!motionOk) {
+        setShouldBeMounted(false);
+        onRemoveRef.current?.();
+      } else {
+        const animation = animateOut();
+
+        // Unmount *after* handling dialog exit.
+        animation?.addEventListener('finish', () => {
+          setShouldBeMounted(false);
+          onRemoveRef.current?.();
+        });
+      }
+    }
+
     function animateIn() {
       if (!motionOk) {
         return;
@@ -385,11 +392,11 @@ const useAnimateToastBasedOnVisibility = (
     animateOutTo,
     motionOk,
     thisElement,
+    setShouldBeMounted,
 
     // Using latest refs to avoid unnecessary effect executions.
     animateOutToRef,
     onRemoveRef,
-    setShouldBeMountedRef,
   ]);
 
   return shouldBeMounted;