From f26f30ccb9d7b84946d6867d1631de90318b5a5f Mon Sep 17 00:00:00 2001
From: Alexander Zinchuk <alx.zinchuk@gmail.com>
Date: Tue, 11 May 2021 23:18:02 +0300
Subject: [PATCH] Teact: Fix missing update

---
 src/lib/teact/teact.ts | 31 ++++++++++++++++++++-----------
 1 file changed, 20 insertions(+), 11 deletions(-)

diff --git a/src/lib/teact/teact.ts b/src/lib/teact/teact.ts
index 4f5c4c189..21a770d72 100644
--- a/src/lib/teact/teact.ts
+++ b/src/lib/teact/teact.ts
@@ -80,7 +80,7 @@ interface ComponentInstance {
       }[];
     };
   };
-  scheduleNextState?: () => void;
+  prepareForFrame?: () => void;
   forceUpdate?: () => void;
   onUpdate?: () => void;
 }
@@ -372,8 +372,9 @@ function unmountComponent(componentInstance: ComponentInstance) {
 }
 
 // We force cleaning as many objects as possible. Not sure this is needed at all.
-/* eslint-disable no-null/no-null */
 function helpGc(componentInstance: ComponentInstance) {
+  /* eslint-disable no-null/no-null */
+
   componentInstance.hooks.effects.byCursor.forEach((hook) => {
     hook.cleanup = null as any;
     hook.effect = null as any;
@@ -396,11 +397,11 @@ function helpGc(componentInstance: ComponentInstance) {
   componentInstance.props = null as any;
   componentInstance.forceUpdate = null as any;
   componentInstance.onUpdate = null as any;
+
+  /* eslint-enable no-null/no-null */
 }
 
-/* eslint-enable no-null/no-null */
-
-function applyNextState(componentInstance: ComponentInstance) {
+function prepareComponentForFrame(componentInstance: ComponentInstance) {
   if (!componentInstance.isMounted) {
     return;
   }
@@ -408,6 +409,9 @@ function applyNextState(componentInstance: ComponentInstance) {
   componentInstance.hooks.state.byCursor.forEach((hook) => {
     hook.value = hook.nextValue;
   });
+
+  componentInstance.prepareForFrame = throttleWithPrimaryRaf(() => prepareComponentForFrame(componentInstance));
+  componentInstance.forceUpdate = throttleWithRaf(() => forceUpdateComponent(componentInstance));
 }
 
 function forceUpdateComponent(componentInstance: ComponentInstance) {
@@ -416,6 +420,7 @@ function forceUpdateComponent(componentInstance: ComponentInstance) {
   }
 
   const currentElement = componentInstance.$element;
+
   renderComponent(componentInstance);
 
   if (componentInstance.$element !== currentElement) {
@@ -452,11 +457,18 @@ export function useState<T>(initial?: T): [T, StateHookSetter<T>] {
             ? (newValue as (current: T) => T)(byCursor[cursor].value)
             : newValue;
 
-          if (!componentInstance.scheduleNextState || !componentInstance.forceUpdate) {
-            componentInstance.scheduleNextState = throttleWithPrimaryRaf(() => applyNextState(componentInstance));
-            componentInstance.forceUpdate = throttleWithRaf(() => forceUpdateComponent(componentInstance));
+          if (!componentInstance.prepareForFrame || !componentInstance.forceUpdate) {
+            componentInstance.prepareForFrame = throttleWithPrimaryRaf(
+              () => prepareComponentForFrame(componentInstance),
+            );
+            componentInstance.forceUpdate = throttleWithRaf(
+              () => forceUpdateComponent(componentInstance),
+            );
           }
 
+          componentInstance.prepareForFrame();
+          componentInstance.forceUpdate();
+
           if (DEBUG_MORE) {
             // eslint-disable-next-line no-console
             console.log(
@@ -470,9 +482,6 @@ export function useState<T>(initial?: T): [T, StateHookSetter<T>] {
               byCursor[cursor].nextValue,
             );
           }
-
-          componentInstance.scheduleNextState();
-          componentInstance.forceUpdate();
         }
       })(renderingInstance),
     };