Memoize components and improve performance

vercel · Nov 27, 2024 · 2deb633 · 2deb633
1 parent 318927e
commit 2deb633
Show file tree

Hide file tree

Showing 12 changed files with 3,703 additions and 4,254 deletions.
diff --git a/components/block-actions.tsx b/components/block-actions.tsx
@@ -0,0 +1,108 @@
+import { cn } from '@/lib/utils';
+import { CopyIcon, DeltaIcon, RedoIcon, UndoIcon } from './icons';
+import { Button } from './ui/button';
+import { Tooltip, TooltipContent, TooltipTrigger } from './ui/tooltip';
+import { useCopyToClipboard } from 'usehooks-ts';
+import { toast } from 'sonner';
+import { UIBlock } from './block';
+import { memo } from 'react';
+
+interface BlockActionsProps {
+  block: UIBlock;
+  handleVersionChange: (type: 'next' | 'prev' | 'toggle' | 'latest') => void;
+  currentVersionIndex: number;
+  isCurrentVersion: boolean;
+  mode: 'read-only' | 'edit' | 'diff';
+}
+
+export function PureBlockActions({
+  block,
+  handleVersionChange,
+  currentVersionIndex,
+  isCurrentVersion,
+  mode,
+}: BlockActionsProps) {
+  const [_, copyToClipboard] = useCopyToClipboard();
+
+  return (
+    <div className="flex flex-row gap-1">
+      <Tooltip>
+        <TooltipTrigger asChild>
+          <Button
+            variant="outline"
+            className="p-2 h-fit dark:hover:bg-zinc-700"
+            onClick={() => {
+              copyToClipboard(block.content);
+              toast.success('Copied to clipboard!');
+            }}
+            disabled={block.status === 'streaming'}
+          >
+            <CopyIcon size={18} />
+          </Button>
+        </TooltipTrigger>
+        <TooltipContent>Copy to clipboard</TooltipContent>
+      </Tooltip>
+      <Tooltip>
+        <TooltipTrigger asChild>
+          <Button
+            variant="outline"
+            className="p-2 h-fit dark:hover:bg-zinc-700 !pointer-events-auto"
+            onClick={() => {
+              handleVersionChange('prev');
+            }}
+            disabled={currentVersionIndex === 0 || block.status === 'streaming'}
+          >
+            <UndoIcon size={18} />
+          </Button>
+        </TooltipTrigger>
+        <TooltipContent>View Previous version</TooltipContent>
+      </Tooltip>
+      <Tooltip>
+        <TooltipTrigger asChild>
+          <Button
+            variant="outline"
+            className="p-2 h-fit dark:hover:bg-zinc-700 !pointer-events-auto"
+            onClick={() => {
+              handleVersionChange('next');
+            }}
+            disabled={isCurrentVersion || block.status === 'streaming'}
+          >
+            <RedoIcon size={18} />
+          </Button>
+        </TooltipTrigger>
+        <TooltipContent>View Next version</TooltipContent>
+      </Tooltip>
+      <Tooltip>
+        <TooltipTrigger asChild>
+          <Button
+            variant="outline"
+            className={cn(
+              'p-2 h-fit !pointer-events-auto dark:hover:bg-zinc-700',
+              {
+                'bg-muted': mode === 'diff',
+              },
+            )}
+            onClick={() => {
+              handleVersionChange('toggle');
+            }}
+            disabled={block.status === 'streaming' || currentVersionIndex === 0}
+          >
+            <DeltaIcon size={18} />
+          </Button>
+        </TooltipTrigger>
+        <TooltipContent>View changes</TooltipContent>
+      </Tooltip>
+    </div>
+  );
+}
+
+export const BlockActions = memo(PureBlockActions, (prevProps, nextProps) => {
+  if (
+    prevProps.block.status === 'streaming' &&
+    nextProps.block.status === 'streaming'
+  ) {
+    return true;
+  }
+
+  return false;
+});
diff --git a/components/block-close-button.tsx b/components/block-close-button.tsx
@@ -0,0 +1,28 @@
+import { memo, SetStateAction } from 'react';
+import { CrossIcon } from './icons';
+import { Button } from './ui/button';
+import { UIBlock } from './block';
+import equal from 'fast-deep-equal';
+
+interface BlockCloseButtonProps {
+  setBlock: (value: SetStateAction<UIBlock>) => void;
+}
+
+export function PureBlockCloseButton({ setBlock }: BlockCloseButtonProps) {
+  return (
+    <Button
+      variant="outline"
+      className="h-fit p-2 dark:hover:bg-zinc-700"
+      onClick={() => {
+        setBlock((currentBlock) => ({
+          ...currentBlock,
+          isVisible: false,
+        }));
+      }}
+    >
+      <CrossIcon size={18} />
+    </Button>
+  );
+}
+
+export const BlockCloseButton = memo(PureBlockCloseButton, () => true);
diff --git a/components/block.tsx b/components/block.tsx
@@ -9,6 +9,7 @@ import { formatDistance } from 'date-fns';
 import { AnimatePresence, motion } from 'framer-motion';
 import {
   type Dispatch,
+  memo,
   type SetStateAction,
   useCallback,
   useEffect,
@@ -36,6 +37,9 @@ import { Button } from './ui/button';
 import { Tooltip, TooltipContent, TooltipTrigger } from './ui/tooltip';
 import { useScrollToBottom } from './use-scroll-to-bottom';
 import { VersionFooter } from './version-footer';
+import { Markdown } from './markdown';
+import { BlockActions } from './block-actions';
+import { BlockCloseButton } from './block-close-button';
 export interface UIBlock {
   title: string;
   documentId: string;
@@ -50,7 +54,7 @@ export interface UIBlock {
   };
 }
 
-export function Block({
+export function PureBlock({
   chatId,
   input,
   setInput,
@@ -247,8 +251,6 @@ export function Block({
   const { width: windowWidth, height: windowHeight } = useWindowSize();
   const isMobile = windowWidth ? windowWidth < 768 : false;
 
-  const [_, copyToClipboard] = useCopyToClipboard();
-
   return (
     <motion.div
       className="flex flex-row h-dvh w-dvw fixed top-0 left-0 z-50 bg-muted"
@@ -401,18 +403,7 @@ export function Block({
       >
         <div className="p-2 flex flex-row justify-between items-start">
           <div className="flex flex-row gap-4 items-start">
-            <Button
-              variant="outline"
-              className="h-fit p-2 dark:hover:bg-zinc-700"
-              onClick={() => {
-                setBlock((currentBlock) => ({
-                  ...currentBlock,
-                  isVisible: false,
-                }));
-              }}
-            >
-              <CrossIcon size={18} />
-            </Button>
+            <BlockCloseButton setBlock={setBlock} />
 
             <div className="flex flex-col">
               <div className="font-medium">
@@ -439,78 +430,13 @@ export function Block({
             </div>
           </div>
 
-          <div className="flex flex-row gap-1">
-            <Tooltip>
-              <TooltipTrigger asChild>
-                <Button
-                  variant="outline"
-                  className="p-2 h-fit dark:hover:bg-zinc-700"
-                  onClick={() => {
-                    copyToClipboard(block.content);
-                    toast.success('Copied to clipboard!');
-                  }}
-                  disabled={block.status === 'streaming'}
-                >
-                  <CopyIcon size={18} />
-                </Button>
-              </TooltipTrigger>
-              <TooltipContent>Copy to clipboard</TooltipContent>
-            </Tooltip>
-            <Tooltip>
-              <TooltipTrigger asChild>
-                <Button
-                  variant="outline"
-                  className="p-2 h-fit dark:hover:bg-zinc-700 !pointer-events-auto"
-                  onClick={() => {
-                    handleVersionChange('prev');
-                  }}
-                  disabled={
-                    currentVersionIndex === 0 || block.status === 'streaming'
-                  }
-                >
-                  <UndoIcon size={18} />
-                </Button>
-              </TooltipTrigger>
-              <TooltipContent>View Previous version</TooltipContent>
-            </Tooltip>
-            <Tooltip>
-              <TooltipTrigger asChild>
-                <Button
-                  variant="outline"
-                  className="p-2 h-fit dark:hover:bg-zinc-700 !pointer-events-auto"
-                  onClick={() => {
-                    handleVersionChange('next');
-                  }}
-                  disabled={isCurrentVersion || block.status === 'streaming'}
-                >
-                  <RedoIcon size={18} />
-                </Button>
-              </TooltipTrigger>
-              <TooltipContent>View Next version</TooltipContent>
-            </Tooltip>
-            <Tooltip>
-              <TooltipTrigger asChild>
-                <Button
-                  variant="outline"
-                  className={cx(
-                    'p-2 h-fit !pointer-events-auto dark:hover:bg-zinc-700',
-                    {
-                      'bg-muted': mode === 'diff',
-                    },
-                  )}
-                  onClick={() => {
-                    handleVersionChange('toggle');
-                  }}
-                  disabled={
-                    block.status === 'streaming' || currentVersionIndex === 0
-                  }
-                >
-                  <DeltaIcon size={18} />
-                </Button>
-              </TooltipTrigger>
-              <TooltipContent>View changes</TooltipContent>
-            </Tooltip>
-          </div>
+          <BlockActions
+            block={block}
+            currentVersionIndex={currentVersionIndex}
+            handleVersionChange={handleVersionChange}
+            isCurrentVersion={isCurrentVersion}
+            mode={mode}
+          />
         </div>
 
         <div className="prose dark:prose-invert dark:bg-muted bg-background h-full overflow-y-scroll px-4 py-8 md:p-20 !max-w-full pb-40 items-center">
@@ -570,3 +496,7 @@ export function Block({
     </motion.div>
   );
 }
+
+export const Block = memo(PureBlock, (prevProps, nextProps) => {
+  return false;
+});
diff --git a/components/chat-header.tsx b/components/chat-header.tsx
@@ -10,8 +10,13 @@ import { Button } from '@/components/ui/button';
 import { BetterTooltip } from '@/components/ui/tooltip';
 import { PlusIcon, VercelIcon } from './icons';
 import { useSidebar } from './ui/sidebar';
+import { memo } from 'react';
 
-export function ChatHeader({ selectedModelId }: { selectedModelId: string }) {
+export function PureChatHeader({
+  selectedModelId,
+}: {
+  selectedModelId: string;
+}) {
   const router = useRouter();
   const { open } = useSidebar();
 
@@ -54,3 +59,7 @@ export function ChatHeader({ selectedModelId }: { selectedModelId: string }) {
     </header>
   );
 }
+
+export const ChatHeader = memo(PureChatHeader, (prevProps, nextProps) => {
+  return prevProps.selectedModelId === nextProps.selectedModelId;
+});
diff --git a/components/chat.tsx b/components/chat.tsx
@@ -8,15 +8,13 @@ import useSWR, { useSWRConfig } from 'swr';
 import { useWindowSize } from 'usehooks-ts';
 
 import { ChatHeader } from '@/components/chat-header';
-import { PreviewMessage, ThinkingMessage } from '@/components/message';
-import { useScrollToBottom } from '@/components/use-scroll-to-bottom';
 import type { Vote } from '@/lib/db/schema';
 import { fetcher } from '@/lib/utils';
 
 import { Block, type UIBlock } from './block';
 import { BlockStreamHandler } from './block-stream-handler';
 import { MultimodalInput } from './multimodal-input';
-import { Overview } from './overview';
+import { Messages } from './messages';
 
 export function Chat({
   id,
@@ -69,48 +67,22 @@ export function Chat({
     fetcher,
   );
 
-  const [messagesContainerRef, messagesEndRef] =
-    useScrollToBottom<HTMLDivElement>();
-
   const [attachments, setAttachments] = useState<Array<Attachment>>([]);
 
   return (
     <>
       <div className="flex flex-col min-w-0 h-dvh bg-background">
         <ChatHeader selectedModelId={selectedModelId} />
-        <div
-          ref={messagesContainerRef}
-          className="flex flex-col min-w-0 gap-6 flex-1 overflow-y-scroll pt-4"
-        >
-          {messages.length === 0 && <Overview />}
-
-          {messages.map((message, index) => (
-            <PreviewMessage
-              key={message.id}
-              chatId={id}
-              message={message}
-              block={block}
-              setBlock={setBlock}
-              isLoading={isLoading && messages.length - 1 === index}
-              vote={
-                votes
-                  ? votes.find((vote) => vote.messageId === message.id)
-                  : undefined
-              }
-            />
-          ))}
 
-          {isLoading &&
-            messages.length > 0 &&
-            messages[messages.length - 1].role === 'user' && (
-              <ThinkingMessage />
-            )}
+        <Messages
+          chatId={id}
+          block={block}
+          setBlock={setBlock}
+          isLoading={isLoading}
+          votes={votes}
+          messages={messages}
+        />
 
-          <div
-            ref={messagesEndRef}
-            className="shrink-0 min-w-[24px] min-h-[24px]"
-          />
-        </div>
         <form className="flex mx-auto px-4 bg-background pb-4 md:pb-6 gap-2 w-full md:max-w-3xl">
           <MultimodalInput
             chatId={id}