Address PR review comments: improve checkbox reset logic

- Only set timestamp when checkbox is checked (not unchecked)
- Update comment to clarify timestamp behavior
- Use >= for file mtime comparison to handle race conditions
- Fix cleanup logic to not delete folder checkbox states
- Consolidate tree refresh to avoid redundant refreshes
- Add caching for computeFolderCheckboxState performance

Co-authored-by: letmaik <530988+letmaik@users.noreply.github.com>

Author: Copilot

src/treeProvider.ts

@@ -19,7 +19,7 @@ import { API as GitAPI, Repository as GitAPIRepository } from './typings/git';
 
 interface CheckboxStateInfo {
     state: TreeItemCheckboxState;
-    timestamp: number; // When the checkbox was checked
+    timestamp: number; // Timestamp when the checkbox was checked (not set for unchecked states)
 }
 
 class FileElement implements IDiffStatus {
@@ -140,6 +140,7 @@ export class GitTreeCompareProvider implements TreeDataProvider<Element>, Dispos
     private treeView: TreeView<Element>;
     private isPaused: boolean;
     private checkboxStates: Map<string, CheckboxStateInfo> = new Map<string, CheckboxStateInfo>();
+    private folderCheckboxStateCache: Map<string, TreeItemCheckboxState> = new Map<string, TreeItemCheckboxState>();
 
     // Other
     private readonly disposables: Disposable[] = [];
@@ -311,12 +312,33 @@ export class GitTreeCompareProvider implements TreeDataProvider<Element>, Dispos
     }
 
     private async handleChangeCheckboxState(e: TreeCheckboxChangeEvent<Element>) {
+        // Clear the folder checkbox state cache since states are changing
+        this.folderCheckboxStateCache.clear();
+        
         for (let [element, state] of e.items) {
             if (element instanceof FileElement || element instanceof FolderElement) {
-                this.checkboxStates.set(element.dstAbsPath, {
-                    state: state,
-                    timestamp: Date.now()
-                });
+                const key = element.dstAbsPath;
+                const existing = this.checkboxStates.get(key);
+
+                if (state === TreeItemCheckboxState.Checked) {
+                    // Only set timestamp when checking the item
+                    this.checkboxStates.set(key, {
+                        state: state,
+                        timestamp: Date.now()
+                    });
+                } else if (existing) {
+                    // When unchecking, preserve the timestamp of when it was last checked
+                    this.checkboxStates.set(key, {
+                        state: state,
+                        timestamp: existing.timestamp
+                    });
+                } else {
+                    // No existing state, just set unchecked with current timestamp
+                    this.checkboxStates.set(key, {
+                        state: state,
+                        timestamp: Date.now()
+                    });
+                }
             }
         }
     }
@@ -403,10 +425,18 @@ export class GitTreeCompareProvider implements TreeDataProvider<Element>, Dispos
     }
 
     private computeFolderCheckboxState(folder: FolderElement): TreeItemCheckboxState {
+        // Check cache first
+        const cached = this.folderCheckboxStateCache.get(folder.dstAbsPath);
+        if (cached !== undefined) {
+            return cached;
+        }
+        
         // Check if user explicitly set state on this folder
         const explicitState = this.checkboxStates.get(folder.dstAbsPath);
         if (explicitState) {
-            return explicitState.state;
+            const result = explicitState.state;
+            this.folderCheckboxStateCache.set(folder.dstAbsPath, result);
+            return result;
         }
 
         // Otherwise derive from files: folder is checked only if ALL files under it are checked
@@ -429,7 +459,9 @@ export class GitTreeCompareProvider implements TreeDataProvider<Element>, Dispos
             }
         }
 
