From c3687e0fb8df78e1b66ab6387e1230a89e349c35 Mon Sep 17 00:00:00 2001
From: Sulaiman Hamouda <sulai@buildwithfern.com>
Date: Mon, 29 Jul 2024 15:53:28 -0400
Subject: [PATCH] Initialize collapsed/expanded sidebar state within
 CollapseSidebarContext

---
 .../src/sidebar/CollapseSidebarContext.tsx    | 32 +++++++++++++------
 .../src/sidebar/nodes/SidebarSectionNode.tsx  | 13 ++------
 2 files changed, 25 insertions(+), 20 deletions(-)

diff --git a/packages/ui/app/src/sidebar/CollapseSidebarContext.tsx b/packages/ui/app/src/sidebar/CollapseSidebarContext.tsx
index 3b9e42d24a..4b3e206dd1 100644
--- a/packages/ui/app/src/sidebar/CollapseSidebarContext.tsx
+++ b/packages/ui/app/src/sidebar/CollapseSidebarContext.tsx
@@ -7,9 +7,8 @@ import {
     createContext,
     useCallback,
     useContext,
-    useEffect,
     useMemo,
-    useState,
+    useState
 } from "react";
 import { noop } from "ts-essentials";
 import { useCallbackOne } from "use-memo-one";
@@ -84,11 +83,10 @@ export const CollapseSidebarProvider: FC<
     //         }),
     //     [registerListener, scrollContainerRef],
     // );
-
+    
     const { parentIdMap, parentToChildrenMap } = useMemo(() => {
         const parentIdMap = new Map<FernNavigation.NodeId, FernNavigation.NodeId[]>();
         const parentToChildrenMap = new Map<FernNavigation.NodeId, FernNavigation.NodeId[]>();
-
         if (sidebar == null) {
             return { parentIdMap, parentToChildrenMap };
         }
@@ -113,14 +111,28 @@ export const CollapseSidebarProvider: FC<
         return { parentIdMap, parentToChildrenMap };
     }, [sidebar]);
 
-    const [expanded, setExpanded] = useState<FernNavigation.NodeId[]>(() =>
-        selectedNodeId == null ? [] : [selectedNodeId, ...(parentIdMap.get(selectedNodeId) ?? [])],
+    const initializeExpandedSections = (): FernNavigation.NodeId[] => {
+        if(selectedNodeId == null) {
+            return []
+        } else {
+        if(sidebar) {
+            const selectedNodes: FernNavigation.NodeId[] = []
+            FernNavigation.utils.traverseNavigation(sidebar, (node, _index, parents) => {
+                if(FernNavigation.isSection(node) && node.type === "section") {
+                    if(!node?.collapsed) {
+                        selectedNodes.push(...[node.id, ...(parentIdMap.get(node.id) ?? [])])
+                    }
+                }
+            })
+            return selectedNodes
+        } else {
+            return [selectedNodeId, ...(parentIdMap.get(selectedNodeId) ?? [])]
+        }}
+    }
+
+    const [expanded, setExpanded] = useState<FernNavigation.NodeId[]>(() => initializeExpandedSections()
     );
 
-    useEffect(() => {
-        setExpanded(selectedNodeId == null ? [] : [selectedNodeId, ...(parentIdMap.get(selectedNodeId) ?? [])]);
-    }, [selectedNodeId, parentIdMap]);
-
     const checkExpanded = useCallback(
         (expandableId: FernNavigation.NodeId) =>
             expanded.some((id) => id === expandableId || parentIdMap.get(id)?.includes(expandableId)),
diff --git a/packages/ui/app/src/sidebar/nodes/SidebarSectionNode.tsx b/packages/ui/app/src/sidebar/nodes/SidebarSectionNode.tsx
index db41f8d965..f557feba49 100644
--- a/packages/ui/app/src/sidebar/nodes/SidebarSectionNode.tsx
+++ b/packages/ui/app/src/sidebar/nodes/SidebarSectionNode.tsx
@@ -1,6 +1,6 @@
 import { FernNavigation } from "@fern-api/fdr-sdk";
 import clsx from "clsx";
-import { useCallback, useEffect, useState } from "react";
+import { useCallback } from "react";
 import { useCurrentNodeId } from "../../atoms";
 import { useCollapseSidebar } from "../CollapseSidebarContext";
 import { SidebarSlugLink } from "../SidebarLink";
@@ -14,12 +14,8 @@ interface SidebarSectionNodeProps {
 }
 
 export function SidebarSectionNode({ node, className, depth }: SidebarSectionNodeProps): React.ReactElement | null {
-    const [expanded, setExpanded] = useState<boolean>(false)
     const { checkExpanded, toggleExpanded, checkChildSelected, registerScrolledToPathListener } = useCollapseSidebar();
-    const handleToggleExpand = useCallback(() => {
-        toggleExpanded(node.id)
-        setExpanded(!expanded)
-    }, [node.id, toggleExpanded, expanded]);
+    const handleToggleExpand = useCallback(() => toggleExpanded(node.id), [node.id, toggleExpanded]);
     const selectedNodeId = useCurrentNodeId();
 
     if (node.children.length === 0) {
@@ -45,10 +41,7 @@ export function SidebarSectionNode({ node, className, depth }: SidebarSectionNod
         return null;
     }
 
-    useEffect(() => {
-        setExpanded(!!node.id || checkExpanded(node.id) || (childSelected && node.overviewPageId != null) || !node?.collapsed)
-    }, [])
-
+    const expanded = selectedNodeId === node.id || checkExpanded(node.id) || (childSelected && node.overviewPageId != null)
     const showIndicator = childSelected && !expanded;
 
     return (