From b9e1eb48ad42c82f6289aeee46c5d14b960611b1 Mon Sep 17 00:00:00 2001
From: Alessio Gravili <alessio@gravili.de>
Date: Mon, 3 Nov 2025 19:04:29 -0800
Subject: [PATCH] perf(ui): reduce unnecessary remounting of pillselector

---
 .../ui/src/elements/PillSelector/index.tsx    | 89 ++++++++++---------
 1 file changed, 45 insertions(+), 44 deletions(-)

diff --git a/packages/ui/src/elements/PillSelector/index.tsx b/packages/ui/src/elements/PillSelector/index.tsx
index 25b04038f82..3bd5338975c 100644
--- a/packages/ui/src/elements/PillSelector/index.tsx
+++ b/packages/ui/src/elements/PillSelector/index.tsx
@@ -30,52 +30,53 @@ export type Props = {
  * If `draggable` is true, the pills can be reordered by dragging.
  */
 export const PillSelector: React.FC<Props> = ({ draggable, onClick, pills }) => {
-  const Wrapper = React.useMemo(() => {
-    if (draggable) {
-      return ({ children }) => (
-        <DraggableSortable
-          className={baseClass}
-          ids={pills.map((pill) => pill.name)}
-          onDragEnd={({ moveFromIndex, moveToIndex }) => {
-            draggable.onDragEnd({
-              moveFromIndex,
-              moveToIndex,
-            })
+  // IMPORTANT: Do NOT wrap DraggableSortable in a dynamic component function using useMemo.
+  // BAD: useMemo(() => ({ children }) => <DraggableSortable>...</DraggableSortable>, [deps])
+  // This creates a new function reference on each recomputation, causing React to treat it as a
+  // different component type, triggering unmount/mount cycles instead of just updating props.
+  // GOOD: Use conditional rendering directly: draggable ? <DraggableSortable /> : <div />
+  const pillElements = React.useMemo(() => {
+    return pills.map((pill, i) => {
+      return (
+        <Pill
+          alignIcon="left"
+          aria-checked={pill.selected}
+          className={[`${baseClass}__pill`, pill.selected && `${baseClass}__pill--selected`]
+            .filter(Boolean)
+            .join(' ')}
+          draggable={Boolean(draggable)}
+          icon={pill.selected ? <XIcon /> : <PlusIcon />}
+          id={pill.name}
+          key={pill.key ?? `${pill.name}-${i}`}
+          onClick={() => {
+            if (onClick) {
+              void onClick({ pill })
+            }
           }}
+          size="small"
         >
-          {children}
-        </DraggableSortable>
+          {pill.Label ?? <span className={`${baseClass}__pill-label`}>{pill.name}</span>}
+        </Pill>
       )
-    } else {
-      return ({ children }) => <div className={baseClass}>{children}</div>
-    }
-  }, [draggable, pills])
+    })
+  }, [pills, onClick, draggable])
 
-  return (
-    <Wrapper>
-      {pills.map((pill, i) => {
-        return (
-          <Pill
-            alignIcon="left"
-            aria-checked={pill.selected}
-            className={[`${baseClass}__pill`, pill.selected && `${baseClass}__pill--selected`]
-              .filter(Boolean)
-              .join(' ')}
-            draggable={Boolean(draggable)}
-            icon={pill.selected ? <XIcon /> : <PlusIcon />}
-            id={pill.name}
-            key={pill.key ?? `${pill.name}-${i}`}
-            onClick={() => {
-              if (onClick) {
-                void onClick({ pill })
-              }
-            }}
-            size="small"
-          >
-            {pill.Label ?? <span className={`${baseClass}__pill-label`}>{pill.name}</span>}
-          </Pill>
-        )
-      })}
-    </Wrapper>
-  )
+  if (draggable) {
+    return (
+      <DraggableSortable
+        className={baseClass}
+        ids={pills.map((pill) => pill.name)}
+        onDragEnd={({ moveFromIndex, moveToIndex }) => {
+          draggable.onDragEnd({
+            moveFromIndex,
+            moveToIndex,
+          })
+        }}
+      >
+        {pillElements}
+      </DraggableSortable>
+    )
+  }
+
+  return <div className={baseClass}>{pillElements}</div>
 }