From d0cd729b380db67262d27940b9b4122eff296c2c Mon Sep 17 00:00:00 2001
From: Alexander Zinchuk <alx.zinchuk@gmail.com>
Date: Wed, 13 Oct 2021 14:38:45 +0300
Subject: [PATCH] Teact: Fix effect cleanup order

---
 src/components/test/TestCleanupOrder.tsx | 53 ++++++++++++++++++++++++
 src/lib/teact/teact.ts                   | 27 +++++++-----
 src/util/schedulers.ts                   | 35 +++++++++++++---
 3 files changed, 99 insertions(+), 16 deletions(-)
 create mode 100644 src/components/test/TestCleanupOrder.tsx

diff --git a/src/components/test/TestCleanupOrder.tsx b/src/components/test/TestCleanupOrder.tsx
new file mode 100644
index 000000000..8efdbb77d
--- /dev/null
+++ b/src/components/test/TestCleanupOrder.tsx
@@ -0,0 +1,53 @@
+import React, { useState, useEffect, useLayoutEffect } from '../../lib/teact/teact';
+
+const TestCleanupOrder = () => {
+  const [, setRand] = useState(Math.random());
+
+  useEffect(() => {
+    // eslint-disable-next-line no-console
+    console.log('effect 1');
+
+    setTimeout(() => {
+      setRand(Math.random());
+    }, 3000);
+
+    return () => {
+      // eslint-disable-next-line no-console
+      console.log('cleanup 1');
+    };
+  });
+
+  useEffect(() => {
+    // eslint-disable-next-line no-console
+    console.log('effect 2');
+
+    return () => {
+      // eslint-disable-next-line no-console
+      console.log('cleanup 2');
+    };
+  });
+
+  useLayoutEffect(() => {
+    // eslint-disable-next-line no-console
+    console.log('layout effect 1');
+
+    return () => {
+      // eslint-disable-next-line no-console
+      console.log('layout cleanup 1');
+    };
+  });
+
+  useLayoutEffect(() => {
+    // eslint-disable-next-line no-console
+    console.log('layout effect 2');
+
+    return () => {
+      // eslint-disable-next-line no-console
+      console.log('layout cleanup 2');
+    };
+  });
+
+  return <div>Test</div>;
+};
+
+export default TestCleanupOrder;
diff --git a/src/lib/teact/teact.ts b/src/lib/teact/teact.ts
index 1e5bdf084..073665f02 100644
--- a/src/lib/teact/teact.ts
+++ b/src/lib/teact/teact.ts
@@ -1,6 +1,6 @@
 import { DEBUG, DEBUG_MORE } from '../../config';
 import {
-  fastRaf, onTickEnd, throttleWithPrimaryRaf, throttleWithRaf,
+  fastRaf, fastRafPrimary, onTickEnd, onTickEndPrimary, throttleWithPrimaryRaf, throttleWithRaf,
 } from '../../util/schedulers';
 import { flatten, orderBy } from '../../util/iteratees';
 import arePropsShallowEqual, { getUnequalProps } from '../../util/arePropsShallowEqual';
@@ -10,11 +10,9 @@ import { removeAllDelegatedListeners } from './dom-events';
 export type Props = AnyLiteral;
 export type FC<P extends Props = any> = (props: P) => any;
 // eslint-disable-next-line @typescript-eslint/naming-convention
