From 351cb16b8f60c194a764f03a203d39d792439bdf Mon Sep 17 00:00:00 2001
From: EnixCoda <enixcoda@gmail.com>
Date: Tue, 3 Aug 2021 23:02:25 +0800
Subject: [PATCH] feat: restore expanded folders

---
 src/components/FileExplorer.tsx              |  6 ++--
 src/components/settings/FileTreeSettings.tsx |  7 +++++
 src/utils/VisibleNodesGenerator.ts           | 33 +++++++++++++-------
 src/utils/config/helper.ts                   |  4 ++-
 4 files changed, 35 insertions(+), 15 deletions(-)

diff --git a/src/components/FileExplorer.tsx b/src/components/FileExplorer.tsx
index 8abcc59..2e1b8b8 100644
--- a/src/components/FileExplorer.tsx
+++ b/src/components/FileExplorer.tsx
@@ -46,17 +46,17 @@ const RawFileExplorer: React.FC<Props & ConnectorState> = function RawFileExplor
     searched,
   } = props
   const {
-    value: { accessToken, compressSingletonFolder, searchMode, commentToggle },
+    value: { accessToken, compressSingletonFolder, searchMode, commentToggle, restoreExpandedFolders },
   } = useConfigs()
 
   const onSearch = React.useCallback(
     (searchKey: string, searchMode: SearchMode) => {
       updateSearchKey(searchKey)
       if (visibleNodesGenerator) {
-        visibleNodesGenerator.search(searchModes[searchMode].getSearchParams(searchKey))
+        visibleNodesGenerator.search(searchModes[searchMode].getSearchParams(searchKey), restoreExpandedFolders)
       }
     },
-    [updateSearchKey, visibleNodesGenerator],
+    [updateSearchKey, visibleNodesGenerator, restoreExpandedFolders],
   )
 
   const stateContext = useLoadedContext(SideBarStateContext)
diff --git a/src/components/settings/FileTreeSettings.tsx b/src/components/settings/FileTreeSettings.tsx
index 337ab54..3b79af5 100644
--- a/src/components/settings/FileTreeSettings.tsx
+++ b/src/components/settings/FileTreeSettings.tsx
@@ -76,6 +76,13 @@ export function FileTreeSettings(props: React.PropsWithChildren<Props>) {
           tooltip: 'Merge folders and their only child folder to make UI more compact.',
         }}
       />
+      <SimpleToggleField
+        field={{
+          key: 'restoreExpandedFolders',
+          label: 'Restore expanded folders',
+          tooltip: 'Folders will be expanded again when clear search input',
+        }}
+      />
       <SimpleToggleField
         field={{
           key: 'commentToggle',
diff --git a/src/utils/VisibleNodesGenerator.ts b/src/utils/VisibleNodesGenerator.ts
index c14c234..9ae231f 100644
--- a/src/utils/VisibleNodesGenerator.ts
+++ b/src/utils/VisibleNodesGenerator.ts
@@ -124,6 +124,10 @@ class ShakeLayer extends BaseLayer {
   lastSearchParams: SearchParams | null = null
   shakeHub = new EventHub<{ emit: TreeNode | null }>()
 
+  get isSearching() {
+    return this.lastSearchParams !== null
+  }
+
   constructor(options: Options) {
     super(options)
 
@@ -193,6 +197,7 @@ class FlattenLayer extends CompressLayer {
   focusedNode: TreeNode | null = null
   nodes: TreeNode[] = []
   expandedNodes: Set<TreeNode['path']> = new Set()
+  backupExpandedNodes: Set<TreeNode['path']> = new Set()
   flattenHub = new EventHub<{ emit: null }>()
 
   constructor(options: Options) {
@@ -300,18 +305,24 @@ class FlattenLayer extends CompressLayer {
     }
   }, this.generateVisibleNodes)
 
-  search = (searchParams: Pick<SearchParams, 'matchNode'> | null) => {
-    this.shake(
-      searchParams
-        ? {
-            matchNode: searchParams.matchNode,
-            onChildMatch: node => this.$setExpand(node, true),
-          }
-        : null,
-    )
-    // collapse all nodes on clearing search key
-    if (!searchParams) {
+  search = (searchParams: Pick<SearchParams, 'matchNode'> | null, restoreExpandedFolders?: boolean) => {
+    // backup expansion before search
+    if (searchParams) {
+      if (!this.isSearching) {
+        this.backupExpandedNodes.clear()
+        this.expandedNodes.forEach(path => this.backupExpandedNodes.add(path))
+      }
+      this.shake({
+        matchNode: searchParams.matchNode,
+        onChildMatch: node => this.$setExpand(node, true),
+      })
+    } else {
+      this.shake(null)
+      // collapse all nodes on clearing search key
       this.expandedNodes.clear()
+      if (restoreExpandedFolders) {
+        this.backupExpandedNodes.forEach(path => this.expandedNodes.add(path))
+      }
     }
   }
 }
diff --git a/src/utils/config/helper.ts b/src/utils/config/helper.ts
index 9a2b562..001b1be 100644
--- a/src/utils/config/helper.ts
+++ b/src/utils/config/helper.ts
@@ -1,5 +1,4 @@
 import { SearchMode } from 'components/searchModes'
-import { platform, platformName } from 'platforms'
 import { storageHelper } from 'utils/storageHelper'
 import { migrateConfig } from './migrations'
 
@@ -20,6 +19,7 @@ export type Config = {
   commentToggle: boolean
   codeFolding: boolean
   compactFileTree: boolean
+  restoreExpandedFolders: boolean
 }
 
 enum configKeys {
@@ -39,6 +39,7 @@ enum configKeys {
   commentToggle = 'commentToggle',
   codeFolding = 'codeFolding',
   compactFileTree = 'compactFileTree',
+  restoreExpandedFolders = 'restoreExpandedFolders',
 }
 
 // do NOT use platform name
@@ -61,6 +62,7 @@ export const defaultConfigs: Config = {
   commentToggle: true,
   codeFolding: true,
   compactFileTree: false,
+  restoreExpandedFolders: true,
 }
 
 const configKeyArray = Object.values(configKeys)