From 3e48f3fbf9af188e583c1cf63669d4f58f042019 Mon Sep 17 00:00:00 2001
From: Alexander Zinchuk <alx.zinchuk@gmail.com>
Date: Wed, 11 Sep 2024 02:07:15 +0200
Subject: [PATCH] `useShowTransitions`: Workaround for `noMountTransition` not
 working after forced reflow

---
 src/hooks/useShowTransition.ts | 29 ++++++++++++++++++++---------
 src/lib/teact/teact-dom.ts     |  1 +
 src/lib/teact/teact.ts         | 16 ++++++++++------
 3 files changed, 31 insertions(+), 15 deletions(-)

diff --git a/src/hooks/useShowTransition.ts b/src/hooks/useShowTransition.ts
index c68b77323..97c579999 100644
--- a/src/hooks/useShowTransition.ts
+++ b/src/hooks/useShowTransition.ts
@@ -1,4 +1,4 @@
-import type { RefObject } from 'react';
+import type { RefObject } from '../lib/teact/teact';
 import { useLayoutEffect, useRef, useSignal } from '../lib/teact/teact';
 import { addExtraClass, toggleExtraClass } from '../lib/teact/teact-dom';
 
@@ -9,13 +9,14 @@ import useDerivedSignal from './useDerivedSignal';
 import useDerivedState from './useDerivedState';
 import useLastCallback from './useLastCallback';
 import { useStateRef } from './useStateRef';
+import useSyncEffect from './useSyncEffect';
 import useSyncEffectWithPrevDeps from './useSyncEffectWithPrevDeps';
 
 const CLOSE_DURATION = 350;
 
 type BaseHookParams<RefType extends HTMLElement> = {
   isOpen: boolean | undefined;
-  ref?: RefObject<RefType>;
+  ref?: RefObject<RefType | null>;
   noMountTransition?: boolean;
   noOpenTransition?: boolean;
   noCloseTransition?: boolean;
@@ -34,7 +35,7 @@ type HookParamsWithShouldRender<RefType extends HTMLElement> = BaseHookParams<Re
 };
 
 type HookResult<RefType extends HTMLElement> = {
-  ref: RefObject<RefType>;
+  ref: RefObject<RefType | null>;
   getIsClosing: Signal<boolean>;
 };
 
@@ -67,13 +68,10 @@ export default function useShowTransition<RefType extends HTMLElement = HTMLDivE
     prefix = '',
     onCloseAnimationEnd,
   } = params;
-  let ref = params.ref;
-
-  const withShouldRender = 'withShouldRender' in params && params.withShouldRender;
 
   // eslint-disable-next-line no-null/no-null
   const localRef = useRef<RefType>(null);
-  ref ||= localRef;
+  const ref = params.ref || localRef;
   const closingTimeoutRef = useRef<number>();
   const [getState, setState] = useSignal<State | undefined>();
   const optionsRef = useStateRef({
@@ -110,7 +108,7 @@ export default function useShowTransition<RefType extends HTMLElement = HTMLDivE
     }
   }, [isOpen]);
 
-  useLayoutEffect(() => {
+  const applyClassNames = useLastCallback(() => {
     const element = ref.current;
     if (!element) return;
 
@@ -129,12 +127,25 @@ export default function useShowTransition<RefType extends HTMLElement = HTMLDivE
     toggleExtraClass(element, `${prefix}open`, hasOpenClass);
     toggleExtraClass(element, `${prefix}not-open`, !hasOpenClass);
     toggleExtraClass(element, `${prefix}closing`, isClosing);
-  }, [className, getState, prefix, ref]);
+  });
 
+  // Workaround for Chrome causing forced reflow in the middle of mutation phase when unmounting a focused element.
+  // Due to such forced reflow setting initial class names in the first layout effect causes transitions to start.
+  useSyncEffect(() => {
+    ref.onChange = () => {
+      ref.onChange = undefined;
+      applyClassNames();
+    };
+  }, [applyClassNames, ref]);
+
+  useLayoutEffect(applyClassNames, [applyClassNames, getState]);
+
+  const withShouldRender = 'withShouldRender' in params && params.withShouldRender;
   const shouldRender = useDerivedState(
     () => (withShouldRender && getState() !== 'closed'),
     [withShouldRender, getState],
   );
+
   const getIsClosing = useDerivedSignal(() => getState() === 'closing', [getState]);
 
   if (withShouldRender) {
diff --git a/src/lib/teact/teact-dom.ts b/src/lib/teact/teact-dom.ts
index 5aaf78891..cf5a064a9 100644
--- a/src/lib/teact/teact-dom.ts
+++ b/src/lib/teact/teact-dom.ts
@@ -628,6 +628,7 @@ function setElementRef($element: VirtualElementTag, element: DOMElement | undefi
 
   if (typeof ref === 'object') {
     ref.current = element;
+    ref.onChange?.();
   } else if (typeof ref === 'function') {
     ref(element);
   }
diff --git a/src/lib/teact/teact.ts b/src/lib/teact/teact.ts
index bc03afe53..df9fa3f0f 100644
--- a/src/lib/teact/teact.ts
+++ b/src/lib/teact/teact.ts
@@ -57,6 +57,11 @@ export interface VirtualElementFragment {
 
 export type StateHookSetter<T> = (newValue: ((current: T) => T) | T) => void;
 
+export interface RefObject<T = any> {
+  current: T;
+  onChange?: NoneToVoidFunction;
+}
+
 export enum MountState {
   New,
   Mounted,
@@ -99,9 +104,7 @@ interface ComponentInstance {
     };
     refs?: {
       cursor: number;
-      byCursor: {
-        current: any;
-      }[];
+      byCursor: RefObject[];
     };
   };
   prepareForFrame?: () => void;
@@ -547,6 +550,7 @@ function helpGc(componentInstance: ComponentInstance) {
   if (refs) {
     for (const hook of refs.byCursor) {
       hook.current = undefined as any;
+      hook.onChange = undefined as any;
     }
   }
 
@@ -886,9 +890,9 @@ export function useCallback<F extends AnyFunction>(newCallback: F, dependencies:
   return useMemo(() => newCallback, dependencies, debugKey);
 }
 
-export function useRef<T>(initial: T): { current: T };
-export function useRef<T>(): { current: T | undefined }; // TT way (empty is `undefined`)
-export function useRef<T>(initial: null): { current: T | null }; // React way (empty is `null`)
+export function useRef<T>(initial: T): RefObject<T>;
+export function useRef<T>(): RefObject<T | undefined>; // TT way (empty is `undefined`)
+export function useRef<T>(initial: null): RefObject<T | null>; // React way (empty is `null`)
 // eslint-disable-next-line no-null/no-null
 export function useRef<T>(initial?: T | null) {
   if (!renderingInstance.hooks) {