-export type FC_withDebug =
-  FC
-  & {
-    DEBUG_contentComponentName?: string;
-  };
+export type FC_withDebug = FC & {
+  DEBUG_contentComponentName?: string;
+};
 
 export enum VirtualElementTypesEnum {
   Empty,
@@ -511,6 +509,7 @@ export function useState<T>(initial?: T): [T, StateHookSetter<T>] {
 
 function useLayoutEffectBase(
   schedulerFn: typeof onTickEnd | typeof requestAnimationFrame,
+  primarySchedulerFn: typeof onTickEnd | typeof requestAnimationFrame,
   effect: () => Function | void,
   dependencies?: any[],
   debugKey?: string,
@@ -518,7 +517,7 @@ function useLayoutEffectBase(
   const { cursor, byCursor } = renderingInstance.hooks.effects;
   const componentInstance = renderingInstance;
 
-  const exec = () => {
+  function execCleanup() {
     if (!componentInstance.isMounted) {
       return;
     }
@@ -531,9 +530,15 @@ function useLayoutEffectBase(
         handleError(err);
       }
     }
+  }
+
+  function exec() {
+    if (!componentInstance.isMounted) {
+      return;
+    }
 
     byCursor[cursor].cleanup = effect() as Function;
-  };
+  }
 
   if (byCursor[cursor] !== undefined && dependencies && byCursor[cursor].dependencies) {
     if (dependencies.some((dependency, i) => dependency !== byCursor[cursor].dependencies![i])) {
@@ -556,9 +561,11 @@ function useLayoutEffectBase(
         );
       }
 
+      primarySchedulerFn(execCleanup);
       schedulerFn(exec);
     }
   } else {
+    primarySchedulerFn(execCleanup);
     schedulerFn(exec);
   }
 
@@ -572,11 +579,11 @@ function useLayoutEffectBase(
 }
 
 export function useEffect(effect: () => Function | void, dependencies?: any[], debugKey?: string) {
-  return useLayoutEffectBase(fastRaf, effect, dependencies, debugKey);
+  return useLayoutEffectBase(fastRaf, fastRafPrimary, effect, dependencies, debugKey);
 }
 
 export function useLayoutEffect(effect: () => Function | void, dependencies?: any[], debugKey?: string) {
-  return useLayoutEffectBase(onTickEnd, effect, dependencies, debugKey);
+  return useLayoutEffectBase(onTickEnd, onTickEndPrimary, effect, dependencies, debugKey);
 }
 
 export function useMemo<T extends any>(resolver: () => T, dependencies: any[], debugKey?: string): T {
diff --git a/src/util/schedulers.ts b/src/util/schedulers.ts
index 92b536550..6b5a5177f 100644
--- a/src/util/schedulers.ts
+++ b/src/util/schedulers.ts
@@ -74,7 +74,7 @@ export function throttleWithRaf<F extends AnyToVoidFunction>(fn: F) {
 }
 
 export function throttleWithPrimaryRaf<F extends AnyToVoidFunction>(fn: F) {
-  return throttleWith(fastPrimaryRaf, fn);
+  return throttleWith(fastRafPrimary, fn);
 }
 
 export function throttleWithTickEnd<F extends AnyToVoidFunction>(fn: F) {
@@ -104,10 +104,6 @@ export function throttleWith<F extends AnyToVoidFunction>(schedulerFn: Scheduler
   };
 }
 
-export function onTickEnd(cb: NoneToVoidFunction) {
-  Promise.resolve().then(cb);
-}
-
 export function onIdle(cb: NoneToVoidFunction) {
   // eslint-disable-next-line no-restricted-globals
   if (self.requestIdleCallback) {
@@ -156,10 +152,37 @@ export function fastRaf(callback: NoneToVoidFunction, isPrimary = false) {
   }
 }
 
-export function fastPrimaryRaf(callback: NoneToVoidFunction) {
+export function fastRafPrimary(callback: NoneToVoidFunction) {
   fastRaf(callback, true);
 }
 
+let onTickEndCallbacks: NoneToVoidFunction[] | undefined;
+let onTickEndPrimaryCallbacks: NoneToVoidFunction[] | undefined;
+
+export function onTickEnd(callback: NoneToVoidFunction, isPrimary = false) {
+  if (!onTickEndCallbacks) {
+    onTickEndCallbacks = isPrimary ? [] : [callback];
+    onTickEndPrimaryCallbacks = isPrimary ? [callback] : [];
+
+    Promise.resolve().then(() => {
+      const currentCallbacks = onTickEndCallbacks!;
+      const currentPrimaryCallbacks = onTickEndPrimaryCallbacks!;
+      onTickEndCallbacks = undefined;
+      onTickEndPrimaryCallbacks = undefined;
+      currentPrimaryCallbacks.forEach((cb) => cb());
+      currentCallbacks.forEach((cb) => cb());
+    });
+  } else if (isPrimary) {
+    onTickEndPrimaryCallbacks!.push(callback);
+  } else {
+    onTickEndCallbacks.push(callback);
+  }
+}
+
+export function onTickEndPrimary(callback: NoneToVoidFunction) {
+  onTickEnd(callback, true);
+}
+
 let beforeUnloadCallbacks: NoneToVoidFunction[] | undefined;
 
 export function onBeforeUnload(callback: NoneToVoidFunction, isLast = false) {