From 29852d2bf6bc9be7550435b31299bfb0faedad94 Mon Sep 17 00:00:00 2001
From: Alexander Zinchuk <alx.zinchuk@gmail.com>
Date: Fri, 21 Jan 2022 17:29:48 +0100
Subject: [PATCH] [Perf] Teact: Fix memory leaks

---
 src/components/middle/hooks/useScrollHooks.ts |  6 +-
 src/lib/teact/teact-dom.ts                    | 14 ++--
 src/lib/teact/teact.ts                        | 67 ++++++++++---------
 3 files changed, 51 insertions(+), 36 deletions(-)

diff --git a/src/components/middle/hooks/useScrollHooks.ts b/src/components/middle/hooks/useScrollHooks.ts
index f7591019f..07a4cd1b7 100644
--- a/src/components/middle/hooks/useScrollHooks.ts
+++ b/src/components/middle/hooks/useScrollHooks.ts
@@ -60,7 +60,11 @@ export default function useScrollHooks(
       return;
     }
 
-    const { offsetHeight, scrollHeight, scrollTop } = containerRef.current!;
+    if (!containerRef.current) {
+      return;
+    }
+
+    const { offsetHeight, scrollHeight, scrollTop } = containerRef.current;
     const scrollBottom = Math.round(scrollHeight - scrollTop - offsetHeight);
     const isNearBottom = scrollBottom <= FAB_THRESHOLD;
     const isAtBottom = scrollBottom <= NOTCH_THRESHOLD;
diff --git a/src/lib/teact/teact-dom.ts b/src/lib/teact/teact-dom.ts
index 42b594148..67cfdaa74 100644
--- a/src/lib/teact/teact-dom.ts
+++ b/src/lib/teact/teact-dom.ts
@@ -122,14 +122,20 @@ function renderWithVirtual(
       unmountTree($current);
     } else {
       const areComponents = isComponentElement($current) && isComponentElement($new);
+      const currentTarget = getTarget($current);
 
       if (!areComponents) {
-        setTarget($new, getTarget($current)!);
+        setTarget($new, currentTarget!);
+        setTarget($current, undefined as any); // Help GC
+
+        if ('props' in $current && 'props' in $new) {
+          $new.props.ref = $current.props.ref;
+        }
       }
 
       if (isRealElement($current) && isRealElement($new)) {
         if (moveDirection) {
-          const node = getTarget($current)!;
+          const node = currentTarget!;
           const nextSibling = parentEl.childNodes[moveDirection === 'up' ? index : index + 1];
 
           if (nextSibling) {
@@ -140,13 +146,13 @@ function renderWithVirtual(
         }
 
         if (!areComponents) {
-          updateAttributes($current, $new, getTarget($current) as HTMLElement);
+          updateAttributes($current, $new, currentTarget as HTMLElement);
         }
 
         $new.children = renderChildren(
           $current,
           $new,
-          areComponents ? parentEl : getTarget($current) as HTMLElement,
+          areComponents ? parentEl : currentTarget as HTMLElement,
         );
       }
     }
diff --git a/src/lib/teact/teact.ts b/src/lib/teact/teact.ts
index 18f74b60d..61e42beae 100644
--- a/src/lib/teact/teact.ts
+++ b/src/lib/teact/teact.ts
@@ -336,17 +336,26 @@ export function hasElementChanged($old: VirtualElement, $new: VirtualElement) {
 }
 
 export function unmountTree($element: VirtualElement) {
-  if (!isRealElement($element)) {
-    return;
-  }
-
   if (isComponentElement($element)) {
     unmountComponent($element.componentInstance);
-  } else if ($element.target) {
-    removeAllDelegatedListeners($element.target as HTMLElement);
-    // Trying to help GC
-    // eslint-disable-next-line no-null/no-null
-    $element.target = null as any;
+  } else {
+    if (isTagElement($element)) {
+      if ($element.target) {
+        removeAllDelegatedListeners($element.target as HTMLElement);
+      }
+
+      if ($element.props.ref) {
+        $element.props.ref.current = undefined; // Help GC
+      }
+    }
+
+    if ($element.target) {
+      $element.target = undefined; // Help GC
+    }
+
+    if (!isRealElement($element)) {
+      return;
+    }
   }
 
   $element.children.forEach(unmountTree);
@@ -363,9 +372,9 @@ function unmountComponent(componentInstance: ComponentInstance) {
     return;
   }
 
-  componentInstance.hooks.memos.byCursor.forEach((hook) => {
-    // eslint-disable-next-line no-null/no-null
-    hook.current = null;
+  // We need to clean refs before running effect cleanups
+  componentInstance.hooks.memos.byCursor.forEach((memoContainer) => {
+    memoContainer.current = undefined;
   });
 
   componentInstance.hooks.effects.byCursor.forEach(({ cleanup }) => {
@@ -383,35 +392,31 @@ function unmountComponent(componentInstance: ComponentInstance) {
   helpGc(componentInstance);
 }
 
-// We need to remove all references to DOM objects. We also clean all other references, just in case.
+// We need to remove all references to DOM objects. We also clean all other references, just in case
 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;
-    hook.dependencies = null as any;
+    hook.cleanup = undefined;
+    hook.effect = undefined as any;
+    hook.dependencies = undefined;
   });
 
   componentInstance.hooks.state.byCursor.forEach((hook) => {
-    hook.value = null as any;
-    hook.nextValue = null as any;
-    hook.setter = null as any;
+    hook.value = undefined;
+    hook.nextValue = undefined;
+    hook.setter = undefined as any;
   });
 
   componentInstance.hooks.memos.byCursor.forEach((hook) => {
-    hook.dependencies = null as any;
+    hook.dependencies = undefined as any;
   });
 
-  componentInstance.hooks = null as any;
-  componentInstance.$element = null as any;
-  componentInstance.renderedValue = null as any;
-  componentInstance.Component = null as any;
-  componentInstance.props = null as any;
-  componentInstance.forceUpdate = null as any;
-  componentInstance.onUpdate = null as any;
-
-  /* eslint-enable no-null/no-null */
+  componentInstance.hooks = undefined as any;
+  componentInstance.$element = undefined as any;
+  componentInstance.renderedValue = undefined;
+  componentInstance.Component = undefined as any;
+  componentInstance.props = undefined as any;
+  componentInstance.forceUpdate = undefined;
+  componentInstance.onUpdate = undefined;
 }
 
 function prepareComponentForFrame(componentInstance: ComponentInstance) {