From 2c3af5e73ce4d42cdfe90cfbdc400a94fd56f29e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Greg=20Berg=C3=A9?= <berge.greg@gmail.com>
Date: Mon, 3 Mar 2025 14:47:30 +0100
Subject: [PATCH] Improve scrolling with lot of highlighted blocks (#2909)

---
 .../CodeBlock/ClientCodeBlock.tsx             | 76 ++++++++++---------
 .../components/hooks/useInViewportListener.ts | 49 ++++++++++++
 2 files changed, 91 insertions(+), 34 deletions(-)
 create mode 100644 packages/gitbook/src/components/hooks/useInViewportListener.ts

diff --git a/packages/gitbook/src/components/DocumentView/CodeBlock/ClientCodeBlock.tsx b/packages/gitbook/src/components/DocumentView/CodeBlock/ClientCodeBlock.tsx
index 671d9e811..0c1111298 100644
--- a/packages/gitbook/src/components/DocumentView/CodeBlock/ClientCodeBlock.tsx
+++ b/packages/gitbook/src/components/DocumentView/CodeBlock/ClientCodeBlock.tsx
@@ -3,8 +3,9 @@
 import type { DocumentBlockCode } from '@gitbook/api';
 import { useEffect, useRef, useState } from 'react';
 
-import { useHasBeenInViewport } from '@/components/hooks/useHasBeenInViewport';
-
+import { useInViewportListener } from '@/components/hooks/useInViewportListener';
+import { useScrollListener } from '@/components/hooks/useScrollListener';
+import { useDebounceCallback, useEventCallback } from 'usehooks-ts';
 import type { BlockProps } from '../Block';
 import { CodeBlockRenderer } from './CodeBlockRenderer';
 import type { HighlightLine, RenderedInline } from './highlight';
@@ -21,6 +22,8 @@ type ClientBlockProps = Pick<BlockProps<DocumentBlockCode>, 'block' | 'style'> &
 export function ClientCodeBlock(props: ClientBlockProps) {
     const { block, style, inlines } = props;
     const blockRef = useRef<HTMLDivElement>(null);
+    const processedRef = useRef(false);
+    const isInViewportRef = useRef<boolean | null>(null);
     const [lines, setLines] = useState<HighlightLine[]>(() => plainHighlight(block, []));
 
     // Preload the highlighter when the block is mounted.
@@ -28,39 +31,44 @@ export function ClientCodeBlock(props: ClientBlockProps) {
         import('./highlight').then(({ preloadHighlight }) => preloadHighlight(block));
     }, [block]);
 
-    // Check if the block is in the viewport to start highlighting it.
-    const hasBeenInViewport = useHasBeenInViewport(blockRef, {
-        rootMargin: '200px',
-    });
-
-    // Highlight the block when it's in the viewport.
-    useEffect(() => {
-        if (hasBeenInViewport) {
-            let canceled = false;
-            import('./highlight').then(({ highlight }) => {
-                // We use requestIdleCallback to avoid blocking the main thread
-                // when scrolling.
-                if (typeof requestIdleCallback === 'function') {
-                    requestIdleCallback(() =>
-                        highlight(block, inlines).then((result) => {
-                            if (!canceled) {
-                                setLines(result);
-                            }
-                        })
-                    );
-                } else {
-                    highlight(block, inlines).then((result) => {
-                        if (!canceled) {
-                            setLines(result);
-                        }
-                    });
-                }
-            });
-            return () => {
-                canceled = true;
-            };
+    const runHighlight = useEventCallback(() => {
+        if (processedRef.current) {
+            return;
         }
-    }, [hasBeenInViewport, block, inlines]);
+        if (typeof window !== 'undefined') {
+            import('./highlight').then(({ highlight }) => {
+                highlight(block, inlines).then((lines) => {
+                    setLines(lines);
+                    processedRef.current = true;
+                });
+            });
+        }
+    });
+    const debouncedRunHighlight = useDebounceCallback(runHighlight, 1000);
+
+    useInViewportListener(
+        blockRef,
+        (isInViewport, disconnect) => {
+            // Disconnect once in viewport
+            if (isInViewport) {
+                disconnect();
+                // If it's initially in viewport, we need to run the highlight
+                if (isInViewportRef.current === null) {
+                    runHighlight();
+                }
+            }
+            isInViewportRef.current = isInViewport;
+        },
+        { rootMargin: '200px' }
+    );
+
+    const handleScroll = useDebounceCallback(() => {
+        if (isInViewportRef.current) {
+            debouncedRunHighlight();
+        }
+    }, 80);
+
+    useScrollListener(handleScroll, useRef(typeof window !== 'undefined' ? window : null));
 
     return <CodeBlockRenderer ref={blockRef} block={block} style={style} lines={lines} />;
 }
diff --git a/packages/gitbook/src/components/hooks/useInViewportListener.ts b/packages/gitbook/src/components/hooks/useInViewportListener.ts
new file mode 100644
index 000000000..8f44fd826
--- /dev/null
+++ b/packages/gitbook/src/components/hooks/useInViewportListener.ts
@@ -0,0 +1,49 @@
+'use client';
+import { useEffect, useLayoutEffect, useRef } from 'react';
+
+const HAS_INTERSECTION_OBSERVER = typeof IntersectionObserver !== 'undefined';
+
+/**
+ * Watch an element to know when it is in the viewport.
+ */
+export function useInViewportListener(
+    containerRef: React.RefObject<HTMLElement>,
+    listener: (isIntersecting: boolean, disconnect: () => void) => void,
+    options?: Pick<IntersectionObserverInit, 'root' | 'rootMargin' | 'threshold'>
+) {
+    const listenerRef = useRef(listener);
+    useLayoutEffect(() => {
+        listenerRef.current = listener;
+    });
+    const isIntersectingRef = useRef(false);
+    useEffect(() => {
+        // We set the element as visible if the IntersectionObserver API is not available.
+        // we have to do it in the `useEffect` to be SSR compatible.
+        if (!HAS_INTERSECTION_OBSERVER) {
+            listenerRef.current(true, () => {});
+            return;
+        }
+
+        if (!containerRef.current) {
+            return;
+        }
+
+        const observer = new IntersectionObserver(
+            ([entry]) => {
+                isIntersectingRef.current = entry.isIntersecting;
+                listenerRef.current(entry.isIntersecting, () => {
+                    observer.disconnect();
+                });
+            },
+            {
+                root: options?.root,
+                rootMargin: options?.rootMargin,
+                threshold: options?.threshold,
+            }
+        );
+
+        observer.observe(containerRef.current);
+
+        return () => observer.disconnect();
+    }, [containerRef, options?.root, options?.rootMargin, options?.threshold]);
+}