Refactor animate component to reduce complexity

9 files changed, 101 insertions, 88 deletions

diff --git a/desktop/packages/mullvad-vpn/src/renderer/lib/components/animate/Animate.tsx b/desktop/packages/mullvad-vpn/src/renderer/lib/components/animate/Animate.tsx
index b393d88dbe..84c51e7aaa 100644
--- a/desktop/packages/mullvad-vpn/src/renderer/lib/components/animate/Animate.tsx
+++ b/desktop/packages/mullvad-vpn/src/renderer/lib/components/animate/Animate.tsx
@@ -2,14 +2,13 @@ import React from 'react';
 import styled, { css, RuleSet } from 'styled-components';
 
 import { TransientProps } from '../../types';
-import { useMounted } from '../../utility-hooks';
-import { AnimateProvider } from './AnimateContext';
-import { useAnimations, useHandleAnimationEnd, useShow } from './hooks';
+import { AnimateProvider, useAnimateContext } from './AnimateContext';
+import { useAnimate, useAnimations, useHandleAnimationEnd } from './hooks';
 import { Animation } from './types';
 
 type AnimateBaseProps = {
   initial?: boolean;
-  present?: boolean;
+  present: boolean;
   duration?: React.CSSProperties['animationDuration'];
   timingFunction?: React.CSSProperties['animationTimingFunction'];
   direction?: React.CSSProperties['animationDirection'];
@@ -23,7 +22,7 @@ export type AnimateProps = AnimateBaseProps &
   Omit<React.HTMLAttributes<HTMLDivElement>, keyof AnimateBaseProps>;
 
 const StyledDiv = styled.div<
-  TransientProps<Omit<AnimateBaseProps, 'animations'>> & {
+  TransientProps<Omit<AnimateBaseProps, 'animations' | 'present'>> & {
     $animations: RuleSet;
   }
 >`
@@ -69,52 +68,34 @@ const StyledDiv = styled.div<
  * @param present - Whether element is present, i.e rendered or not.
  * @param animations - List of animations to apply.
  */
-export function Animate({ initial, present, children, ...props }: AnimateProps) {
+export function Animate({ animations, initial, present, children, ...props }: AnimateProps) {
   return (
-    <AnimateProvider initial={initial} present={present}>
-      <AnimateImpl {...props} present={present} initial={initial}>
-        {children}
-      </AnimateImpl>
+    <AnimateProvider animations={animations} initial={initial} present={present}>
+      <AnimateImpl {...props}>{children}</AnimateImpl>
     </AnimateProvider>
   );
 }
 
+export type AnimateImplProps = Omit<AnimateProps, 'animations' | 'present'>;
+
 function AnimateImpl({
-  initial = true,
-  present: presentProp,
-  animations: animationsProp,
   duration,
   timingFunction,
   direction,
   iterationCount,
   onAnimationEnd,
   ...props
-}: AnimateProps) {
-  const animations = useAnimations(animationsProp);
-  const show = useShow();
-  const [initialPresent] = React.useState(presentProp);
-  const [presentChanged, setPresentChanged] = React.useState(false);
-  const [present, setPresent] = React.useState(presentProp);
-
-  React.useEffect(() => {
-    if (presentProp !== initialPresent && !presentChanged) {
-      setPresentChanged(true);
-    }
-  }, [initialPresent, presentProp, presentChanged]);
-
-  React.useEffect(() => {
-    setPresent(presentProp);
-  }, [presentProp]);
-
-  const handleAnimationEnd = useHandleAnimationEnd(onAnimationEnd);
-  const mounted = useMounted();
-  if (!show) return null;
+}: AnimateImplProps) {
+  const { animatePresent } = useAnimateContext();
+  const animations = useAnimations();
+  const animate = useAnimate();
+  const handleAnimationEnd = useHandleAnimationEnd();
 
   return (
     <StyledDiv
-      data-animate={initial || (mounted() && presentChanged)}
-      data-present={present}
-      data-is-present-animation={presentProp !== undefined}
+      data-animate={animate}
+      data-present={animatePresent}
+      data-is-present-animation={true}
       onAnimationEnd={handleAnimationEnd}
       $animations={animations}
       $duration={duration}
diff --git a/desktop/packages/mullvad-vpn/src/renderer/lib/components/animate/AnimateContext.tsx b/desktop/packages/mullvad-vpn/src/renderer/lib/components/animate/AnimateContext.tsx
index e9a41e8079..d168f07a2e 100644
--- a/desktop/packages/mullvad-vpn/src/renderer/lib/components/animate/AnimateContext.tsx
+++ b/desktop/packages/mullvad-vpn/src/renderer/lib/components/animate/AnimateContext.tsx
@@ -1,10 +1,15 @@
-import React from 'react';
+import React, { useState } from 'react';
+
+import { Animation } from './types';
 
 interface AnimateContextProps {
-  present?: boolean;
+  animations: Animation[];
+  animate: boolean;
+  animatePresent: boolean;
+  present: boolean;
   initial?: boolean;
-  show: boolean;
-  setShow: React.Dispatch<React.SetStateAction<boolean>>;
+  setAnimate: React.Dispatch<React.SetStateAction<boolean>>;
+  setAnimatePresent: React.Dispatch<React.SetStateAction<boolean>>;
 }
 
 const AnimateContext = React.createContext<AnimateContextProps | undefined>(undefined);
@@ -18,15 +23,32 @@ export const useAnimateContext = (): AnimateContextProps => {
 };
 
 interface AnimateProviderProps {
-  present?: boolean;
+  animations: Animation[];
+  present: boolean;
   initial?: boolean;
   children: React.ReactNode;
 }
 
-export const AnimateProvider = ({ present, initial, children }: AnimateProviderProps) => {
-  const [show, setShow] = React.useState<boolean>(present || false);
+export const AnimateProvider = ({
+  animations,
+  present,
+  initial,
+  children,
+}: AnimateProviderProps) => {
+  const [animate, setAnimate] = React.useState<boolean>((initial && present) || false);
+  const [animatePresent, setAnimatePresent] = useState<boolean>(present);
+
   return (
-    <AnimateContext.Provider value={{ present, initial, show, setShow }}>
+    <AnimateContext.Provider
+      value={{
+        animate,
+        animatePresent,
+        animations,
+        initial,
+        present,
+        setAnimate,
+        setAnimatePresent,
+      }}>
       {children}
     </AnimateContext.Provider>
   );
diff --git a/desktop/packages/mullvad-vpn/src/renderer/lib/components/animate/hooks/index.ts b/desktop/packages/mullvad-vpn/src/renderer/lib/components/animate/hooks/index.ts
index 5ba626bda4..8d6128d192 100644
--- a/desktop/packages/mullvad-vpn/src/renderer/lib/components/animate/hooks/index.ts
+++ b/desktop/packages/mullvad-vpn/src/renderer/lib/components/animate/hooks/index.ts
@@ -1,4 +1,4 @@
+export * from './useAnimate';
 export * from './useAnimations';
-export * from './useIsInitialRender';
-export * from './useShow';
 export * from './useHandleAnimationEnd';
+export * from './usePreviousValue';
diff --git a/desktop/packages/mullvad-vpn/src/renderer/lib/components/animate/hooks/useAnimate.ts b/desktop/packages/mullvad-vpn/src/renderer/lib/components/animate/hooks/useAnimate.ts
new file mode 100644
index 0000000000..f67405c72d
--- /dev/null
+++ b/desktop/packages/mullvad-vpn/src/renderer/lib/components/animate/hooks/useAnimate.ts
@@ -0,0 +1,18 @@
+import { useEffect } from 'react';
+
+import { useAnimateContext } from '../AnimateContext';
+import { usePreviousValue } from './usePreviousValue';
+
+export const useAnimate = () => {
+  const { animate, present, setAnimate, setAnimatePresent } = useAnimateContext();
+  const previousPresent = usePreviousValue(present);
+
+  useEffect(() => {
+    if (present !== previousPresent) {
+      setAnimate(true);
+      setAnimatePresent(present);
+    }
+  }, [present, previousPresent, setAnimate, setAnimatePresent]);
+
+  return animate;
+};
diff --git a/desktop/packages/mullvad-vpn/src/renderer/lib/components/animate/hooks/useAnimations.ts b/desktop/packages/mullvad-vpn/src/renderer/lib/components/animate/hooks/useAnimations.ts
index b9a887182d..232bcd99b5 100644
--- a/desktop/packages/mullvad-vpn/src/renderer/lib/components/animate/hooks/useAnimations.ts
+++ b/desktop/packages/mullvad-vpn/src/renderer/lib/components/animate/hooks/useAnimations.ts
@@ -1,10 +1,12 @@
 import { css, RuleSet } from 'styled-components';
 
+import { useAnimateContext } from '../AnimateContext';
 import { animations } from '../animations';
-import { Animation } from '../types';
 import { createAnimationDeclaration } from '../utils';
 
-export const useAnimations = (values: Animation[]) => {
+export const useAnimations = () => {
+  const { animations: values } = useAnimateContext();
+
   const inAnimations: Array<{ name: string; rule: RuleSet }> = [];
   const outAnimations: Array<{ name: string; rule: RuleSet }> = [];
 
diff --git a/desktop/packages/mullvad-vpn/src/renderer/lib/components/animate/hooks/useHandleAnimationEnd.ts b/desktop/packages/mullvad-vpn/src/renderer/lib/components/animate/hooks/useHandleAnimationEnd.ts
index 2b4d83441c..53389437bc 100644
--- a/desktop/packages/mullvad-vpn/src/renderer/lib/components/animate/hooks/useHandleAnimationEnd.ts
+++ b/desktop/packages/mullvad-vpn/src/renderer/lib/components/animate/hooks/useHandleAnimationEnd.ts
@@ -1,19 +1,23 @@
-import React from 'react';
+import { useCallback, useRef } from 'react';
 
-import { AnimateProps } from '../Animate';
 import { useAnimateContext } from '../AnimateContext';
 
-export const useHandleAnimationEnd = (onAnimationEnd: AnimateProps['onAnimationEnd']) => {
-  const { present, show, setShow } = useAnimateContext();
-  return React.useCallback(
-    (e: React.AnimationEvent<HTMLDivElement>) => {
-      if (!present && show) {
-        setShow(false);
-      }
-      if (onAnimationEnd) {
-        onAnimationEnd(e);
-      }
-    },
-    [onAnimationEnd, present, setShow, show],
-  );
+export const useHandleAnimationEnd = () => {
+  const { animations, present, setAnimate, setAnimatePresent } = useAnimateContext();
+  const animationsCount = animations.length;
+  const animationsFinishedCount = useRef(0);
+
+  const handleAnimationEnd = useCallback(() => {
+    const nextAnimationsFinishedCount = animationsFinishedCount.current + 1;
+
+    if (nextAnimationsFinishedCount === animationsCount) {
+      animationsFinishedCount.current = 0;
+      setAnimate(false);
+      setAnimatePresent(present);
+    } else {
+      animationsFinishedCount.current = nextAnimationsFinishedCount;
+    }
+  }, [animationsCount, animationsFinishedCount, present, setAnimate, setAnimatePresent]);
+
+  return handleAnimationEnd;
 };
diff --git a/desktop/packages/mullvad-vpn/src/renderer/lib/components/animate/hooks/useIsInitialRender.ts b/desktop/packages/mullvad-vpn/src/renderer/lib/components/animate/hooks/useIsInitialRender.ts
deleted file mode 100644
index 546eba8152..0000000000
--- a/desktop/packages/mullvad-vpn/src/renderer/lib/components/animate/hooks/useIsInitialRender.ts
+++ /dev/null
@@ -1,11 +0,0 @@
-import React from 'react';
-
-export const useIsInitialRender = () => {
-  const isInitialRender = React.useRef(true);
-
-  React.useEffect(() => {
-    isInitialRender.current = false;
-  }, []);
-
-  return React.useCallback(() => isInitialRender.current, [isInitialRender]);
-};
diff --git a/desktop/packages/mullvad-vpn/src/renderer/lib/components/animate/hooks/usePreviousValue.ts b/desktop/packages/mullvad-vpn/src/renderer/lib/components/animate/hooks/usePreviousValue.ts
new file mode 100644
index 0000000000..1b835a7823
--- /dev/null
+++ b/desktop/packages/mullvad-vpn/src/renderer/lib/components/animate/hooks/usePreviousValue.ts
@@ -0,0 +1,11 @@
+import { useEffect, useState } from 'react';
+
+export const usePreviousValue = <T>(value: T) => {
+  const [previousValue, setPreviousValue] = useState(value);
+
+  useEffect(() => {
+    setPreviousValue(value);
+  }, [setPreviousValue, value]);
+
+  return previousValue;
+};
diff --git a/desktop/packages/mullvad-vpn/src/renderer/lib/components/animate/hooks/useShow.ts b/desktop/packages/mullvad-vpn/src/renderer/lib/components/animate/hooks/useShow.ts
deleted file mode 100644
index 2a452ffbb4..0000000000
--- a/desktop/packages/mullvad-vpn/src/renderer/lib/components/animate/hooks/useShow.ts
+++ /dev/null
@@ -1,14 +0,0 @@
-import React from 'react';
-
-import { useAnimateContext } from '../AnimateContext';
-
-export const useShow = () => {
-  const { present, show, setShow } = useAnimateContext();
-
-  React.useEffect(() => {
-    if (present && !show) {
-      setShow(true);
-    }
-  }, [present, setShow, show]);
-  return show;
-};