-        return (hasFiles && allChecked) ? TreeItemCheckboxState.Checked : TreeItemCheckboxState.Unchecked;
+        const result = (hasFiles && allChecked) ? TreeItemCheckboxState.Checked : TreeItemCheckboxState.Unchecked;
+        this.folderCheckboxStateCache.set(folder.dstAbsPath, result);
+        return result;
     }
 
     async getChildren(element?: Element): Promise<Element[]> {
@@ -555,8 +587,10 @@ export class GitTreeCompareProvider implements TreeDataProvider<Element>, Dispos
         this.log(`${diff.length} diff entries (${untrackedCount} untracked)`);
 
         const newFilePaths = new Set<string>();
+        const newFolderPaths = new Set<string>();
         // Collect files that need mtime checking for async batch processing
         const filesToCheckMtime: Array<{filePath: string, stateInfo: CheckboxStateInfo}> = [];
+        let checkboxesWereReset = false;
 
         for (const entry of diff) {
             const folder = path.dirname(entry.dstAbsPath);
@@ -575,6 +609,7 @@ export class GitTreeCompareProvider implements TreeDataProvider<Element>, Dispos
                 if (!files.has(currentFolder)) {
                     files.set(currentFolder, new Array());
                 }
+                newFolderPaths.add(currentFolder);
                 currentFolder = path.dirname(currentFolder)
             }
 
@@ -600,8 +635,8 @@ export class GitTreeCompareProvider implements TreeDataProvider<Element>, Dispos
                     const stats = await fs.promises.stat(filePath);
                     const fileMtime = stats.mtimeMs;
 
-                    // If file was modified after checkbox was checked, reset it
-                    if (fileMtime > stateInfo.timestamp) {
+                    // If file was modified at or after the checkbox was checked, reset it
+                    if (fileMtime >= stateInfo.timestamp) {
                         return filePath;
                     }
                 } catch (error: unknown) {
@@ -614,23 +649,30 @@ export class GitTreeCompareProvider implements TreeDataProvider<Element>, Dispos
 
             const pathsToReset = await Promise.all(statPromises);
             const actualPathsToReset = pathsToReset.filter((filePath): filePath is string => filePath !== null);
-            actualPathsToReset.forEach(filePath => this.checkboxStates.delete(filePath));
 
-            // Fire tree refresh to update checkbox UI
             if (actualPathsToReset.length > 0) {
-                this._onDidChangeTreeData.fire();
+                // Clear folder checkbox state cache before deleting checkbox states
+                this.folderCheckboxStateCache.clear();
+                actualPathsToReset.forEach(filePath => this.checkboxStates.delete(filePath));
             }
+            
+            // Track if checkboxes were reset to consolidate tree refresh later
+            checkboxesWereReset = actualPathsToReset.length > 0;
         }
 
-        // Clear checkbox state for files that no longer exist in the diff
+        // Clear checkbox state for files and folders that no longer exist in the diff
         const pathsToDelete: string[] = [];
-        for (const [filePath] of this.checkboxStates) {
-            if (!newFilePaths.has(filePath)) {
-                pathsToDelete.push(filePath);
+        for (const [itemPath] of this.checkboxStates) {
+            // Only delete if it's not in files AND not in folders
+            if (!newFilePaths.has(itemPath) && !newFolderPaths.has(itemPath)) {
+                pathsToDelete.push(itemPath);
             }
         }
-        for (const filePath of pathsToDelete) {
-            this.checkboxStates.delete(filePath);
+        if (pathsToDelete.length > 0) {
+            this.folderCheckboxStateCache.clear();
+            for (const itemPath of pathsToDelete) {
+                this.checkboxStates.delete(itemPath);
+            }
         }
 
         let treeHasChanged = false;
@@ -679,7 +721,7 @@ export class GitTreeCompareProvider implements TreeDataProvider<Element>, Dispos
         this.filesInsideTreeRoot = filesInsideTreeRoot;
         this.filesOutsideTreeRoot = filesOutsideTreeRoot;
 
-        if (fireChangeEvents && treeHasChanged) {
+        if (fireChangeEvents && (treeHasChanged || checkboxesWereReset)) {
             this.log('Refreshing tree')
             this._onDidChangeTreeData.fire();
         }