perf: Sidebar virtualization and re-render optimization (#12443)

* perf: Prevent action context invalidation on location change

* PR feedback

* virtualization

* fix: Initial visiblity incorrect

* PR feedback

app/components/Sidebar/App.tsx

@@ -1,8 +1,12 @@
 import { observer } from "mobx-react";
 import { SearchIcon, HomeIcon, SidebarIcon } from "outline-icons";
-import { useEffect, useState, useCallback, useMemo } from "react";
+import { useEffect, useState, useCallback, useMemo, useRef } from "react";
 import { DndProvider } from "react-dnd";
 import { HTML5Backend } from "react-dnd-html5-backend";
+import {
+  DragActiveProvider,
+  SidebarScrollProvider,
+} from "./components/DragActiveContext";
 import { useTranslation } from "react-i18next";
 import { useHistory } from "react-router-dom";
 import styled from "styled-components";
@@ -67,78 +71,92 @@ function AppSidebar() {
     [dndArea]
   );
 
+  // Scrollable reads ref.current internally for its shadow/ResizeObserver
+  // logic, so we must pass an object ref — a callback ref would leave those
+  // reads undefined. We mirror the attached node into state so the
+  // SidebarScrollProvider can re-render descendants with the scroll element.
+  const scrollRef = useRef<HTMLDivElement>(null);
+  const [scrollArea, setScrollArea] = useState<HTMLElement | null>(null);
+  useEffect(() => {
+    setScrollArea(scrollRef.current);
+  }, []);
+
   return (
     <Sidebar hidden={!ui.readyToShow} ref={handleSidebarRef}>
       {dndArea && (
         <DndProvider backend={HTML5Backend} options={html5Options}>
-          <DragPlaceholder />
+          <DragActiveProvider>
+            <DragPlaceholder />
 
-          <TeamMenu>
-            <SidebarButton
-              title={team.name}
-              image={<TeamLogo model={team} size={24} alt={t("Logo")} />}
-            >
-              {isMobile ? null : (
-                <Tooltip
-                  content={t("Toggle sidebar")}
-                  shortcut={`${metaDisplay}+.`}
-                >
-                  <ToggleButton
-                    position="bottom"
-                    image={<SidebarIcon />}
-                    aria-label={
-                      ui.sidebarCollapsed
-                        ? t("Expand sidebar")
-                        : t("Collapse sidebar")
-                    }
-                    style={{ paddingInline: 4 }}
-                    onClick={() => {
-                      ui.toggleCollapsedSidebar();
-                      (document.activeElement as HTMLElement)?.blur();
-                    }}
-                  />
-                </Tooltip>
-              )}
-            </SidebarButton>
-          </TeamMenu>
-          <Overflow>
-            <Section>
-              <SidebarLink
-                to={homePath()}
-                icon={<HomeIcon />}
-                exact={false}
-                label={t("Home")}
-              />
-              <SidebarLink
-                to={searchPath()}
-                icon={<SearchIcon />}
-                label={t("Search")}
-                exact={false}
-                onClick={handleSearchClick}
-              />
-              {can.createDocument && <DraftsLink />}
-            </Section>
-          </Overflow>
-          <Scrollable flex shadow>
-            <Section>
-              <Starred />
-            </Section>
-            <Section>
-              <SharedWithMe />
-            </Section>
-            <Section>
-              <Collections />
-            </Section>
-            {can.createDocument && (
-              <Section auto>
-                <ArchiveLink />
+            <TeamMenu>
+              <SidebarButton
+                title={team.name}
+                image={<TeamLogo model={team} size={24} alt={t("Logo")} />}
+              >
+                {isMobile ? null : (
+                  <Tooltip
+                    content={t("Toggle sidebar")}
+                    shortcut={`${metaDisplay}+.`}
+                  >
+                    <ToggleButton
+                      position="bottom"
+                      image={<SidebarIcon />}
+                      aria-label={
+                        ui.sidebarCollapsed
+                          ? t("Expand sidebar")
+                          : t("Collapse sidebar")
+                      }
+                      style={{ paddingInline: 4 }}
+                      onClick={() => {
+                        ui.toggleCollapsedSidebar();
+                        (document.activeElement as HTMLElement)?.blur();
+                      }}
+                    />
+                  </Tooltip>
+                )}
+              </SidebarButton>
+            </TeamMenu>
+            <Overflow>
+              <Section>
+                <SidebarLink
+                  to={homePath()}
+                  icon={<HomeIcon />}
+                  exact={false}
+                  label={t("Home")}
+                />
+                <SidebarLink
+                  to={searchPath()}
+                  icon={<SearchIcon />}
+                  label={t("Search")}
+                  exact={false}
+                  onClick={handleSearchClick}
+                />
+                {can.createDocument && <DraftsLink />}
               </Section>
-            )}
-            <Section>
-              {can.createDocument && <TrashLink />}
-              <SidebarAction action={inviteUser} />
-            </Section>
-          </Scrollable>
+            </Overflow>
+            <Scrollable flex shadow ref={scrollRef}>
+              <SidebarScrollProvider value={scrollArea}>
+                <Section>
+                  <Starred />
+                </Section>
+                <Section>
+                  <SharedWithMe />
+                </Section>
+                <Section>
+                  <Collections />
+                </Section>
+                {can.createDocument && (
+                  <Section auto>
+                    <ArchiveLink />
+                  </Section>
+                )}
+                <Section>
+                  {can.createDocument && <TrashLink />}
+                  <SidebarAction action={inviteUser} />
+                </Section>
+              </SidebarScrollProvider>
+            </Scrollable>
+          </DragActiveProvider>
         </DndProvider>
       )}
       <HistoryNavigation />

app/components/Sidebar/components/DocumentLink.tsx

@@ -16,6 +16,7 @@ import type { RefHandle } from "~/components/EditableTitle";
 import useBoolean from "~/hooks/useBoolean";
 import useCurrentUser from "~/hooks/useCurrentUser";
 import { useDocumentMenuAction } from "~/hooks/useDocumentMenuAction";
+import useOnScreen from "~/hooks/useOnScreen";
 import usePolicy from "~/hooks/usePolicy";
 import useStores from "~/hooks/useStores";
 import DocumentMenu from "~/menus/DocumentMenu";
@@ -25,6 +26,7 @@ import {
   useDropToReorderDocument,
   useDropToReparentDocument,
 } from "../hooks/useDragAndDrop";
+import { useIsDragActive, useSidebarScrollElement } from "./DragActiveContext";
 import { useSidebarExpansion } from "./SidebarExpansionContext";
 import DocumentRow from "./DocumentRow";
 import DropCursor from "./DropCursor";
@@ -44,34 +46,28 @@ type Props = {
   parentId?: string;
 };
 
-const DocumentLink = observer(function DocumentLinkInner({
-  node,
-  collection,
-  membership,
-  activeDocument,
-  prefetchDocument,
-  isDraft,
-  depth,
-  index,
-  parentId,
-}: Props) {
+// Approximate rendered row height; used to reserve space for unmounted rows so
+// the scroll container stays the right height and IntersectionObserver triggers
+// correctly as the user scrolls.
+const ROW_HEIGHT = 30;
+
+// Pre-mount rows just outside the viewport so scrolling stays smooth and drop
+// targets exist a screen ahead when a drag starts.
+const ROOT_MARGIN = "300px 0px";
+
+const DocumentLink = observer(function DocumentLink(props: Props) {
+  const { node, collection, activeDocument } = props;
   const { documents } = useStores();
-  const { t } = useTranslation();
-  const history = useHistory();
-  const can = usePolicy(node.id);
-  const canUpdate = can.update;
+  const expansion = useSidebarExpansion();
+  const expanded = expansion.isExpanded(node.id);
   const isActiveDocument = activeDocument && activeDocument.id === node.id;
   const hasChildDocuments =
     !!node.children.length || activeDocument?.parentDocumentId === node.id;
-  const document = documents.get(node.id);
-  const { fetchChildDocuments } = documents;
-  const [isEditing, setIsEditing] = React.useState(false);
-  const editableTitleRef = React.useRef<RefHandle>(null);
   const sidebarContext = useSidebarContext();
-  const user = useCurrentUser();
-  const expansion = useSidebarExpansion();
-  const expanded = expansion.isExpanded(node.id);
+  const { fetchChildDocuments } = documents;
 
+  // Keep expansion/data effects on the outer so they run regardless of whether
+  // the heavy row content is currently mounted.
   React.useEffect(() => {
     if (expanded && !hasChildDocuments) {
       expansion.collapse(node.id);
@@ -93,6 +89,105 @@ const DocumentLink = observer(function DocumentLinkInner({
     isActiveDocument,
   ]);
 
+  const insertDraftChild = !!(
+    activeDocument?.isDraft &&
+    activeDocument?.isActive &&
+    activeDocument?.parentDocumentId === node.id
+  );
+
+  const draftNavNode = insertDraftChild
+    ? activeDocument?.asNavigationNode
+    : undefined;
+
+  const nodeChildren = React.useMemo(
+    () =>
+      collection && draftNavNode
+        ? sortNavigationNodes(
+            [draftNavNode, ...node.children],
+            collection.sort,
+            false
+          )
+        : node.children,
+    [draftNavNode, collection, node.children]
+  );
+
+  // Visibility gate: only mount the heavy inner content when scrolled near the
+  // viewport, but keep it mounted while a drag is in progress so the dragged
+  // source (or a drop target the user is heading toward) isn't yanked.
+  const scrollRoot = useSidebarScrollElement();
+  const placeholderRef = React.useRef<HTMLDivElement>(null);
+  const observerOptions = React.useMemo(
+    () => ({ root: scrollRoot, rootMargin: ROOT_MARGIN }),
+    [scrollRoot]
+  );
+  const isOnScreen = useOnScreen(placeholderRef, observerOptions);
+  const isDragActive = useIsDragActive();
+  const [mounted, setMounted] = React.useState(false);
+
+  // Flip mount state during render (not in an effect) so the first paint
+  // already contains the row content when the placeholder is on screen,
+  // avoiding a blank frame.
+  if (isOnScreen && !mounted) {
+    setMounted(true);
+  } else if (!isOnScreen && !isDragActive && mounted) {
+    setMounted(false);
+  }
+
+  return (
+    <>
+      <div ref={placeholderRef} style={{ minHeight: ROW_HEIGHT }}>
+        {mounted ? (
+          <DocumentLinkInner {...props} hasChildren={nodeChildren.length > 0} />
+        ) : null}
+      </div>
+      <Folder expanded={expanded}>
+        {nodeChildren.map((childNode, childIndex) => (
+          <DocumentLink
+            key={childNode.id}
+            collection={collection}
+            membership={props.membership}
+            node={childNode}
+            activeDocument={activeDocument}
+            prefetchDocument={props.prefetchDocument}
+            isDraft={childNode.isDraft}
+            depth={props.depth + 1}
+            index={childIndex}
+            parentId={node.id}
+          />
+        ))}
+      </Folder>
+    </>
+  );
+});
+
+type InnerProps = Props & {
+  hasChildren: boolean;
+};
+
+const DocumentLinkInner = observer(function DocumentLinkInner({
+  node,
+  collection,
+  membership,
+  prefetchDocument,
+  isDraft,
+  depth,
+  index,
+  parentId,
+  hasChildren,
+}: InnerProps) {
+  const { documents } = useStores();
+  const { t } = useTranslation();
+  const history = useHistory();
+  const can = usePolicy(node.id);
+  const canUpdate = can.update;
+  const document = documents.get(node.id);
+  const [isEditing, setIsEditing] = React.useState(false);
+  const editableTitleRef = React.useRef<RefHandle>(null);
+  const sidebarContext = useSidebarContext();
+  const user = useCurrentUser();
+  const expansion = useSidebarExpansion();
+  const expanded = expansion.isExpanded(node.id);
+
   const handleDisclosureClick = React.useCallback(
     (ev?: React.MouseEvent<HTMLElement>) => {
       if (expanded) {
@@ -226,30 +321,7 @@ const DocumentLink = observer(function DocumentLinkInner({
       };
     });
 
-  const insertDraftChild = !!(
-    activeDocument?.isDraft &&
-    activeDocument?.isActive &&
-    activeDocument?.parentDocumentId === node.id
-  );
-
-  const draftNavNode = insertDraftChild
-    ? activeDocument?.asNavigationNode
-    : undefined;
-
-  const nodeChildren = React.useMemo(
-    () =>
-      collection && draftNavNode
-        ? sortNavigationNodes(
-            [draftNavNode, ...node.children],
-            collection.sort,
-            false
-          )
-        : node.children,
-    [draftNavNode, collection, node.children]
-  );
-
   const title = document?.title || node.title || t("Untitled");
-  const hasChildren = nodeChildren.length > 0;
 
   const handleNewDoc = React.useCallback(
     async (input: string) => {
@@ -348,24 +420,7 @@ const DocumentLink = observer(function DocumentLinkInner({
       contextAction={contextMenuAction}
       isActiveOverride={isActiveCheck}
       onClickIntent={handlePrefetch}
-    >
-      <Folder expanded={expanded && !isDragging}>
-        {nodeChildren.map((childNode, childIndex) => (
-          <DocumentLink
-            key={childNode.id}
-            collection={collection}
-            membership={membership}
-            node={childNode}
-            activeDocument={activeDocument}
-            prefetchDocument={prefetchDocument}
-            isDraft={childNode.isDraft}
-            depth={depth + 1}
-            index={childIndex}
-            parentId={node.id}
-          />
-        ))}
-      </Folder>
-    </DocumentRow>
+    />
   );
 });
 

app/components/Sidebar/components/DragActiveContext.tsx

@@ -0,0 +1,47 @@
+import * as React from "react";
+import { useDragLayer } from "react-dnd";
+
+const DragActiveContext = React.createContext(false);
+
+const SidebarScrollContext = React.createContext<HTMLElement | null>(null);
+
+/**
+ * Provides the sidebar's scroll container so descendants can use it as the
+ * IntersectionObserver root when deciding whether to render heavy content.
+ */
+export const SidebarScrollProvider = SidebarScrollContext.Provider;
+
+/**
+ * Returns the sidebar scroll container element, or null if not within a
+ * SidebarScrollProvider.
+ */
+export function useSidebarScrollElement(): HTMLElement | null {
+  return React.useContext(SidebarScrollContext);
+}
+
+/**
+ * Subscribes once to react-dnd's drag state and exposes a boolean via context.
+ *
+ * Visibility-gated sidebar rows read this to keep their inner content mounted
+ * for the duration of a drag, so that scrolling away from the dragged source
+ * (or a drop target the user is heading toward) does not unmount it mid-drag.
+ */
+export function DragActiveProvider({
+  children,
+}: {
+  children: React.ReactNode;
+}) {
+  const isDragging = useDragLayer((monitor) => monitor.isDragging());
+  return (
+    <DragActiveContext.Provider value={isDragging}>
+      {children}
+    </DragActiveContext.Provider>
+  );
+}
+
+/**
+ * Returns whether any react-dnd drag is currently active.
+ */
+export function useIsDragActive(): boolean {
+  return React.useContext(DragActiveContext);
+}

app/hooks/useActionContext.tsx

@@ -1,8 +1,8 @@
 import { observer } from "mobx-react";
 import type { ReactNode } from "react";
-import React, { createContext, useContext } from "react";
+import React, { createContext, useCallback, useContext, useMemo } from "react";
 import { useTranslation } from "react-i18next";
-import { useLocation } from "react-router";
+import { useHistory } from "react-router";
 import useStores from "~/hooks/useStores";
 import type Model from "~/models/base/Model";
 import type Policy from "~/models/Policy";
@@ -55,98 +55,159 @@ export const ActionContextProvider = observer(function ActionContextProvider_({
   const parentContext = useContext(ActionContext);
   const stores = useStores();
   const { t } = useTranslation();
-  const location = useLocation();
-  const { activeModels: valueModels, ...overrides } = value;
 
-  // Create the base context if we don't have a parent context
-  const baseContext: ActionContextType = parentContext ?? {
-    isMenu: false,
-    isCommandBar: false,
-    isButton: false,
+  // Use history (stable reference) and read location lazily via a getter so
+  // navigation does not invalidate the context value. Action perform/visible
+  // callbacks see the current location at call time via history.location,
+  // which react-router updates on every navigation.
+  const history = useHistory();
 
-    // Legacy (backward compatibility)
-    activeCollectionId: stores.ui.activeCollectionId ?? undefined,
-    activeDocumentId: stores.ui.activeDocumentId ?? undefined,
+  const {
+    activeModels: valueModels,
+    isMenu,
+    isCommandBar,
+    isButton,
+    sidebarContext,
+    event,
+  } = value;
 
-    getActiveModels: <T extends Model>(
-      modelClass: new (...args: never[]) => T
-    ): T[] => stores.ui.getActiveModels<T>(modelClass),
+  // Track membership of stores.ui.activeModels so memos invalidate when it changes.
+  // Reading inside the observer-wrapped render keeps MobX subscriptions intact.
+  const activeModelsKey = Array.from(stores.ui.activeModels.keys()).join(",");
+  const activeCollectionIdFromStore = stores.ui.activeCollectionId ?? undefined;
+  const activeDocumentIdFromStore = stores.ui.activeDocumentId ?? undefined;
+  const currentUserId = stores.auth.user?.id;
+  const currentTeamId = stores.auth.team?.id;
 
-    getActiveModel: <T extends Model>(
-      modelClass: new (...args: never[]) => T
-    ): T | undefined => stores.ui.getActiveModels<T>(modelClass)[0],
+  const getActiveModels = useCallback(
+    <T extends Model>(modelClass: new (...args: never[]) => T): T[] => {
+      if (valueModels && valueModels.length > 0) {
+        const matching = valueModels.filter(
+          (model): model is T => model instanceof modelClass
+        );
+        if (matching.length > 0) {
+          return matching;
+        }
+      }
+      if (parentContext) {
+        return parentContext.getActiveModels(modelClass);
+      }
+      return stores.ui.getActiveModels<T>(modelClass);
+    },
+    [valueModels, parentContext, stores]
+  );
 
-    getActivePolicies: <T extends Model>(
-      modelClass: new (...args: never[]) => T
-    ): Policy[] =>
-      stores.ui
-        .getActiveModels<T>(modelClass)
+  const getActiveModel = useCallback(
+    <T extends Model>(modelClass: new (...args: never[]) => T): T | undefined =>
+      getActiveModels(modelClass)[0],
+    [getActiveModels]
+  );
+
+  const getActivePolicies = useCallback(
+    <T extends Model>(modelClass: new (...args: never[]) => T): Policy[] =>
+      getActiveModels(modelClass)
         .map((node) => stores.policies.get(node.id))
         .filter((policy): policy is Policy => policy !== undefined),
+    [getActiveModels, stores]
+  );
 
-    isModelActive: (model: Model): boolean => stores.ui.isModelActive(model),
-    activeModels: new Set(stores.ui.activeModels.values()),
+  const allActiveModels = useMemo(() => {
+    const base = parentContext
+      ? parentContext.activeModels
+      : new Set(stores.ui.activeModels.values());
+    if (valueModels && valueModels.length > 0) {
+      return new Set([...base, ...valueModels]);
+    }
+    return base;
+    // eslint-disable-next-line react-hooks/exhaustive-deps
+  }, [parentContext, stores, valueModels, activeModelsKey]);
 
-    currentUserId: stores.auth.user?.id,
-    currentTeamId: stores.auth.team?.id,
-    location,
+  const isModelActive = useCallback(
+    (model: Model): boolean => allActiveModels.has(model),
+    [allActiveModels]
+  );
+
+  const contextValue = useMemo<ActionContextType>(() => {
+    const baseContext: ActionContextType = parentContext ?? {
+      isMenu: false,
+      isCommandBar: false,
+      isButton: false,
+
+      // Legacy (backward compatibility)
+      activeCollectionId: activeCollectionIdFromStore,
+      activeDocumentId: activeDocumentIdFromStore,
+
+      getActiveModels,
+      getActiveModel,
+      getActivePolicies,
+      isModelActive,
+      activeModels: allActiveModels,
+
+      currentUserId,
+      currentTeamId,
+      // Consumers reading `ctx.location` get the current location at access time.
+      location: history.location,
+      stores,
+      t,
+    };
+
+    // Derive legacy IDs from value models, falling back to base context
+    const activeCollectionId =
+      valueModels?.find(
+        (m) => (m.constructor as typeof Model).modelName === "Collection"
+      )?.id ?? baseContext.activeCollectionId;
+
+    const activeDocumentId =
+      valueModels?.find(
+        (m) => (m.constructor as typeof Model).modelName === "Document"
+      )?.id ?? baseContext.activeDocumentId;
+
+    const result = {
+      ...baseContext,
+      ...(isMenu !== undefined ? { isMenu } : {}),
+      ...(isCommandBar !== undefined ? { isCommandBar } : {}),
+      ...(isButton !== undefined ? { isButton } : {}),
+      ...(sidebarContext !== undefined ? { sidebarContext } : {}),
+      ...(event !== undefined ? { event } : {}),
+      activeCollectionId,
+      activeDocumentId,
+      getActiveModels,
+      getActiveModel,
+      getActivePolicies,
+      isModelActive,
+      activeModels: allActiveModels,
+    };
+
+    // Define `location` as a getter so reads always return the current
+    // location without invalidating this memo on navigation.
+    Object.defineProperty(result, "location", {
+      get: () => history.location,
+      enumerable: true,
+      configurable: true,
+    });
+
+    return result;
+  }, [
+    parentContext,
     stores,
     t,
-  };
-
-  // Override model accessors when models are provided in value
-  const getActiveModels =
-    valueModels && valueModels.length > 0
-      ? <T extends Model>(modelClass: new (...args: never[]) => T): T[] => {
-          const matching = valueModels.filter(
-            (model): model is T => model instanceof modelClass
-          );
-          return matching.length > 0
-            ? matching
-            : baseContext.getActiveModels(modelClass);
-        }
-      : baseContext.getActiveModels;
-
-  const getActiveModel = <T extends Model>(
-    modelClass: new (...args: never[]) => T
-  ): T | undefined => getActiveModels(modelClass)[0];
-
-  const getActivePolicies = <T extends Model>(
-    modelClass: new (...args: never[]) => T
-  ): Policy[] =>
-    getActiveModels(modelClass)
-      .map((node) => stores.policies.get(node.id))
-      .filter((policy): policy is Policy => policy !== undefined);
-
-  const allActiveModels =
-    valueModels && valueModels.length > 0
-      ? new Set([...baseContext.activeModels, ...valueModels])
-      : baseContext.activeModels;
-
-  const isModelActive = (model: Model): boolean => allActiveModels.has(model);
-
-  // Derive legacy IDs from value models, falling back to base context
-  const activeCollectionId =
-    valueModels?.find(
-      (m) => (m.constructor as typeof Model).modelName === "Collection"
-    )?.id ?? baseContext.activeCollectionId;
-
-  const activeDocumentId =
-    valueModels?.find(
-      (m) => (m.constructor as typeof Model).modelName === "Document"
-    )?.id ?? baseContext.activeDocumentId;
-
-  const contextValue: ActionContextType = {
-    ...baseContext,
-    ...overrides,
-    activeCollectionId,
-    activeDocumentId,
+    history,
+    valueModels,
+    isMenu,
+    isCommandBar,
+    isButton,
+    sidebarContext,
+    event,
+    activeCollectionIdFromStore,
+    activeDocumentIdFromStore,
+    currentUserId,
+    currentTeamId,
     getActiveModels,
     getActiveModel,
     getActivePolicies,
     isModelActive,
-    activeModels: allActiveModels,
-  };
+    allActiveModels,
+  ]);
 
   return (
     <ActionContext.Provider value={contextValue}>
@@ -173,15 +234,20 @@ export default function useActionContext(
 ): ActionContextType {
   const contextValue = useContext(ActionContext);
 
-  // If we have a context value from a provider, use it as the base
-  if (contextValue) {
-    return {
-      ...contextValue,
-      ...overrides,
-    };
+  if (!contextValue) {
+    throw new Error(
+      "useActionContext must be used within an ActionContextProvider"
+    );
   }
 
-  throw new Error(
-    "useActionContext must be used within an ActionContextProvider"
-  );
+  // Short-circuit when no overrides are provided so consumers get a stable
+  // reference and don't re-render unnecessarily.
+  if (!overrides || Object.keys(overrides).length === 0) {
+    return contextValue;
+  }
+
+  return {
+    ...contextValue,
+    ...overrides,
+  };
 }

app/hooks/useOnScreen.ts

@@ -2,34 +2,80 @@ import * as React from "react";
 
 const isSupported = "IntersectionObserver" in window;
 
+// Parses a rootMargin string ("10px 20px" / "10px" / "10px 20px 30px 40px")
+// into [top, right, bottom, left] in pixels. Percentages are not supported in
+// the synchronous fast path and fall back to 0.
+function parseRootMargin(
+  rootMargin: string | undefined
+): [number, number, number, number] {
+  if (!rootMargin) {
+    return [0, 0, 0, 0];
+  }
+  const parts = rootMargin
+    .split(/\s+/)
+    .map((p) => (p.endsWith("px") ? parseFloat(p) : 0));
+  const [t = 0, r = t, b = t, l = r] = parts;
+  return [t, r, b, l];
+}
+
 /**
  * Hook to return if a given ref is visible on screen.
  *
  * @returns boolean if the node is visible
  */
-export default function useOnScreen(ref: React.RefObject<HTMLElement>) {
+export default function useOnScreen(
+  ref: React.RefObject<HTMLElement>,
+  options?: IntersectionObserverInit
+) {
+  const root = options?.root;
+  const rootMargin = options?.rootMargin;
+  const threshold = Array.isArray(options?.threshold)
+    ? options?.threshold.join(",")
+    : options?.threshold;
+
   const [isIntersecting, setIntersecting] = React.useState(!isSupported);
 
-  React.useEffect(() => {
+  React.useLayoutEffect(() => {
     const element = ref.current;
-    let observer: IntersectionObserver | undefined;
-
-    if (isSupported) {
-      observer = new IntersectionObserver(([entry]) => {
-        // Update our state when observer callback fires
-        setIntersecting(entry.isIntersecting);
-      });
+    if (!element) {
+      return undefined;
     }
 
-    if (element) {
-      observer?.observe(element);
+    // Synchronous initial check so the first paint is correct.
+    const [mt, mr, mb, ml] = parseRootMargin(rootMargin);
+    const rect = element.getBoundingClientRect();
+    const rootRect =
+      root instanceof Element
+        ? root.getBoundingClientRect()
+        : {
+            top: 0,
+            left: 0,
+            bottom: window.innerHeight,
+            right: window.innerWidth,
+          };
+    const initialVisible =
+      rect.bottom >= rootRect.top - mt &&
+      rect.top <= rootRect.bottom + mb &&
+      rect.right >= rootRect.left - ml &&
+      rect.left <= rootRect.right + mr;
+
+    setIntersecting(initialVisible);
+
+    if (!isSupported) {
+      return undefined;
     }
+
+    const observer = new IntersectionObserver(([entry]) => {
+      setIntersecting(entry.isIntersecting);
+    }, options);
+    observer.observe(element);
+
     return () => {
-      if (element) {
-        observer?.unobserve(element);
-      }
+      observer.unobserve(element);
     };
-  }, [ref]);
+    // Re-create when option primitives change; options object identity ignored
+    // eslint-disable-next-line react-hooks/exhaustive-deps
+  }, [ref, root, rootMargin, threshold]);
 
   return isIntersecting;
 }