perf: improve vi.mock performance (#2594)

Author: sheremet-va

packages/vitest/src/integrations/vi.ts

@@ -1,9 +1,10 @@
 import type { FakeTimerInstallOpts } from '@sinonjs/fake-timers'
-import { parseStacktrace } from '../utils/source-map'
+import { parseSingleStack } from '../utils/source-map'
 import type { VitestMocker } from '../runtime/mocker'
 import type { ResolvedConfig, RuntimeConfig } from '../types'
 import { getWorkerState, resetModules, waitForImportsToResolve } from '../utils'
 import type { MockFactoryWithHelper } from '../types/mocker'
+import { createSimpleStackTrace } from '../utils/error'
 import { FakeTimers } from './mock/timers'
 import type { EnhancedSpy, MaybeMocked, MaybeMockedDeep, MaybePartiallyMocked, MaybePartiallyMockedDeep } from './spy'
 import { fn, isMockFunction, spies, spyOn } from './spy'
@@ -110,9 +111,10 @@ class VitestUtils {
   fn = fn
 
   private getImporter() {
-    const err = new Error('mock')
-    const [,, importer] = parseStacktrace(err, true)
-    return importer.file
+    const stackTrace = createSimpleStackTrace({ stackTraceLimit: 4 })
+    const importerStack = stackTrace.split('\n')[4]
+    const stack = parseSingleStack(importerStack)
+    return stack?.file || ''
   }
 
   /**

packages/vitest/src/utils/error.ts

@@ -0,0 +1,22 @@
+interface ErrorOptions {
+  message?: string
+  stackTraceLimit?: number
+}
+
+/**
+ * Get original stacktrace without source map support the most performant way.
+ * - Create only 1 stack frame.
+ * - Rewrite prepareStackTrace to bypass "support-stack-trace" (usually takes ~250ms).
+ */
+export function createSimpleStackTrace(options?: ErrorOptions) {
+  const { message = 'error', stackTraceLimit = 1 } = options || {}
+  const limit = Error.stackTraceLimit
+  const prepareStackTrace = Error.prepareStackTrace
+  Error.stackTraceLimit = stackTraceLimit
+  Error.prepareStackTrace = e => e.stack
+  const err = new Error(message)
+  const stackTrace = err.stack || ''
+  Error.prepareStackTrace = prepareStackTrace
+  Error.stackTraceLimit = limit
+  return stackTrace
+}

packages/vitest/src/utils/source-map.ts

@@ -27,63 +27,69 @@ function extractLocation(urlLike: string) {
   return [parts[1], parts[2] || undefined, parts[3] || undefined]
 }
 
-export function parseStacktrace(e: ErrorWithDiff, full = false): ParsedStack[] {
-  if (!e)
-    return []
+// Based on https://github.com/stacktracejs/error-stack-parser
+// Credit to stacktracejs
+export function parseSingleStack(raw: string): ParsedStack | null {
+  let line = raw.trim()
 
-  if (e.stacks)
-    return e.stacks
+  if (line.includes('(eval '))
+    line = line.replace(/eval code/g, 'eval').replace(/(\(eval at [^()]*)|(,.*$)/g, '')
 
-  const stackStr = e.stack || e.stackStr || ''
-  const stackFrames = stackStr
-    .split('\n')
-    // Based on https://github.com/stacktracejs/error-stack-parser
-    // Credit to stacktracejs
-    .map((raw): ParsedStack | null => {
-      let line = raw.trim()
+  let sanitizedLine = line
+    .replace(/^\s+/, '')
+    .replace(/\(eval code/g, '(')
+    .replace(/^.*?\s+/, '')
 
-      if (line.includes('(eval '))
-        line = line.replace(/eval code/g, 'eval').replace(/(\(eval at [^()]*)|(,.*$)/g, '')
+  // capture and preserve the parenthesized location "(/foo/my bar.js:12:87)" in
+  // case it has spaces in it, as the string is split on \s+ later on
+  const location = sanitizedLine.match(/ (\(.+\)$)/)
 
-      let sanitizedLine = line
-        .replace(/^\s+/, '')
-        .replace(/\(eval code/g, '(')
-        .replace(/^.*?\s+/, '')
+  // remove the parenthesized location from the line, if it was matched
+  sanitizedLine = location ? sanitizedLine.replace(location[0], '') : sanitizedLine
 
-      // capture and preserve the parenthesized location "(/foo/my bar.js:12:87)" in
-      // case it has spaces in it, as the string is split on \s+ later on
-      const location = sanitizedLine.match(/ (\(.+\)$)/)
+  // if a location was matched, pass it to extractLocation() otherwise pass all sanitizedLine
+  // because this line doesn't have function name
+  const [url, lineNumber, columnNumber] = extractLocation(location ? location[1] : sanitizedLine)
+  let method = (location && sanitizedLine) || ''
+  let file = url && ['eval', '<anonymous>'].includes(url) ? undefined : url
 
-      // remove the parenthesized location from the line, if it was matched
-      sanitizedLine = location ? sanitizedLine.replace(location[0], '') : sanitizedLine
+  if (!file || !lineNumber || !columnNumber)
+    return null
 
-      // if a location was matched, pass it to extractLocation() otherwise pass all sanitizedLine
-      // because this line doesn't have function name
-      const [url, lineNumber, columnNumber] = extractLocation(location ? location[1] : sanitizedLine)
-      let method = (location && sanitizedLine) || ''
-      let file = url && ['eval', '<anonymous>'].includes(url) ? undefined : url
+  if (method.startsWith('async '))
+    method = method.slice(6)
 
-      if (!file || !lineNumber || !columnNumber)
-        return null
+  if (file.startsWith('file://'))
+    file = file.slice(7)
 
-      if (method.startsWith('async '))
-        method = method.slice(6)
+  // normalize Windows path (\ -> /)
+  file = resolve(file)
 
-      if (file.startsWith('file://'))
-        file = file.slice(7)
+  return {
+    method,
+    file,
+    line: parseInt(lineNumber),
+    column: parseInt(columnNumber),
+  }
+}
 
-      // normalize Windows path (\ -> /)
-      file = resolve(file)
+export function parseStacktrace(e: ErrorWithDiff, full = false): ParsedStack[] {
+  if (!e)
+    return []
+
+  if (e.stacks)
+    return e.stacks
+
+  const stackStr = e.stack || e.stackStr || ''
+  const stackFrames = stackStr
+    .split('\n')
+    .map((raw): ParsedStack | null => {
+      const stack = parseSingleStack(raw)
 
-      if (!full && stackIgnorePatterns.some(p => file && file.includes(p)))
+      if (!stack || (!full && stackIgnorePatterns.some(p => stack.file.includes(p))))
         return null
 
-      return {
-        method,
-        file,
-        line: parseInt(lineNumber),
-        column: parseInt(columnNumber),
-      }
+      return stack
     })
     .filter(notNullish)