fix(core): bound lastSeen vnode tracking (#630)

Author: JoviDeCroock

.changeset/tidy-dogs-smile.md

@@ -0,0 +1,5 @@
+---
+'@prefresh/core': patch
+---
+
+Prevent stale component vnodes from being retained across rerenders and unmounts

packages/core/src/runtime/unmount.js

@@ -1,15 +1,19 @@
 import { options } from 'preact';
-import { vnodesForComponent } from './vnodesForComponent';
+import { vnodesForComponent, clearLastSeen } from './vnodesForComponent';
 
 const oldUnmount = options.unmount;
 options.unmount = vnode => {
   const type = (vnode || {}).type;
-  if (typeof type === 'function' && vnodesForComponent.has(type)) {
-    const vnodes = vnodesForComponent.get(type);
-    if (vnodes) {
-      const index = vnodes.indexOf(vnode);
-      if (index !== -1) {
-        vnodes.splice(index, 1);
+  if (typeof type === 'function') {
+    clearLastSeen(vnode);
+
+    if (vnodesForComponent.has(type)) {
+      const vnodes = vnodesForComponent.get(type);
+      if (vnodes) {
+        const index = vnodes.indexOf(vnode);
+        if (index !== -1) {
+          vnodes.splice(index, 1);
+        }
       }
     }
   }

packages/core/src/runtime/vnode.js

@@ -2,7 +2,7 @@ import { options } from 'preact';
 import {
   vnodesForComponent,
   mappedVNodes,
-  lastSeen,
+  setLastSeen,
 } from './vnodesForComponent';
 import { VNODE_COMPONENT } from '../constants';
 
@@ -57,7 +57,7 @@ const oldDiffed = options.diffed;
 options.diffed = vnode => {
   if (vnode && typeof vnode.type === 'function') {
     const vnodes = vnodesForComponent.get(vnode.type);
-    lastSeen.set(vnode.__v, vnode);
+    setLastSeen(vnode);
     if (vnodes) {
       const matchingDom = vnodes.filter(p => p.__c === vnode.__c);
       if (matchingDom.length > 1) {

packages/core/src/runtime/vnodesForComponent.js

@@ -2,3 +2,27 @@
 export const vnodesForComponent = new WeakMap();
 export const mappedVNodes = new WeakMap();
 export const lastSeen = new Map();
+
+const lastSeenByInstance = new WeakMap();
+
+export const setLastSeen = vnode => {
+  const component = vnode.__c;
+  if (component) {
+    if (lastSeenByInstance.has(component)) {
+      lastSeen.delete(lastSeenByInstance.get(component));
+    }
+    lastSeenByInstance.set(component, vnode.__v);
+  }
+
+  lastSeen.set(vnode.__v, vnode);
+};
+
+export const clearLastSeen = vnode => {
+  lastSeen.delete(vnode.__v);
+
+  const component = vnode.__c;
+  if (component && lastSeenByInstance.has(component)) {
+    lastSeen.delete(lastSeenByInstance.get(component));
+    lastSeenByInstance.delete(component);
+  }
+};

test/fixture/vite-babel/src/app.jsx

@@ -1,4 +1,5 @@
 import { h } from 'preact';
+import { useState } from 'preact/hooks';
 import { useCounter } from './useCounter';
 import { StoreProvider } from './context';
 import { Products } from './products';
@@ -10,6 +11,9 @@ import { Style } from './styles';
 
 setup(h);
 
+const lastSeenPayloads = [];
+self.__lastSeenPayloads = lastSeenPayloads;
+
 function Test() {
   const [count, increment] = useCounter();
   return (
@@ -20,10 +24,30 @@ function Test() {
   )
 }
 
+function LastSeenChild({ payload }) {
+  return <span className="last-seen-child">{payload.count}</span>;
+}
+
+function LastSeenTest() {
+  const [count, setCount] = useState(0);
+  const payload = { count };
+  lastSeenPayloads.push(new WeakRef(payload));
+
+  return (
+    <div>
+      <button className="last-seen-increment" onClick={() => setCount(count + 1)}>
+        Increment tracked component
+      </button>
+      <LastSeenChild payload={payload} />
+    </div>
+  );
+}
+
 export function App(props) {
   return (
     <Style id="color">
       <Test />
+      <LastSeenTest />
       <Greeting />
       <StoreProvider>
         <Products />

test/fixture/vite/src/app.jsx

@@ -1,4 +1,5 @@
 import { h } from 'preact';
+import { useState } from 'preact/hooks';
 import { useCounter } from './useCounter';
 import { StoreProvider } from './context';
 import { Products } from './products';
@@ -10,6 +11,9 @@ import { Style } from './styles';
 
 setup(h);
 
+const lastSeenPayloads = [];
+self.__lastSeenPayloads = lastSeenPayloads;
+
 function Test() {
   const [count, increment] = useCounter();
   return (
@@ -20,10 +24,30 @@ function Test() {
   )
 }
 
+function LastSeenChild({ payload }) {
+  return <span className="last-seen-child">{payload.count}</span>;
+}
+
+function LastSeenTest() {
+  const [count, setCount] = useState(0);
+  const payload = { count };
+  lastSeenPayloads.push(new WeakRef(payload));
+
+  return (
+    <div>
+      <button className="last-seen-increment" onClick={() => setCount(count + 1)}>
+        Increment tracked component
+      </button>
+      <LastSeenChild payload={payload} />
+    </div>
+  );
+}
+
 export function App(props) {
   return (
     <Style id="color">
       <Test />
+      <LastSeenTest />
       <Greeting />
       <StoreProvider>
         <Products />

test/index.test.js

@@ -127,6 +127,29 @@ describe('Prefresh integrations', () => {
         await expectByPolling(() => getText(button), 'Increment (+)');
       });
 
+      test('does not retain stale vnodes after rerenders', async () => {
+        const increment = await page.$('.last-seen-increment');
+
+        await page.evaluate(() => {
+          self.__lastSeenPayloads.length = 0;
+        });
+
+        for (let i = 0; i < 20; i++) {
+          await increment.click();
+        }
+
+        await expectByPolling(() => getText('.last-seen-child'), '20');
+
+        const client = await page.target().createCDPSession();
+        await client.send('HeapProfiler.collectGarbage');
+        expect(
+          await page.evaluate(
+            () => self.__lastSeenPayloads.filter(ref => ref.deref()).length
+          )
+        ).toBe(1);
+        await client.detach();
+      });
+
       test('add export', async () => {
         const button = await page.$('.button');
         await expectByPolling(() => getText(button), 'Increment (+)');