Fix checkbox state when toggling on large directory (#8516)

Author: alexr00

src/view/treeNodes/treeUtils.ts

@@ -4,6 +4,7 @@
  *--------------------------------------------------------------------------------------------*/
 
 import * as vscode from 'vscode';
+import { DirectoryTreeNode } from './directoryTreeNode';
 import { FileChangeNode } from './fileChangeNode';
 import { TreeNode } from './treeNode';
 
@@ -14,16 +15,26 @@ export namespace TreeUtils {
 		const checkedNodes: FileChangeNode[] = [];
 		const uncheckedNodes: FileChangeNode[] = [];
 
+		// The first item is the one the user actually clicked.
+		// Only collect missing descendants if a directory was clicked directly.
+		const firstNode = checkboxUpdates.items[0]?.[0];
+
+		const eventNodes = new Set<TreeNode>(checkboxUpdates.items.map(([node]) => node));
+
 		checkboxUpdates.items.forEach(checkboxUpdate => {
 			const node = checkboxUpdate[0];
 			const newState = checkboxUpdate[1];
 
 			if (node instanceof FileChangeNode) {
-				if (newState == vscode.TreeItemCheckboxState.Checked) {
+				if (newState === vscode.TreeItemCheckboxState.Checked) {
 					checkedNodes.push(node);
 				} else {
 					uncheckedNodes.push(node);
 				}
+			} else if (firstNode instanceof DirectoryTreeNode && node === firstNode) {
+				// VS Code auto-propagates to rendered children, but unrendered children
+				// (due to virtual scrolling) won't be in the event. Collect those missing ones.
+				collectMissingDescendants(firstNode, newState, checkedNodes, uncheckedNodes, eventNodes);
 			}
 
 			node.updateFromCheckboxChanged(newState);
@@ -56,4 +67,32 @@ export namespace TreeUtils {
 			prModel.markFiles(filenames, true, 'unviewed');
 		}
 	}
+
+	/**
+	 * Collect descendant FileChangeNodes that are NOT already in the event.
+	 * These are children VS Code missed because they weren't rendered (virtual scrolling).
+	 */
+	function collectMissingDescendants(
+		dirNode: DirectoryTreeNode,
+		newState: vscode.TreeItemCheckboxState,
+		checkedNodes: FileChangeNode[],
+		uncheckedNodes: FileChangeNode[],
+		eventNodes: Set<TreeNode>
+	): void {
+		for (const child of dirNode._children) {
+			if (eventNodes.has(child)) {
+				continue;
+			}
+			if (child instanceof FileChangeNode) {
+				if (newState === vscode.TreeItemCheckboxState.Checked) {
+					checkedNodes.push(child);
+				} else {
+					uncheckedNodes.push(child);
+				}
+				child.updateFromCheckboxChanged(newState);
+			} else if (child instanceof DirectoryTreeNode) {
+				collectMissingDescendants(child, newState, checkedNodes, uncheckedNodes, eventNodes);
+			}
+		}
+	}
 }