From 096e2784d2d556fe557ba597bcbbb69fee4eed38 Mon Sep 17 00:00:00 2001
From: Enrico Ros <enrico.ros@gmail.com>
Date: Tue, 23 Jul 2024 02:52:08 -0700
Subject: [PATCH] AttachmentButton: massive performance boost

---
 .../llmattachments/LLMAttachmentButton.tsx    |  4 +-
 .../llmattachments/LLMAttachmentsList.tsx     |  4 +-
 .../llmattachments/useLLMAttachmentDrafts.ts  | 52 +++++++++++++++----
 3 files changed, 48 insertions(+), 12 deletions(-)

diff --git a/src/apps/chat/components/composer/llmattachments/LLMAttachmentButton.tsx b/src/apps/chat/components/composer/llmattachments/LLMAttachmentButton.tsx
index b02fe5652..4a1004259 100644
--- a/src/apps/chat/components/composer/llmattachments/LLMAttachmentButton.tsx
+++ b/src/apps/chat/components/composer/llmattachments/LLMAttachmentButton.tsx
@@ -165,7 +165,9 @@ function attachmentLabelText(attachmentDraft: AttachmentDraft): string {
 }
 
 
-export function LLMAttachmentButton(props: {
+export const LLMAttachmentButtonMemo = React.memo(LLMAttachmentButton);
+
+function LLMAttachmentButton(props: {
   llmAttachment: LLMAttachmentDraft,
   menuShown: boolean,
   onToggleMenu: (attachmentDraftId: AttachmentDraftId, anchor: HTMLAnchorElement) => void,
diff --git a/src/apps/chat/components/composer/llmattachments/LLMAttachmentsList.tsx b/src/apps/chat/components/composer/llmattachments/LLMAttachmentsList.tsx
index 0e4cbf67e..070478554 100644
--- a/src/apps/chat/components/composer/llmattachments/LLMAttachmentsList.tsx
+++ b/src/apps/chat/components/composer/llmattachments/LLMAttachmentsList.tsx
@@ -13,7 +13,7 @@ import type { AttachmentDraftId } from '~/common/attachment-drafts/attachment.ty
 import type { AttachmentDraftsStoreApi } from '~/common/attachment-drafts/store-attachment-drafts-slice';
 
 import type { LLMAttachmentDrafts } from './useLLMAttachmentDrafts';
-import { LLMAttachmentButton } from './LLMAttachmentButton';
+import { LLMAttachmentButtonMemo } from './LLMAttachmentButton';
 import { LLMAttachmentMenu } from './LLMAttachmentMenu';
 
 
@@ -107,7 +107,7 @@ export function LLMAttachmentsList(props: {
       {/* Horizontally scrollable Attachments */}
       <Box sx={{ display: 'flex', overflowX: 'auto', gap: 1, height: '100%', pr: 5 }}>
         {llmAttachmentDrafts.map((llmAttachment) =>
-          <LLMAttachmentButton
+          <LLMAttachmentButtonMemo
             key={llmAttachment.attachmentDraft.id}
             llmAttachment={llmAttachment}
             menuShown={llmAttachment.attachmentDraft.id === itemMenuAttachmentDraftId}
diff --git a/src/apps/chat/components/composer/llmattachments/useLLMAttachmentDrafts.ts b/src/apps/chat/components/composer/llmattachments/useLLMAttachmentDrafts.ts
index 24d1b4968..8076f5685 100644
--- a/src/apps/chat/components/composer/llmattachments/useLLMAttachmentDrafts.ts
+++ b/src/apps/chat/components/composer/llmattachments/useLLMAttachmentDrafts.ts
@@ -24,22 +24,52 @@ export interface LLMAttachmentDraft {
 
 
 export function useLLMAttachmentDrafts(attachmentDrafts: AttachmentDraft[], chatLLM: DLLM | null): LLMAttachmentDrafts {
+
+  /* [Optimization] Use a Ref to store the previous state of llmAttachmentDrafts and chatLLM
+   *
+   * Note that this works on 2 levels:
+   * - 1. avoids recomputation, but more importantly,
+   * - 2. avoids re-rendering by keeping those llmAttachmentDrafts objects stable.
+   *
+   * Important to notice that the attachmentDraft objects[] are stable to start with, so we can
+   * safely use reference equality to check if internal properties (or order) have changed.
+   */
+  const prevStateRef = React.useRef<{
+    chatLLM: DLLM | null;
+    llmAttachmentDrafts: LLMAttachmentDraft[];
+  }>({ llmAttachmentDrafts: [], chatLLM: null });
+
   return React.useMemo(() => {
 
+    // [Optimization]
+    const equalChatLLM = chatLLM === prevStateRef.current.chatLLM;
+
     // LLM-dependent multi-modal enablement
     const supportsImages = !!chatLLM?.interfaces?.includes(LLM_IF_OAI_Vision);
     const supportedTypes: DMessageAttachmentFragment['part']['pt'][] = supportsImages ? ['image_ref', 'doc'] : ['doc'];
     const supportedTextTypes: DMessageAttachmentFragment['part']['pt'][] = supportedTypes.filter(pt => pt === 'doc');
 
     // Add LLM-specific properties to each attachment draft
-    const llmAttachmentDrafts = attachmentDrafts.map((a): LLMAttachmentDraft => ({
-      attachmentDraft: a,
-      llmSupportsAllFragments: !a.outputFragments ? false : a.outputFragments.every(op => supportedTypes.includes(op.part.pt)),
-      llmSupportsTextFragments: !a.outputFragments ? false : a.outputFragments.some(op => supportedTextTypes.includes(op.part.pt)),
-      llmTokenCountApprox: chatLLM
-        ? estimateTokensForFragments(chatLLM, 'user', a.outputFragments, true, 'useLLMAttachmentDrafts')
-        : null,
-    }));
+    const llmAttachmentDrafts = attachmentDrafts.map((a, index) => {
+
+      // [Optimization] If not change in LLM and the attachmentDraft is the same object reference, reuse the previous LLMAttachmentDraft
+      let prevDraft: LLMAttachmentDraft | undefined = prevStateRef.current.llmAttachmentDrafts[index];
+      // if not found, search by id
+      if (!prevDraft)
+        prevDraft = prevStateRef.current.llmAttachmentDrafts.find(_pd => _pd.attachmentDraft.id === a.id);
+      if (equalChatLLM && prevDraft && prevDraft.attachmentDraft === a)
+        return prevDraft;
+
+      // Otherwise, create a new LLMAttachmentDraft
+      return {
+        attachmentDraft: a,
+        llmSupportsAllFragments: !a.outputFragments ? false : a.outputFragments.every(op => supportedTypes.includes(op.part.pt)),
+        llmSupportsTextFragments: !a.outputFragments ? false : a.outputFragments.some(op => supportedTextTypes.includes(op.part.pt)),
+        llmTokenCountApprox: chatLLM
+          ? estimateTokensForFragments(chatLLM, 'user', a.outputFragments, true, 'useLLMAttachmentDrafts')
+          : null,
+      };
+    });
 
     // Calculate the overall properties
     const canAttachAllFragments = llmAttachmentDrafts.every(a => a.llmSupportsAllFragments);
@@ -48,11 +78,15 @@ export function useLLMAttachmentDrafts(attachmentDrafts: AttachmentDraft[], chat
       ? llmAttachmentDrafts.reduce((acc, a) => acc + (a.llmTokenCountApprox || 0), 0)
       : null;
 
+    // [Optimization] Update the ref with the new state
+    prevStateRef.current = { llmAttachmentDrafts, chatLLM };
+
     return {
       llmAttachmentDrafts,
       canAttachAllFragments,
       canInlineSomeFragments,
       llmTokenCountApprox,
     };
-  }, [attachmentDrafts, chatLLM]);
+
+  }, [attachmentDrafts, chatLLM]); // Dependencies for the outer useMemo
 }