perf(build): Use WebWorker when removing private fields (#3821)

* perf(build): Use WebWorker when removing private fields

* Update build/remove-private-fields.ts

Co-authored-by: Shotaro Nakamura <79000684+nakasyou@users.noreply.github.com>

---------

Co-authored-by: Shotaro Nakamura <79000684+nakasyou@users.noreply.github.com>

build/build.ts

@@ -14,7 +14,7 @@ import type { Plugin, PluginBuild, BuildOptions } from 'esbuild'
 import * as glob from 'glob'
 import fs from 'fs'
 import path from 'path'
-import { removePrivateFields } from './remove-private-fields'
+import { cleanupWorkers, removePrivateFields } from './remove-private-fields'
 import { validateExports } from './validate-exports'
 
 const args = arg({
@@ -102,14 +102,18 @@ const dtsEntries = glob.globSync('./dist/types/**/*.d.ts')
 const writer = stdout.writer()
 writer.write('\n')
 let lastOutputLength = 0
-for (let i = 0; i < dtsEntries.length; i++) {
-  const entry = dtsEntries[i]
+let removedCount = 0
 
-  const message = `Removing private fields(${i + 1}/${dtsEntries.length}): ${entry}`
-  writer.write(`\r${' '.repeat(lastOutputLength)}`)
-  lastOutputLength = message.length
-  writer.write(`\r${message}`)
+await Promise.all(
+  dtsEntries.map(async (e) => {
+    await fs.promises.writeFile(e, await removePrivateFields(e))
+
+    const message = `Private fields removed(${++removedCount}/${dtsEntries.length}): ${e}`
+    writer.write(`\r${' '.repeat(lastOutputLength)}`)
+    lastOutputLength = message.length
+    writer.write(`\r${message}`)
+  })
+)
 
-  fs.writeFileSync(entry, removePrivateFields(entry))
-}
 writer.write('\n')
+cleanupWorkers()

build/remove-private-fields-worker.test.ts

@@ -3,7 +3,7 @@
 import fs from 'node:fs/promises'
 import os from 'node:os'
 import path from 'node:path'
-import { removePrivateFields } from './remove-private-fields'
+import { removePrivateFields } from './remove-private-fields-worker'
 
 describe('removePrivateFields', () => {
   it('Works', async () => {

build/remove-private-fields-worker.ts

@@ -0,0 +1,85 @@
+import * as ts from 'typescript'
+
+export type WorkerInput = {
+  file: string
+  taskId: number
+}
+
+export type WorkerOutput =
+  | {
+      type: 'success'
+      value: string
+      taskId: number
+    }
+  | {
+      type: 'error'
+      value: unknown
+      taskId: number
+    }
+
+const removePrivateTransformer = <T extends ts.Node>(ctx: ts.TransformationContext) => {
+  const visit: ts.Visitor = (node) => {
+    if (ts.isClassDeclaration(node)) {
+      const newMembers = node.members.filter((elem) => {
+        if (ts.isPropertyDeclaration(elem) || ts.isMethodDeclaration(elem)) {
+          for (const modifier of elem.modifiers ?? []) {
+            if (modifier.kind === ts.SyntaxKind.PrivateKeyword) {
+              return false
+            }
+          }
+        }
+        if (elem.name && ts.isPrivateIdentifier(elem.name)) {
+          return false
+        }
+        return true
+      })
+      return ts.factory.createClassDeclaration(
+        node.modifiers,
+        node.name,
+        node.typeParameters,
+        node.heritageClauses,
+        newMembers
+      )
+    }
+    return ts.visitEachChild(node, visit, ctx)
+  }
+
+  return (node: T) => {
+    const visited = ts.visitNode(node, visit)
+    if (!visited) {
+      throw new Error('The result visited is undefined.')
+    }
+    return visited
+  }
+}
+
+export const removePrivateFields = (tsPath: string) => {
+  const program = ts.createProgram([tsPath], {
+    target: ts.ScriptTarget.ESNext,
+    module: ts.ModuleKind.ESNext,
+  })
+  const file = program.getSourceFile(tsPath)
+
+  const transformed = ts.transform(file!, [removePrivateTransformer])
+  const printer = ts.createPrinter()
+  const transformedSourceFile = transformed.transformed[0] as ts.SourceFile
+  const code = printer.printFile(transformedSourceFile)
+  transformed.dispose()
+  return code
+}
+
+declare const self: Worker
+
+if (globalThis.self) {
+  self.addEventListener('message', function (e) {
+    const { file, taskId } = e.data as WorkerInput
+
+    try {
+      const result = removePrivateFields(file)
+      self.postMessage({ type: 'success', value: result, taskId } satisfies WorkerOutput)
+    } catch (e) {
+      console.error(e)
+      self.postMessage({ type: 'error', value: e, taskId } satisfies WorkerOutput)
+    }
+  })
+}

build/remove-private-fields.ts

@@ -1,52 +1,40 @@
-import * as ts from 'typescript'
+import { cpus } from 'node:os'
+import type { WorkerInput, WorkerOutput } from './remove-private-fields-worker'
 
-const removePrivateTransformer = <T extends ts.Node>(ctx: ts.TransformationContext) => {
-  const visit: ts.Visitor = (node) => {
-    if (ts.isClassDeclaration(node)) {
-      const newMembers = node.members.filter((elem) => {
-        if (ts.isPropertyDeclaration(elem) || ts.isMethodDeclaration(elem)) {
-          for (const modifier of elem.modifiers ?? []) {
-            if (modifier.kind === ts.SyntaxKind.PrivateKeyword) {
-              return false
-            }
+const workers = Array.from({ length: Math.ceil(cpus().length / 2) }).map(
+  () => new Worker(`${import.meta.dirname}/remove-private-fields-worker.ts`, { type: 'module' })
+)
+let workerIndex = 0
+let taskId = 0
+
+export async function removePrivateFields(file: string): Promise<string> {
+  const currentTaskId = taskId++
+  const worker = workers[workerIndex]
+  workerIndex = (workerIndex + 1) % workers.length
+
+  return new Promise<string>((resolve, reject) => {
+    const abortController = new AbortController()
+    worker.addEventListener(
+      'message',
+      ({ data: { type, value, taskId } }: { data: WorkerOutput }) => {
+        if (taskId === currentTaskId) {
+          if (type === 'success') {
+            resolve(value)
+          } else {
+            reject(value)
           }
-        }
-        if (elem.name && ts.isPrivateIdentifier(elem.name)) {
-          return false
-        }
-        return true
-      })
-      return ts.factory.createClassDeclaration(
-        node.modifiers,
-        node.name,
-        node.typeParameters,
-        node.heritageClauses,
-        newMembers
-      )
-    }
-    return ts.visitEachChild(node, visit, ctx)
-  }
 
-  return (node: T) => {
-    const visited = ts.visitNode(node, visit)
-    if (!visited) {
-      throw new Error('The result visited is undefined.')
-    }
-    return visited
-  }
-}
-
-export const removePrivateFields = (tsPath: string) => {
-  const program = ts.createProgram([tsPath], {
-    target: ts.ScriptTarget.ESNext,
-    module: ts.ModuleKind.ESNext,
+          abortController.abort()
+        }
+      },
+      { signal: abortController.signal }
+    )
+    worker.postMessage({ file, taskId: currentTaskId } satisfies WorkerInput)
   })
-  const file = program.getSourceFile(tsPath)
+}
 
-  const transformed = ts.transform(file!, [removePrivateTransformer])
-  const printer = ts.createPrinter()
-  const transformedSourceFile = transformed.transformed[0] as ts.SourceFile
-  const code = printer.printFile(transformedSourceFile)
-  transformed.dispose()
-  return code
+export function cleanupWorkers() {
+  for (const worker of workers) {
+    worker.terminate()
+  }
 }