perf: improve trace performance (#15)

Author: typeholes

messages/src/messages.ts

@@ -25,12 +25,25 @@ export const gotoLocation = z.object({
 })
 export type GotoLocation = z.infer<typeof gotoPosition>
 
-export const traceFile = z.object({
-  message: z.literal('traceFile'),
+export const traceFileStart = z.object({
+  message: z.literal('traceFileStart'),
   fileName: z.string(),
-  traceString: z.string(),
+  size: z.number(),
 })
-export type TraceFile = z.infer<typeof traceFile>
+export type TraceFileStart = z.infer<typeof traceFileStart>
+
+export const traceFileEnd = z.object({
+  message: z.literal('traceFileEnd'),
+  fileName: z.string(),
+})
+export type TraceFileEnd = z.infer<typeof traceFileEnd>
+
+export const traceFileChunk = z.object({
+  message: z.literal('traceFileChunk'),
+  fileName: z.string(),
+  chunk: z.string(),
+})
+export type TraceFileChunk = z.infer<typeof traceFileChunk>
 
 export const gotoTracePosition = z.object({
   message: z.literal('gotoTracePosition'),
@@ -62,4 +75,4 @@ export const fileStats = z.object({
 export type FileStats = z.infer<typeof fileStats>
 
 export type Message = z.infer<typeof message>
-export const message = z.union([ping, pong, gotoLocation, gotoPosition, traceFile, gotoTracePosition, positionTypeCounts, fileStats])
+export const message = z.union([ping, pong, gotoLocation, gotoPosition, traceFileChunk, traceFileEnd, traceFileStart, gotoTracePosition, positionTypeCounts, fileStats])

playground/index.js

@@ -15,3 +15,7 @@ function lazy() {
 }
 //
 //
+const multiLine = `a
+b
+c
+d`

playground/index.ts

@@ -51,3 +51,14 @@ function lazy() {
 }
 //
 //
+
+const multiLine = `a
+b
+c
+d`
+
+type MultiLine = typeof multiLine | `
+1
+2
+3
+4`

playground/traces/trace.json

@@ -1,6 +1,6 @@
 import { join } from 'node:path'
 import { spawnSync } from 'node:child_process'
-import { mkdirSync, readFileSync, readdirSync } from 'node:fs'
+import { createReadStream, mkdirSync, readdirSync, statSync } from 'node:fs'
 import * as vscode from 'vscode'
 import { getTracePanel, postMessage, prepareWebView } from './webview'
 import { getCurrentConfig } from './configuration'
@@ -21,11 +21,32 @@ const commandHandlers = {
     if (!document)
       return
     const fileName = document.fileName
-    const traceString = document.getText()
-    postMessage({ message: 'traceFile', fileName, traceString })
+    sendTrace('', fileName)
   },
 } as const
 
+function sendTrace(dirName: string, fileName: string) {
+  const fullFileName = join(dirName, fileName)
+  const stat = statSync(fullFileName)
+  const size = stat.size
+
+  const stream = createReadStream(fullFileName, { autoClose: true, emitClose: true, encoding: 'utf-8' })
+
+  postMessage({ message: 'traceFileStart', fileName, size })
+
+  stream.on('end', () => postMessage({ message: 'traceFileEnd', fileName }))
+
+  function readChunks() {
+    const chunk = stream.read()
+    if (chunk === null)
+      return
+
+    postMessage({ message: 'traceFileChunk', fileName, chunk })
+    setImmediate(readChunks)
+  }
+  stream.on('readable', readChunks)
+}
+
 export function registerCommands(context: vscode.ExtensionContext) {
   for (const cmd in commandHandlers) {
     const disposable = vscode.commands.registerCommand(
@@ -67,7 +88,7 @@ function runTrace(context: vscode.ExtensionContext) {
 
   const projectPath = getProjectPath()
   if (!projectPath) {
-    vscode.window.showErrorMessage('could not get procject path from workspace folders')
+    vscode.window.showErrorMessage('could not get project path from workspace folders')
     return
   }
 
@@ -79,10 +100,8 @@ function runTrace(context: vscode.ExtensionContext) {
 
   try {
     const fileNames = readdirSync(traceDir)
-    for (const fileName of fileNames) {
-      const traceString = readFileSync(join(traceDir, fileName)).toString()
-      postMessage({ message: 'traceFile', fileName, traceString })
-    }
+    for (const fileName of fileNames)
+      sendTrace(traceDir, fileName)
   }
   catch (e) {
     vscode.window.showErrorMessage(e instanceof Error ? e.message : `${e}`)

playground/traces/types.json

@@ -5,33 +5,65 @@ import { type Tree, toTree } from '~/src/traceTree'
 
 const files = useState('files', () => ({}) as Record<string, TraceData>)
 
-const traceTree = useState('traceTree', () => undefined as Tree | undefined)
+const traceTree = useState('traceTree', () => shallowRef(undefined as Tree | undefined))
+
+const tmpTraceStrings: Record<string, string> = {}
+
+const fileProgress = reactive({} as Record<string, string>)
+
+const fileSizes = reactive({} as Record<string, [number, number]>)
+
+const inProcess = ref(0)
 
 function handleMessage(e: MessageEvent<unknown>) {
   const parsed = Messages.message.safeParse(e.data)
   if (!parsed.success)
     return
 
-  if (parsed.data.message === 'traceFile') {
-    try {
-      const json = JSON.parse(parsed.data.traceString)
+  switch (parsed.data.message) {
+    case 'traceFileStart':
+      inProcess.value++
+      fileSizes[parsed.data.fileName] = [0, parsed.data.size]
+      tmpTraceStrings[parsed.data.fileName] = ''
+      files.value[parsed.data.fileName] ??= []
+      break
+
+    case 'traceFileChunk': {
+      const fileName = parsed.data.fileName
+      const size = fileSizes[fileName]
+      size[0] += parsed.data.chunk.length
+      fileProgress[fileName] = size.join(' / ')
+      tmpTraceStrings[fileName] += parsed.data.chunk
+      break
+    }
+
+    case 'traceFileEnd': {
+      inProcess.value--
+      fileProgress[parsed.data.fileName] = 'Received'
+
+      try {
+        const json = JSON.parse(tmpTraceStrings[parsed.data.fileName])
 
-      const arr = traceData.safeParse(json)
-      if (!arr.success)
-        return
+        const arr = traceData.safeParse(json)
+        if (!arr.success)
+          return
 
-      files.value[parsed.data.fileName] = arr.data
+        files.value[parsed.data.fileName] = shallowReactive(arr.data)
+      }
+      catch (_e) {}
+
+      break
     }
-    catch (_e) {}
   }
 }
 
+const filters = useState('treeFilters', () => ({ startsWith: 'check', sourceFileName: '', position: 0 }))
+
 function processTraces() {
   const values = Object.values(files.value).flat(1) as unknown as TraceData
-  // eslint-disable-next-line no-debugger, no-restricted-syntax
-  debugger
   const tree = toTree(values)
   traceTree.value = tree
+  filters.value = { startsWith: 'check', sourceFileName: '', position: 0 }
 }
 
 onMounted(() => {
@@ -42,9 +74,9 @@ onMounted(() => {
 <template>
   <div>
     <div v-for="(_data, fileName) in files" :key="fileName">
-      <div>{{ fileName }}</div>
+      <div>{{ fileName }}  {{ fileProgress[fileName] }}</div>
     </div>
-    <vscode-button @click="processTraces">
+    <vscode-button :disable="inProcess > 0" @click="processTraces">
       Process Traces
     </vscode-button>
     <div />