feat(ui): decouple tree nodes from the UI representation (#2783)

## Problem
While prototyping features, I kept running into the situation where I wanted to update how a node was rendered independently of the node's children. The current implementation does not easily support this use-case, forcing us to propagate events up to a parent node to get the desired behavior.

## Solution
Fully generalize the TreeNode interface. It was already very close to the general solution, it just needed a few tweaks to get it there. Of course, implementing this correctly with a tree data provider required more substantial changes which tests were added for.

I wouldn't recommend that all tree node implementations use this pattern because it can quickly spiral out of control from a maintainability point of view. Still, I'd rather have the proper solution readily available than a bunch of one-off bespoke solutions.

Author: JadenSimon

src/awsexplorer/localExplorer.ts

@@ -5,10 +5,9 @@
 
 import * as vscode from 'vscode'
 import * as telemetry from '../shared/telemetry/telemetry'
-import { cdkNode } from '../cdk/explorer/rootNode'
+import { cdkNode, CdkRootNode } from '../cdk/explorer/rootNode'
 import { ResourceTreeDataProvider, TreeNode } from '../shared/treeview/resourceTreeDataProvider'
 import { once } from '../shared/utilities/functionUtils'
-import { AppNode } from '../cdk/explorer/nodes/appNode'
 import { isCloud9 } from '../shared/extensionUtilities'
 import { codewhispererNode } from '../codewhisperer/explorer/codewhispererNode'
 
@@ -49,7 +48,7 @@ export function createLocalExplorerView(): vscode.TreeView<TreeNode> {
     // Legacy CDK metric, remove this when we add something generic
     const recordExpandCdkOnce = once(telemetry.recordCdkAppExpanded)
     view.onDidExpandElement(e => {
-        if (e.element.resource instanceof AppNode) {
+        if (e.element.resource instanceof CdkRootNode) {
             recordExpandCdkOnce()
         }
     })

src/cdk/explorer/nodes/appNode.ts

@@ -21,11 +21,9 @@ import { getIcon } from '../../../shared/icons'
 export class AppNode implements TreeNode {
     public readonly id = this.location.cdkJsonUri.toString()
     public readonly resource = this.location
-    public readonly treeItem: vscode.TreeItem
+    public readonly label = vscode.workspace.asRelativePath(vscode.Uri.joinPath(this.location.cdkJsonUri, '..'))
 
-    public constructor(private readonly location: CdkAppLocation) {
-        this.treeItem = this.createTreeItem()
-    }
+    public constructor(private readonly location: CdkAppLocation) {}
 
     public async getChildren(): Promise<(ConstructNode | TreeNode)[]> {
         const constructs = []
@@ -59,9 +57,8 @@ export class AppNode implements TreeNode {
         }
     }
 
-    private createTreeItem() {
-        const label = vscode.workspace.asRelativePath(vscode.Uri.joinPath(this.location.cdkJsonUri, '..'))
-        const item = new vscode.TreeItem(label, vscode.TreeItemCollapsibleState.Collapsed)
+    public getTreeItem() {
+        const item = new vscode.TreeItem(this.label, vscode.TreeItemCollapsibleState.Collapsed)
 
         item.contextValue = 'awsCdkAppNode'
         item.iconPath = getIcon('aws-cdk-logo')

src/cdk/explorer/nodes/constructNode.ts

@@ -13,13 +13,10 @@ import { generatePropertyNodes } from './propertyNode'
 
 export class ConstructNode implements TreeNode {
     public readonly id = this.construct.id
-    public readonly treeItem: vscode.TreeItem
     private readonly type = treeInspector.getTypeAttributeOrDefault(this.construct, '')
     private readonly properties = treeInspector.getProperties(this.construct)
 
-    public constructor(private readonly location: CdkAppLocation, private readonly construct: ConstructTreeEntity) {
-        this.treeItem = this.createTreeItem()
-    }
+    public constructor(private readonly location: CdkAppLocation, private readonly construct: ConstructTreeEntity) {}
 
     public get resource() {
         return {
@@ -36,7 +33,7 @@ export class ConstructNode implements TreeNode {
         return [...propertyNodes, ...constructNodes]
     }
 
-    private createTreeItem() {
+    public getTreeItem() {
         const collapsibleState =
             this.construct.children || this.construct.attributes
                 ? vscode.TreeItemCollapsibleState.Collapsed

src/cdk/explorer/nodes/propertyNode.ts

@@ -14,11 +14,8 @@ import { TreeNode } from '../../../shared/treeview/resourceTreeDataProvider'
 export class PropertyNode implements TreeNode {
     public readonly id = this.key
     public readonly resource = this.value
-    public readonly treeItem: vscode.TreeItem
 
-    public constructor(private readonly key: string, private readonly value: unknown) {
-        this.treeItem = this.createTreeItem()
-    }
+    public constructor(private readonly key: string, private readonly value: unknown) {}
 
     public async getChildren(): Promise<TreeNode[]> {
         if (this.value instanceof Array || this.value instanceof Object) {
@@ -28,7 +25,7 @@ export class PropertyNode implements TreeNode {
         }
     }
 
-    private createTreeItem() {
+    public getTreeItem() {
         const item = new vscode.TreeItem(`${this.key}: ${this.value}`)
 
         item.contextValue = 'awsCdkPropertyNode'

src/cdk/explorer/rootNode.ts

@@ -20,25 +20,24 @@ export async function getAppNodes(): Promise<TreeNode[]> {
         return [createPlaceholderItem(localize('AWS.cdk.explorerNode.noApps', '[No CDK Apps found in Workspaces]'))]
     }
 
-    return appsFound.map(appLocation => new AppNode(appLocation))
+    return appsFound.map(appLocation => new AppNode(appLocation)).sort((a, b) => a.label.localeCompare(b.label) ?? 0)
 }
 
 export class CdkRootNode implements TreeNode {
     public readonly id = 'cdk'
-    public readonly treeItem = this.createTreeItem()
     public readonly resource = this
     private readonly onDidChangeChildrenEmitter = new vscode.EventEmitter<void>()
     public readonly onDidChangeChildren = this.onDidChangeChildrenEmitter.event
 
-    public async getChildren() {
-        return (await getAppNodes()).sort((a, b) => a.treeItem?.label?.localeCompare(b.treeItem?.label ?? '') ?? 0)
+    public getChildren() {
+        return getAppNodes()
     }
 
     public refresh(): void {
         this.onDidChangeChildrenEmitter.fire()
     }
 
-    private createTreeItem() {
+    public getTreeItem() {
         const item = new vscode.TreeItem('CDK')
         item.collapsibleState = vscode.TreeItemCollapsibleState.Collapsed
         item.contextValue = 'awsCdkRootNode'

src/codewhisperer/explorer/codewhispererNode.ts

@@ -24,7 +24,6 @@ import { isCloud9 } from '../../shared/extensionUtilities'
 import { Cloud9AccessState } from '../models/model'
 export class CodeWhispererNode implements RootNode {
     public readonly id = 'codewhisperer'
-    public readonly treeItem = this.createTreeItem()
     public readonly resource = this
     private readonly onDidChangeChildrenEmitter = new vscode.EventEmitter<void>()
     public readonly onDidChangeChildren = this.onDidChangeChildrenEmitter.event
@@ -45,7 +44,7 @@ export class CodeWhispererNode implements RootNode {
         })
     }
 
-    private createTreeItem() {
+    public getTreeItem() {
         const item = new vscode.TreeItem('CodeWhisperer (Preview)')
         item.collapsibleState = vscode.TreeItemCollapsibleState.Collapsed
         item.contextValue = 'awsCodeWhispererNode'

src/shared/treeview/resource.ts

@@ -24,7 +24,6 @@ export interface ResourceProvider<T extends Resource = Resource> {
 
 export class ResourceTreeNode<T extends Resource> implements TreeNode<T> {
     public readonly id = this.resource.id
-    public readonly treeItem = this.createTreeItem()
 
     public constructor(
         public readonly resource: T,
@@ -40,7 +39,7 @@ export class ResourceTreeNode<T extends Resource> implements TreeNode<T> {
         return this.children?.listResources() ?? []
     }
 
-    private createTreeItem(): vscode.TreeItem {
+    public getTreeItem(): vscode.TreeItem {
         const collapsed =
             this.children !== undefined
                 ? vscode.TreeItemCollapsibleState.Collapsed

src/shared/treeview/resourceTreeDataProvider.ts

@@ -25,14 +25,19 @@ export interface TreeNode<T = unknown> {
     readonly resource: T
 
     /**
-     * A tree item used to display the node in a tree view.
+     * An optional event to signal that this node's children has changed.
      */
-    readonly treeItem: vscode.TreeItem
+    readonly onDidChangeChildren?: vscode.Event<void>
 
     /**
-     * Optional event to signal that this node's children has changed.
+     * An optional event to signal that this node's tree item has changed.
      */
-    readonly onDidChangeChildren?: vscode.Event<void>
+    readonly onDidChangeTreeItem?: vscode.Event<void>
+
+    /**
+     * Returns a tree item used to display the node in a tree view.
+     */
+    getTreeItem(): Promise<vscode.TreeItem> | vscode.TreeItem
 
     /**
      * Optional method to provide child nodes.
@@ -46,38 +51,52 @@ export function isTreeNode(obj: unknown): obj is TreeNode {
         typeof obj === 'object' &&
         'resource' in obj &&
         typeof (obj as TreeNode).id === 'string' &&
-        (obj as TreeNode).treeItem instanceof vscode.TreeItem
+        typeof (obj as TreeNode).getTreeItem === 'function'
     )
 }
 
 function copyNode<T>(id: string, node: Omit<TreeNode<T>, 'id'>): TreeNode<T> {
     return {
         id,
         resource: node.resource,
-        treeItem: node.treeItem,
         onDidChangeChildren: node.onDidChangeChildren,
+        onDidChangeTreeItem: node.onDidChangeTreeItem,
+        getTreeItem: node.getTreeItem?.bind(node),
         getChildren: node.getChildren?.bind(node),
     }
 }
 
 export class ResourceTreeDataProvider implements vscode.TreeDataProvider<TreeNode> {
-    private readonly children: Map<string, TreeNode[]> = new Map()
-    private readonly listeners: Map<string, vscode.Disposable> = new Map()
+    private readonly items = new Map<string, vscode.TreeItem>()
+    private readonly children = new Map<string, TreeNode[]>()
+    private readonly listeners = new Map<string, vscode.Disposable>()
     private readonly onDidChangeTreeDataEmitter = new vscode.EventEmitter<TreeNode | void>()
     public readonly onDidChangeTreeData = this.onDidChangeTreeDataEmitter.event
 
     public constructor(private readonly root: Required<Pick<TreeNode, 'getChildren'>>) {}
 
-    public getTreeItem(element: TreeNode): vscode.TreeItem {
-        const item = element.treeItem
+    public async getTreeItem(element: TreeNode): Promise<vscode.TreeItem> {
+        const previousItem = this.items.get(element.id)
+        if (previousItem) {
+            return previousItem
+        }
+
+        const item = await element.getTreeItem()
         item.id = element.id
+        this.items.set(element.id, item)
 
         return item
     }
 
     public async getChildren(element?: TreeNode): Promise<TreeNode[]> {
         if (element) {
-            this.children.get(element.id)?.forEach(n => this.clear(n))
+            const previousChildren = this.children.get(element.id)
+
+            if (previousChildren !== undefined) {
+                return previousChildren
+            } else {
+                this.children.get(element.id)?.forEach(n => this.clear(n))
+            }
         }
 
         const getId = (id: string) => (element ? `${element.id}/${id}` : id)
@@ -91,6 +110,7 @@ export class ResourceTreeDataProvider implements vscode.TreeDataProvider<TreeNod
     public refresh(): void {
         vscode.Disposable.from(...this.listeners.values()).dispose()
 
+        this.items.clear()
         this.children.clear()
         this.listeners.clear()
         this.onDidChangeTreeDataEmitter.fire()
@@ -99,6 +119,7 @@ export class ResourceTreeDataProvider implements vscode.TreeDataProvider<TreeNod
     private clear(node: TreeNode): void {
         const children = this.children.get(node.id)
 
+        this.items.delete(node.id)
         this.children.delete(node.id)
         this.listeners.get(node.id)?.dispose()
         this.listeners.delete(node.id)
@@ -108,14 +129,29 @@ export class ResourceTreeDataProvider implements vscode.TreeDataProvider<TreeNod
 
     private insert(id: string, resource: TreeNode): TreeNode {
         const node = copyNode(id, resource)
+        const listeners: vscode.Disposable[] = []
 
         if (node.onDidChangeChildren) {
-            const listener = node.onDidChangeChildren?.(() => {
-                this.children.get(node.id)?.forEach(n => this.clear(n))
-                this.onDidChangeTreeDataEmitter.fire(node)
-            })
+            listeners.push(
+                node.onDidChangeChildren?.(() => {
+                    this.children.get(node.id)?.forEach(n => this.clear(n))
+                    this.children.delete(node.id)
+                    this.onDidChangeTreeDataEmitter.fire(node)
+                })
+            )
+        }
+
+        if (node.onDidChangeTreeItem) {
+            listeners.push(
+                node.onDidChangeTreeItem?.(() => {
+                    this.items.delete(node.id)
+                    this.onDidChangeTreeDataEmitter.fire(node)
+                })
+            )
+        }
 
-            this.listeners.set(id, listener)
+        if (listeners.length !== 0) {
+            this.listeners.set(id, vscode.Disposable.from(...listeners))
         }
 
         return node

src/shared/treeview/utils.ts

@@ -77,7 +77,7 @@ export function createPlaceholderItem(message: string): TreeNode {
     return {
         id: 'placeholder',
         resource: message,
-        treeItem: new vscode.TreeItem(message, vscode.TreeItemCollapsibleState.None),
+        getTreeItem: () => new vscode.TreeItem(message, vscode.TreeItemCollapsibleState.None),
     }
 }
 
@@ -96,16 +96,43 @@ export function unboxTreeNode<T>(node: TreeNode, predicate: (resource: unknown)
  * would be passed in as-is.
  */
 export class TreeShim extends AWSTreeNodeBase {
-    public constructor(public readonly node: TreeNode) {
-        super(node.treeItem.label ?? '[No label]')
-        assign(node.treeItem, this)
+    private children?: AWSTreeNodeBase[]
 
-        this.node.onDidChangeChildren?.(() => this.refresh())
+    public constructor(public readonly node: TreeNode) {
+        super('Loading...')
+        this.updateTreeItem()
+
+        this.node.onDidChangeChildren?.(() => {
+            this.children = undefined
+            this.refresh()
+        })
+
+        this.node.onDidChangeTreeItem?.(async () => {
+            const { didRefresh } = await this.updateTreeItem()
+            !didRefresh && this.refresh()
+        })
     }
 
     public override async getChildren(): Promise<AWSTreeNodeBase[]> {
+        if (this.children) {
+            return this.children
+        }
+
         const children = (await this.node.getChildren?.()) ?? []
 
-        return children.map(n => new TreeShim(n))
+        return (this.children = children.map(n => new TreeShim(n)))
+    }
+
+    private async updateTreeItem(): Promise<{ readonly didRefresh: boolean }> {
+        const item = this.node.getTreeItem()
+        if (item instanceof Promise) {
+            assign(await item, this)
+            this.refresh()
+
+            return { didRefresh: true }
+        }
+
+        assign(item, this)
+        return { didRefresh: false }
     }
 }

src/shared/vscode/commands2.ts

@@ -283,12 +283,15 @@ class CommandResource<T extends Callback = Callback, U extends any[] = any[]> {
     private buildTreeNode(id: string, args: unknown[]) {
         return (content: PartialTreeItem) => {
             const treeItem = new vscode.TreeItem(content.label, vscode.TreeItemCollapsibleState.None)
-            treeItem.command = { command: id, arguments: args, title: content.label }
+            Object.assign(treeItem, {
+                ...content,
+                command: { command: id, arguments: args, title: content.label },
+            })
 
             return {
                 id: `${id}-${(this.idCounter += 1)}`,
-                treeItem: Object.assign(treeItem, content),
                 resource: this,
+                getTreeItem: () => treeItem,
             }
         }
     }