From 5b03801304200f12178cf2b6b2c39a5f4fd532fb Mon Sep 17 00:00:00 2001
From: Denise Li <dli@squareup.com>
Date: Mon, 7 Oct 2024 15:11:34 -0400
Subject: [PATCH] fix: module tree expansion state when closing the module of
 the selected decl (#3022)

Fixes: https://github.com/TBD54566975/ftl/issues/2982

This fix delves into the nuance of component re-rendering lifecycle a
bit. Basically, before, we would add the selected module to local
storage AND update the component state on every single re-render of the
component (at `const [expandedModules, setExpandedModules] =
useState(...)`. The problem with that was that when we toggled the
selected module to close it, the re-render would cause the selected
module to still be added back to local storage, but the component state
would be out of sync. Then when you went to expand that module again,
because the states got unsynced, it would stay stuck in the
visibly-closed state. Now, we only add the selected module to local
storage in the effect, which doesn't get triggered on toggle a module's
expansion state. As a result, the state doesn't de-sync.

So now, you can do the following:
1. Select a decl
2. Close the module of the selected decl in the module tree
3. Re-expand that same module in the module tree. It works.

While the following still works correctly:
1. Close all the modules (i.e. to-be-selected module will not be cached
as expanded)
2. Go to a decl page (e.g. paste a link into the browser)
3. The selected module should be expanded in the tree.
---
 .../src/features/modules/ModulesTree.tsx      | 24 +++++++++----------
 1 file changed, 12 insertions(+), 12 deletions(-)

diff --git a/frontend/console/src/features/modules/ModulesTree.tsx b/frontend/console/src/features/modules/ModulesTree.tsx
index ba41fbab9..b3e764612 100644
--- a/frontend/console/src/features/modules/ModulesTree.tsx
+++ b/frontend/console/src/features/modules/ModulesTree.tsx
@@ -127,13 +127,6 @@ const ModuleSection = ({
   )
 }
 
-function getInitialExpandedModules(moduleName?: string, declName?: string): string[] {
-  if (moduleName && declName) {
-    addModuleToLocalStorageIfMissing(moduleName)
-  }
-  return listExpandedModulesFromLocalStorage()
-}
-
 const declTypesSearchParamKey = 'dt'
 
 export const ModulesTree = ({ modules }: { modules: ModuleTreeItem[] }) => {
@@ -144,9 +137,13 @@ export const ModulesTree = ({ modules }: { modules: ModuleTreeItem[] }) => {
   const declTypesFromUrl = declTypeMultiselectOpts.filter((o) => declTypeKeysFromUrl.includes(o.key))
   const [selectedDeclTypes, setSelectedDeclTypes] = useState(declTypesFromUrl.length === 0 ? declTypeMultiselectOpts : declTypesFromUrl)
 
-  const [expandedModules, setExpandedModules] = useState(getInitialExpandedModules(moduleName, declName))
+  const initialExpanded = listExpandedModulesFromLocalStorage()
+  const [expandedModules, setExpandedModules] = useState(initialExpanded)
   useEffect(() => {
-    setExpandedModules(getInitialExpandedModules(moduleName, declName))
+    if (moduleName && declName) {
+      addModuleToLocalStorageIfMissing(moduleName)
+    }
+    setExpandedModules(listExpandedModulesFromLocalStorage())
   }, [moduleName, declName])
 
   function msOnChange(opts: MultiselectOpt[]) {
@@ -160,14 +157,17 @@ export const ModulesTree = ({ modules }: { modules: ModuleTreeItem[] }) => {
     setSelectedDeclTypes(opts)
   }
 
-  function toggle(moduleName: string) {
-    toggleModuleExpansionInLocalStorage(moduleName)
+  function toggle(toggledModule: string) {
+    toggleModuleExpansionInLocalStorage(toggledModule)
     setExpandedModules(listExpandedModulesFromLocalStorage())
   }
 
   function collapseAll() {
     collapseAllModulesInLocalStorage()
-    setExpandedModules([])
+    if (moduleName && declName) {
+      addModuleToLocalStorageIfMissing(moduleName)
+    }
+    setExpandedModules(listExpandedModulesFromLocalStorage())
   }
 
   modules.sort((m1, m2) => Number(m1.isBuiltin) - Number(m2.isBuiltin))