fix: memoize components and improve performance (#579)

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';
+}
+
+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;
+});

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;
+}
+
+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);

components/block-messages.tsx

@@ -0,0 +1,71 @@
+import { Dispatch, memo, SetStateAction } from 'react';
+import { UIBlock } from './block';
+import { PreviewMessage } from './message';
+import { useScrollToBottom } from './use-scroll-to-bottom';
+import { Vote } from '@/lib/db/schema';
+import { Message } from 'ai';
+
+interface BlockMessagesProps {
+  chatId: string;
+  block: UIBlock;
+  setBlock: Dispatch<SetStateAction<UIBlock>>;
+  isLoading: boolean;
+  votes: Array<Vote> | undefined;
+  messages: Array<Message>;
+}
+
+function PureBlockMessages({
+  chatId,
+  block,
+  setBlock,
+  isLoading,
+  votes,
+  messages,
+}: BlockMessagesProps) {
+  const [messagesContainerRef, messagesEndRef] =
+    useScrollToBottom<HTMLDivElement>();
+
+  return (
+    <div
+      ref={messagesContainerRef}
+      className="flex flex-col gap-4 h-full items-center overflow-y-scroll px-4 pt-20"
+    >
+      {messages.map((message, index) => (
+        <PreviewMessage
+          chatId={chatId}
+          key={message.id}
+          message={message}
+          block={block}
+          setBlock={setBlock}
+          isLoading={isLoading && index === messages.length - 1}
+          vote={
+            votes
+              ? votes.find((vote) => vote.messageId === message.id)
+              : undefined
+          }
+        />
+      ))}
+
+      <div
+        ref={messagesEndRef}
+        className="shrink-0 min-w-[24px] min-h-[24px]"
+      />
+    </div>
+  );
+}
+
+function areEqual(
+  prevProps: BlockMessagesProps,
+  nextProps: BlockMessagesProps,
+) {
+  if (
+    prevProps.block.status === 'streaming' &&
+    nextProps.block.status === 'streaming'
+  ) {
+    return true;
+  }
+
+  return false;
+}
+
+export const BlockMessages = memo(PureBlockMessages, areEqual);

components/block-stream-handler.tsx

@@ -9,7 +9,7 @@ interface BlockStreamHandlerProps {
   streamingData: JSONValue[] | undefined;
 }
 
-export function PureBlockStreamHandler({
+function PureBlockStreamHandler({
   setBlock,
   streamingData,
 }: BlockStreamHandlerProps) {