fix: virtualize tree children to fix UI lag with many nodes (Fixes #47)

renovate.json

@@ -6,4 +6,4 @@
   "ignoreDeps": [
     "@types/vscode"
   ]
-}
+}

shared/src/messages.ts

@@ -145,7 +145,10 @@ export type SaveNames = z.infer<typeof saveNames>
 export const childrenById = z.object({
   message: z.literal('childrenById'),
   id: z.number(),
+  limit: z.number().optional().default(50),
+  offset: z.number().optional().default(0),
   children: z.array(zodTree).optional(),
+  total: z.number().optional(),
 })
 export type ChildrenById = z.infer<typeof childrenById>
 

src/handleMessages.ts

@@ -43,7 +43,8 @@ export function handleMessage(panel: vscode.WebviewPanel, message: unknown): voi
       break
     }
     case 'childrenById': {
-      postMessage({ ...data, children: getChildrenById(data.id) }) // TODO: stream these
+      const { children, total } = getChildrenById(data.id, data.limit, data.offset)
+      postMessage({ ...data, children, total })
       break
     }
     case 'typesById': {

src/traceTree.ts

@@ -133,14 +133,16 @@ export function showTree(startsWith: string, sourceFileName: string, position: n
   return nodes
 }
 
-export function getChildrenById(id: number) {
+export function getChildrenById(id: number, limit: number = 50, offset: number = 0) {
   const nodes = treeIdNodes.get(id)?.children ?? []
-  const ret: typeof nodes = []
-  nodes.forEach((node) => {
+  const total = nodes.length
+  const paginatedNodes = nodes.slice(offset, offset + limit)
+  const ret: typeof paginatedNodes = []
+  paginatedNodes.forEach((node) => {
     treeIdNodes.set(node.id, node)
     ret.push({ ...node, children: [], types: [] })
   })
-  return ret
+  return { children: ret, total }
 }
 
 export function getTypesById(id: number) {

ui/components/TreeNode.vue

@@ -1,17 +1,29 @@
 <script setup lang="ts">
 import type { Tree } from '../../src/traceTree'
-import { childrenById, typesById } from '~/src/appState'
+import { childrenById, childrenTotalById, typesById } from '~/src/appState'
 
 const props = defineProps<{ tree: Tree, depth: number }>()
 
 const sendMessage = useNuxtApp().$sendMessage
 
 const children = computed(() => childrenById.get(props.tree.id) ?? [])
+const totalChildren = computed(() => childrenTotalById.get(props.tree.id))
 const types = computed(() => typesById.get(props.tree.id) ?? [])
 
+// Pagination state
+const offset = ref(0)
+const limit = 50
+
 function fetchChildren() {
   if (children.value.length === 0)
-    sendMessage('childrenById', { id: props.tree.id })
+    sendMessage('childrenById', { id: props.tree.id, limit, offset: offset.value })
+}
+
+function loadMoreChildren() {
+  // Load next batch of children
+  const newOffset = offset.value + limit
+  offset.value = newOffset
+  sendMessage('childrenById', { id: props.tree.id, limit, offset: newOffset })
 }
 
 function fetchTypes() {
@@ -78,6 +90,11 @@ const insetClass = `border-e min-w-2 border-[var(--vscode-tree-inactiveIndentGui
       <template v-for="(node, idx) of children" :key="idx">
         <TreeNode :depth="depth + 1" :tree="node" />
       </template>
+      <div v-if="totalChildren !== undefined && children.length + offset < totalChildren" class="flex justify-center py-2 pl-8">
+        <button class="px-4 py-1 bg-[var(--vscode-button-background)] text-[var(--vscode-button-foreground)] rounded hover:opacity-80 transition-opacity text-sm" @click="loadMoreChildren">
+          Load More ({{ children.length + offset }}/{{ totalChildren }})
+        </button>
+      </div>
     </UExpand>
   </div>
 </template>

ui/src/appState.ts

@@ -3,6 +3,7 @@ import type { Tree } from '../../src/traceTree'
 import * as Messages from '../../shared/src/messages'
 
 export const childrenById = shallowReactive(new Map<number, Tree[]>())
+export const childrenTotalById = shallowReactive(new Map<number, number>())
 export const typesById = shallowReactive(new Map<number, TypeLine[]>())
 export const nodes = ref([] as Tree[])
 export const sortBy = ref('Timestamp' as keyof typeof sortValue)
@@ -40,6 +41,8 @@ function handleMessage(e: MessageEvent<unknown>) {
       const id = parsed.data.id
       const children = childrenById.get(id) ?? []
       childrenById.set(id, [...children, ...parsed.data.children])
+      if (parsed.data.total !== undefined)
+        childrenTotalById.set(id, parsed.data.total)
       break
     }
     case 'typesById': {