From 153c1b2c55c2a1e4d0c63d7ef9daac8344247ddd Mon Sep 17 00:00:00 2001
From: pengx17 <pengxiao@outlook.com>
Date: Fri, 27 Dec 2024 08:34:04 +0000
Subject: [PATCH] refactor(core): make lit portal a bit cleaner (#9354)

---
 .../src/lit-react/lit-portal/lit-portal.tsx   | 90 ++++++++-----------
 .../bi-directional-link-panel.tsx             |  2 +-
 .../block-suite-editor/lit-adaper.tsx         |  2 +-
 3 files changed, 41 insertions(+), 53 deletions(-)

diff --git a/packages/frontend/component/src/lit-react/lit-portal/lit-portal.tsx b/packages/frontend/component/src/lit-react/lit-portal/lit-portal.tsx
index c22eb4ff87..6282361320 100644
--- a/packages/frontend/component/src/lit-react/lit-portal/lit-portal.tsx
+++ b/packages/frontend/component/src/lit-react/lit-portal/lit-portal.tsx
@@ -7,12 +7,11 @@ import ReactDOM from 'react-dom';
 type PortalEvent = {
   name: 'connectedCallback' | 'disconnectedCallback' | 'willUpdate';
   target: LitReactPortal;
-  previousPortalId?: string;
 };
 
 type PortalListener = (event: PortalEvent) => void;
 
-export function createLitPortalAnchor(callback: (event: PortalEvent) => void) {
+function createLitPortalAnchor(callback: (event: PortalEvent) => void) {
   return html`<lit-react-portal
     .notify=${callback}
     portalId=${nanoid()}
@@ -24,7 +23,6 @@ export const LIT_REACT_PORTAL = 'lit-react-portal';
 @customElement(LIT_REACT_PORTAL)
 class LitReactPortal extends LitElement {
   portalId!: string;
-
   notify!: PortalListener;
 
   static override get properties() {
@@ -34,6 +32,14 @@ class LitReactPortal extends LitElement {
     };
   }
 
+  override connectedCallback() {
+    super.connectedCallback();
+    this.notify({
+      name: 'connectedCallback',
+      target: this,
+    });
+  }
+
   override attributeChangedCallback(
     name: string,
     oldVal: string,
@@ -44,7 +50,6 @@ class LitReactPortal extends LitElement {
       this.notify({
         name: 'willUpdate',
         target: this,
-        previousPortalId: oldVal,
       });
     }
   }
@@ -54,13 +59,6 @@ class LitReactPortal extends LitElement {
     return this;
   }
 
-  override connectedCallback() {
-    this.notify({
-      name: 'connectedCallback',
-      target: this,
-    });
-  }
-
   override disconnectedCallback() {
     super.disconnectedCallback();
     this.notify({
@@ -81,6 +79,7 @@ export type ElementOrFactory = React.ReactElement | (() => React.ReactElement);
 type LitPortal = {
   id: string;
   portal: React.ReactPortal;
+  litElement: LitReactPortal;
 };
 
 // returns a factory function that renders a given element to a lit template
@@ -98,48 +97,37 @@ export const useLitPortalFactory = () => {
             ? elementOrFactory()
             : elementOrFactory;
         return createLitPortalAnchor(event => {
-          const { name, target } = event;
-          const id = target.portalId;
-
-          if (name === 'connectedCallback') {
-            setPortals(portals => [
-              ...portals,
-              {
-                id,
-                portal: ReactDOM.createPortal(element, target, id),
-              },
-            ]);
-            return;
-          }
-
-          if (name === 'disconnectedCallback') {
-            setPortals(portals => portals.filter(p => p.id !== id));
-            return;
-          }
-
-          const prevId = event.previousPortalId;
-
-          // Ignores first `willUpdate`
-          if (!prevId) {
-            return;
-          }
-
           setPortals(portals => {
-            const portal = portals.find(p => p.id === prevId);
-            if (!portal) return [...portals];
-
-            // Updates `ID`
-            // Used for portal queries in `disconnectedCallback`
-            portal.id = id;
-
-            // Re-rendering
-            //  true: `inline link`
-            //  false: `pdf embed view`
-            if (rerendering) {
-              portal.portal = ReactDOM.createPortal(element, target, id);
+            const { name, target } = event;
+            const id = target.portalId;
+            let newPortals = portals;
+            const updatePortals = () => {
+              let oldPortalIndex = portals.findIndex(
+                p => p.litElement === target
+              );
+              oldPortalIndex =
+                oldPortalIndex === -1 ? portals.length : oldPortalIndex;
+              newPortals = portals.toSpliced(oldPortalIndex, 1, {
+                id,
+                portal: ReactDOM.createPortal(element, target),
+                litElement: target,
+              });
+            };
+            switch (name) {
+              case 'connectedCallback':
+                updatePortals();
+                break;
+              case 'disconnectedCallback':
+                newPortals = portals.filter(p => p.litElement.isConnected);
+                break;
+              case 'willUpdate':
+                if (!target.isConnected || !rerendering) {
+                  break;
+                }
+                updatePortals();
+                break;
             }
-
-            return [...portals];
+            return newPortals;
           });
         });
       },
diff --git a/packages/frontend/core/src/components/blocksuite/block-suite-editor/bi-directional-link-panel.tsx b/packages/frontend/core/src/components/blocksuite/block-suite-editor/bi-directional-link-panel.tsx
index e45e5b7875..9225fa6fac 100644
--- a/packages/frontend/core/src/components/blocksuite/block-suite-editor/bi-directional-link-panel.tsx
+++ b/packages/frontend/core/src/components/blocksuite/block-suite-editor/bi-directional-link-panel.tsx
@@ -357,7 +357,7 @@ export const BiDirectionalLinkPanel = () => {
       {
         <>
           {portals.map(p => (
-            <Fragment key={p.portal.key}>{p.portal}</Fragment>
+            <Fragment key={p.id}>{p.portal}</Fragment>
           ))}
         </>
       }
diff --git a/packages/frontend/core/src/components/blocksuite/block-suite-editor/lit-adaper.tsx b/packages/frontend/core/src/components/blocksuite/block-suite-editor/lit-adaper.tsx
index f9fd1acff5..0da7345fdf 100644
--- a/packages/frontend/core/src/components/blocksuite/block-suite-editor/lit-adaper.tsx
+++ b/packages/frontend/core/src/components/blocksuite/block-suite-editor/lit-adaper.tsx
@@ -194,7 +194,7 @@ const usePatchSpecs = (shared: boolean, mode: DocMode) => {
       () => (
         <>
           {portals.map(p => (
-            <Fragment key={p.portal.key}>{p.portal}</Fragment>
+            <Fragment key={p.id}>{p.portal}</Fragment>
           ))}
         </>
       ),