chore: memoize components (#607)

components/block-actions.tsx

@@ -97,12 +97,10 @@ function PureBlockActions({
 }
 
 export const BlockActions = memo(PureBlockActions, (prevProps, nextProps) => {
-  if (
-    prevProps.block.status === 'streaming' &&
-    nextProps.block.status === 'streaming'
-  ) {
-    return true;
-  }
-
+  if (prevProps.block.status !== nextProps.block.status) return false;
+  if (prevProps.currentVersionIndex !== nextProps.currentVersionIndex)
     return false;
+  if (prevProps.isCurrentVersion !== nextProps.isCurrentVersion) return false;
+
+  return true;
 });

components/message-actions.tsx

@@ -14,8 +14,10 @@ import {
   TooltipProvider,
   TooltipTrigger,
 } from './ui/tooltip';
+import { memo } from 'react';
+import equal from 'fast-deep-equal';
 
-export function MessageActions({
+export function PureMessageActions({
   chatId,
   message,
   vote,
@@ -164,3 +166,13 @@ export function MessageActions({
     </TooltipProvider>
   );
 }
+
+export const MessageActions = memo(
+  PureMessageActions,
+  (prevProps, nextProps) => {
+    if (!equal(prevProps.vote, nextProps.vote)) return false;
+    if (prevProps.isLoading !== nextProps.isLoading) return false;
+
+    return true;
+  },
+);

components/message.tsx

@@ -223,9 +223,16 @@ export const PreviewMessage = memo(
   PurePreviewMessage,
   (prevProps, nextProps) => {
     if (prevProps.isLoading !== nextProps.isLoading) return false;
-    if (prevProps.isLoading && nextProps.isLoading) return false;
-    if (prevProps.message.content && nextProps.message.content) return false;
+    if (prevProps.message.content !== nextProps.message.content) return false;
+    if (
+      !equal(
+        prevProps.message.toolInvocations,
+        nextProps.message.toolInvocations,
+      )
+    )
+      return false;
     if (!equal(prevProps.vote, nextProps.vote)) return false;
+
     return true;
   },
 );

components/messages.tsx

@@ -5,6 +5,7 @@ import { Overview } from './overview';
 import { UIBlock } from './block';
 import { Dispatch, memo, SetStateAction } from 'react';
 import { Vote } from '@/lib/db/schema';
+import equal from 'fast-deep-equal';
 
 interface MessagesProps {
   chatId: string;
@@ -74,15 +75,11 @@ function PureMessages({
   );
 }
 
-function areEqual(prevProps: MessagesProps, nextProps: MessagesProps) {
-  if (
-    prevProps.block.status === 'streaming' &&
-    nextProps.block.status === 'streaming'
-  ) {
+export const Messages = memo(PureMessages, (prevProps, nextProps) => {
+  if (prevProps.isLoading !== nextProps.isLoading) return false;
+  if (prevProps.isLoading && nextProps.isLoading) return false;
+  if (prevProps.messages.length !== nextProps.messages.length) return false;
+  if (!equal(prevProps.votes, nextProps.votes)) return false;
+
   return true;
-  }
-
-  return false;
-}
-
-export const Messages = memo(PureMessages, areEqual);
+});

components/multimodal-input.tsx

@@ -247,40 +247,16 @@ function PureMultimodalInput({
       />
 
       {isLoading ? (
-        <Button
-          className="rounded-full p-1.5 h-fit absolute bottom-2 right-2 m-0.5 border dark:border-zinc-600"
-          onClick={(event) => {
-            event.preventDefault();
-            stop();
-            setMessages((messages) => sanitizeUIMessages(messages));
-          }}
-        >
-          <StopIcon size={14} />
-        </Button>
+        <StopButton stop={stop} setMessages={setMessages} />
       ) : (
-        <Button
-          className="rounded-full p-1.5 h-fit absolute bottom-2 right-2 m-0.5 border dark:border-zinc-600"
-          onClick={(event) => {
-            event.preventDefault();
-            submitForm();
-          }}
-          disabled={input.length === 0 || uploadQueue.length > 0}
-        >
-          <ArrowUpIcon size={14} />
-        </Button>
+        <SendButton
+          input={input}
+          submitForm={submitForm}
+          uploadQueue={uploadQueue}
+        />
       )}
 
-      <Button
-        className="rounded-full p-1.5 h-fit absolute bottom-2 right-11 m-0.5 dark:border-zinc-700"
-        onClick={(event) => {
-          event.preventDefault();
-          fileInputRef.current?.click();
-        }}
-        variant="outline"
-        disabled={isLoading}
-      >
-        <PaperclipIcon size={14} />
-      </Button>
+      <AttachmentsButton fileInputRef={fileInputRef} isLoading={isLoading} />
     </div>
   );
 }
@@ -295,3 +271,80 @@ export const MultimodalInput = memo(
     return true;
   },
 );
+
+function PureAttachmentsButton({
+  fileInputRef,
+  isLoading,
+}: {
+  fileInputRef: React.MutableRefObject<HTMLInputElement | null>;
+  isLoading: boolean;
+}) {
+  return (
+    <Button
+      className="rounded-full p-1.5 h-fit absolute bottom-2 right-11 m-0.5 dark:border-zinc-700"
+      onClick={(event) => {
+        event.preventDefault();
+        fileInputRef.current?.click();
+      }}
+      variant="outline"
+      disabled={isLoading}
+    >
+      <PaperclipIcon size={14} />
+    </Button>
+  );
+}
+
+const AttachmentsButton = memo(PureAttachmentsButton);
+
+function PureStopButton({
+  stop,
+  setMessages,
+}: {
+  stop: () => void;
+  setMessages: Dispatch<SetStateAction<Array<Message>>>;
+}) {
+  return (
+    <Button
+      className="rounded-full p-1.5 h-fit absolute bottom-2 right-2 m-0.5 border dark:border-zinc-600"
+      onClick={(event) => {
+        event.preventDefault();
+        stop();
+        setMessages((messages) => sanitizeUIMessages(messages));
+      }}
+    >
+      <StopIcon size={14} />
+    </Button>
+  );
+}
+
+const StopButton = memo(PureStopButton);
+
+function PureSendButton({
+  submitForm,
+  input,
+  uploadQueue,
+}: {
+  submitForm: () => void;
+  input: string;
+  uploadQueue: Array<string>;
+}) {
+  return (
+    <Button
+      className="rounded-full p-1.5 h-fit absolute bottom-2 right-2 m-0.5 border dark:border-zinc-600"
+      onClick={(event) => {
+        event.preventDefault();
+        submitForm();
+      }}
+      disabled={input.length === 0 || uploadQueue.length > 0}
+    >
+      <ArrowUpIcon size={14} />
+    </Button>
+  );
+}
+
+const SendButton = memo(PureSendButton, (prevProps, nextProps) => {
+  if (prevProps.uploadQueue.length !== nextProps.uploadQueue.length)
+    return false;
+  if (!prevProps.input !== !nextProps.input) return false;
+  return true;
+});

components/toolbar.tsx

@@ -469,5 +469,8 @@ const PureToolbar = ({
 };
 
 export const Toolbar = memo(PureToolbar, (prevProps, nextProps) => {
-  return equal(prevProps, nextProps);
+  if (prevProps.isLoading !== nextProps.isLoading) return false;
+  if (prevProps.isToolbarVisible !== nextProps.isToolbarVisible) return false;
+
+  return true;
 });