From c4277b9ef02b0028ec2cd1df571fd1cf62425b45 Mon Sep 17 00:00:00 2001 From: Enrico Ros Date: Wed, 7 Feb 2024 01:53:28 -0800 Subject: [PATCH] Optimization on the message being typed - recycles references to speed up React. Fixes #402 --- .../message/blocks/BlocksRenderer.tsx | 50 +++++++++++++++---- .../chat/components/message/blocks/blocks.ts | 21 ++++++++ src/modules/aifn/digrams/DiagramsModal.tsx | 3 -- 3 files changed, 61 insertions(+), 13 deletions(-) diff --git a/src/apps/chat/components/message/blocks/BlocksRenderer.tsx b/src/apps/chat/components/message/blocks/BlocksRenderer.tsx index 7445451af..d896fae40 100644 --- a/src/apps/chat/components/message/blocks/BlocksRenderer.tsx +++ b/src/apps/chat/components/message/blocks/BlocksRenderer.tsx @@ -16,7 +16,7 @@ import { RenderLatex } from './RenderLatex'; import { RenderMarkdown } from './RenderMarkdown'; import { RenderText } from './RenderText'; import { RenderTextDiff } from './RenderTextDiff'; -import { parseMessageBlocks } from './blocks'; +import { areBlocksEqual, Block, parseMessageBlocks } from './blocks'; // How long is the user collapsed message @@ -67,6 +67,7 @@ export function BlocksRenderer(props: { // state const [forceUserExpanded, setForceUserExpanded] = React.useState(false); + const prevBlocksRef = React.useRef([]); // derived state const { text: _text, errorMessage, renderTextDiff, wasUserEdited = false } = props; @@ -75,7 +76,12 @@ export function BlocksRenderer(props: { const fromUser = props.fromRole === 'user'; - // Memo text, blocks and styles + const handleTextUncollapse = React.useCallback(() => { + setForceUserExpanded(true); + }, []); + + + // Memo text, which could be 'collapsed' to a few lines in case of user messages const { text, isTextCollapsed } = React.useMemo(() => { if (fromUser && !forceUserExpanded) { @@ -86,15 +92,12 @@ export function BlocksRenderer(props: { return { text: _text, isTextCollapsed: false }; }, [forceUserExpanded, fromUser, _text]); - const blocks = React.useMemo(() => { - const blocks = errorMessage ? [] : parseMessageBlocks(text, fromSystem, renderTextDiff); - return props.specialDiagramMode ? blocks.filter(block => block.type === 'code' || blocks.length === 1) : blocks; - }, [errorMessage, fromSystem, props.specialDiagramMode, renderTextDiff, text]); + // Memo the code style, to minimize re-renders const codeSx: SxProps = React.useMemo(() => ( { backgroundColor: props.specialDiagramMode ? 'background.surface' : fromAssistant ? 'neutral.plainHoverBg' : 'primary.plainActiveBg', - boxShadow: 'xs', + boxShadow: props.specialDiagramMode ? 'md' : 'xs', fontFamily: 'code', fontSize: '0.875rem', fontVariantLigatures: 'none', @@ -104,9 +107,36 @@ export function BlocksRenderer(props: { ), [fromAssistant, props.specialDiagramMode]); - const handleTextUncollapse = React.useCallback(() => { - setForceUserExpanded(true); - }, []); + // Block splitter, with memoand special recycle of former blocks, to help React minimize render work + + const blocks = React.useMemo(() => { + // split the complete input text into blocks + const newBlocks = errorMessage ? [] : parseMessageBlocks(text, fromSystem, renderTextDiff); + + // recycle the previous blocks if they are the same, for stable references to React + const recycledBlocks: Block[] = []; + for (let i = 0; i < newBlocks.length; i++) { + const newBlock = newBlocks[i]; + const prevBlock = prevBlocksRef.current[i]; + + // Check if the new block can be replaced by the previous block to maintain reference stability + if (prevBlock && areBlocksEqual(prevBlock, newBlock)) { + recycledBlocks.push(prevBlock); + } else { + // Once a block doesn't match, we use the new blocks from this point forward. + recycledBlocks.push(...newBlocks.slice(i)); + break; + } + } + + // Update prevBlocksRef with the current blocks for the next render + prevBlocksRef.current = recycledBlocks; + + // Apply specialDiagramMode filter if applicable + return props.specialDiagramMode + ? recycledBlocks.filter(block => block.type === 'code' || recycledBlocks.length === 1) + : recycledBlocks; + }, [errorMessage, fromSystem, props.specialDiagramMode, renderTextDiff, text]); return ( diff --git a/src/apps/chat/components/message/blocks/blocks.ts b/src/apps/chat/components/message/blocks/blocks.ts index 9b1de1f97..c64529edb 100644 --- a/src/apps/chat/components/message/blocks/blocks.ts +++ b/src/apps/chat/components/message/blocks/blocks.ts @@ -13,6 +13,27 @@ export type LatexBlock = { type: 'latex'; latex: string; }; export type TextBlock = { type: 'text'; content: string; }; // for Text or Markdown +export function areBlocksEqual(a: Block, b: Block): boolean { + if (a.type !== b.type) + return false; + + switch (a.type) { + case 'code': + return a.blockTitle === (b as CodeBlock).blockTitle && a.blockCode === (b as CodeBlock).blockCode && a.complete === (b as CodeBlock).complete; + case 'diff': + return false; // diff blocks are never equal + case 'html': + return a.html === (b as HtmlBlock).html; + case 'image': + return a.url === (b as ImageBlock).url && a.alt === (b as ImageBlock).alt; + case 'latex': + return a.latex === (b as LatexBlock).latex; + case 'text': + return a.content === (b as TextBlock).content; + } +} + + export function parseMessageBlocks(text: string, disableParsing: boolean, forceTextDiffs?: TextDiff[]): Block[] { if (disableParsing) return [{ type: 'text', content: text }]; diff --git a/src/modules/aifn/digrams/DiagramsModal.tsx b/src/modules/aifn/digrams/DiagramsModal.tsx index 2686d718a..f12ae79a3 100644 --- a/src/modules/aifn/digrams/DiagramsModal.tsx +++ b/src/modules/aifn/digrams/DiagramsModal.tsx @@ -188,9 +188,6 @@ export function DiagramsModal(props: { config: DiagramConfig, onClose: () => voi marginX: 'calc(-1 * var(--Card-padding))', minHeight: 96, p: { xs: 1, md: 2 }, - '& > div > div > code': { - boxShadow: 'md', - }, }}>