perf: optimize isGriffelCSSModule, add processAssets timing

change/@griffel-webpack-plugin-df2684c1-fec6-44ca-bc49-56fe55560cfa.json

@@ -0,0 +1,7 @@
+{
+  "type": "patch",
+  "comment": "perf: optimize isGriffelCSSModule check, add processAssets timing",
+  "packageName": "@griffel/webpack-plugin",
+  "email": "olfedias@microsoft.com",
+  "dependentChangeType": "patch"
+}

packages/transform-shaker/src/DepsGraph.ts

@@ -6,20 +6,16 @@
 import type { Node } from 'oxc-parser';
 
 import type { IdentifierNode, StringLiteralNode } from './ast.js';
-import type { PromisedNode } from './scope.js';
+import { PromisedNode } from './scope.js';
 import type ScopeManager from './scope.js';
 import { resolveNode } from './scope.js';
 
 type Action = (this: DepsGraph, a: Node, b: Node) => void;
 
-function addEdge(this: DepsGraph, a: Node, b: Node) {
-  if (this.dependencies.has(a) && this.dependencies.get(a)!.has(b)) {
-    // edge has been already added
-    return;
-  }
-
-  if (this.dependencies.has(a)) {
-    this.dependencies.get(a)!.add(b);
+function addEdgeResolved(this: DepsGraph, a: Node, b: Node) {
+  let deps = this.dependencies.get(a);
+  if (deps) {
+    deps.add(b);
   } else {
     this.dependencies.set(a, new Set([b]));
   }
@@ -61,7 +57,12 @@ export default class DepsGraph {
   constructor(protected scope: ScopeManager) {}
 
   addEdge(dependent: Node | PromisedNode, dependency: Node | PromisedNode) {
-    this.actionQueue.push([addEdge, dependent, dependency]);
+    // Fast path: if both are already resolved nodes, add edge directly
+    if (!PromisedNode.is(dependent) && !PromisedNode.is(dependency)) {
+      addEdgeResolved.call(this, dependent, dependency);
+      return;
+    }
+    this.actionQueue.push([addEdgeResolved, dependent, dependency]);
   }
 
   addExport(name: string, node: Node) {

packages/transform-shaker/src/Visitors.ts

@@ -79,12 +79,23 @@ const visitors = {
 
 const isKeyOfVisitors = (type: string): type is string & keyof Visitors => type in visitors;
 
+const visitorsCache = new Map<string, Visitor<Node>[]>();
+
 export function getVisitors<TNode extends Node>(node: TNode): Visitor<TNode>[] {
-  const aliases = ALIAS_KEYS[node.type] || [];
-  const aliasVisitors = aliases
-    .map(type => (isKeyOfVisitors(type) ? visitors[type] : null))
-    .filter(i => i) as Visitor<TNode>[];
-  return [...aliasVisitors, visitors[node.type] as Visitor<TNode>].filter(v => v);
+  const nodeType = node.type;
+  let cached = visitorsCache.get(nodeType);
+  if (cached) return cached as Visitor<TNode>[];
+
+  const aliases = ALIAS_KEYS[nodeType] || [];
+  const result: Visitor<Node>[] = [];
+  for (let i = 0; i < aliases.length; i++) {
+    const type = aliases[i];
+    if (isKeyOfVisitors(type)) result.push(visitors[type]!);
+  }
+  if (visitors[nodeType]) result.push(visitors[nodeType]!);
+
+  visitorsCache.set(nodeType, result);
+  return result as Visitor<TNode>[];
 }
 
 export default visitors;

packages/transform-shaker/src/graphBuilder.ts

@@ -124,34 +124,33 @@ class GraphBuilder {
    * both of them are required for evaluating the value of the expression
    */
   baseVisit<TNode extends Node>(node: TNode, ignoreDeps = false) {
-    const dependencies: Node[] = [];
-    const isExpr = isExpression(node);
+    const isExpr = !ignoreDeps && isExpression(node);
     const keys = getVisitorKeys(node);
-    keys.forEach(key => {
+
+    for (let ki = 0; ki < keys.length; ki++) {
+      const key = keys[ki];
       // Ignore all types
       if (key === 'typeArguments' || key === 'typeParameters') {
-        return;
+        continue;
       }
 
       const subNode = node[key as keyof TNode];
 
       if (Array.isArray(subNode)) {
         for (let i = 0; i < subNode.length; i++) {
           const child = subNode[i];
-          if (child && this.visit(child, node, key, i) !== 'ignore') {
-            dependencies.push(child);
+          if (child && this.visit(child, node, key, i) !== 'ignore' && isExpr) {
+            this.graph.addEdge(node, child);
           }
         }
-      } else if (isNode(subNode) && this.visit(subNode, node, key) !== 'ignore') {
-        dependencies.push(subNode);
+      } else if (isNode(subNode) && this.visit(subNode, node, key) !== 'ignore' && isExpr) {
+        this.graph.addEdge(node, subNode);
       }
-    });
-
-    if (isExpr && !ignoreDeps) {
-      dependencies.forEach(dep => this.graph.addEdge(node, dep));
     }
 
-    this.callbacks.forEach(callback => callback(node));
+    for (let ci = 0; ci < this.callbacks.length; ci++) {
+      this.callbacks[ci](node);
+    }
   }
 
   visit<TNode extends Node, TParent extends Node>(
@@ -249,10 +248,8 @@ class GraphBuilder {
     const visitors = getVisitors(node);
     let action: VisitorAction = undefined;
     if (visitors.length > 0) {
-      let visitor: Visitor<TNode> | undefined;
-      // eslint-disable-next-line no-cond-assign
-      while (!action && (visitor = visitors.shift())) {
-        const method: Visitor<TNode> = visitor.bind(this);
+      for (let vi = 0; vi < visitors.length && !action; vi++) {
+        const method: Visitor<TNode> = visitors[vi].bind(this);
         action = method(node, parent, parentKey, listIdx);
       }
     } else {

packages/transform-shaker/src/scope.ts

@@ -94,12 +94,12 @@ export default class ScopeManager {
     scopeIds.set(scope, scopeId);
     this.map.set(scopeId, scope);
     this.handlers.set(scopeId, []);
-    this.stack.unshift(scope);
+    this.stack.push(scope);
     return scope;
   }
 
   dispose(): Scope | undefined {
-    const disposed = this.stack.shift();
+    const disposed = this.stack.pop();
     if (disposed) {
       this.map.delete(scopeIds.get(disposed)!);
     }
@@ -129,7 +129,15 @@ export default class ScopeManager {
 
     const identifier = identifierOrMemberExpression;
     const idName = identifier.name;
-    const scope = this.stack.slice(stack).find(s => !isHoistable || functionScopes.has(s))!;
+    // Search from innermost (end) towards outermost, skipping `stack` levels from the end
+    let scope: Scope | undefined;
+    for (let i = this.stack.length - 1 - stack; i >= 0; i--) {
+      if (!isHoistable || functionScopes.has(this.stack[i])) {
+        scope = this.stack[i];
+        break;
+      }
+    }
+    scope = scope!;
     if (this.global.has(idName) && !globalIdentifiers.has(idName)) {
       // It's probably a declaration of a previous referenced identifier
       // Let's use naïve implementation of hoisting
@@ -152,7 +160,14 @@ export default class ScopeManager {
     const name = isIdentifier(identifierOrMemberExpression)
       ? identifierOrMemberExpression.name
       : getExportName(identifierOrMemberExpression);
-    const scope = this.stack.find(s => s.has(name)) ?? this.global;
+    let scope: Scope | undefined;
+    for (let i = this.stack.length - 1; i >= 0; i--) {
+      if (this.stack[i].has(name)) {
+        scope = this.stack[i];
+        break;
+      }
+    }
+    if (!scope) scope = this.global;
     const id = getId(scope, name);
     if (scope === this.global && !scope.has(name)) {
       scope.set(name, new Set());
@@ -165,9 +180,11 @@ export default class ScopeManager {
 
   whereIsDeclared(identifier: IdentifierNode): ScopeId | undefined {
     const { name } = identifier;
-    const scope = this.stack.find(s => s.has(name) && s.get(name)!.has(identifier));
-    if (scope) {
-      return scopeIds.get(scope);
+    for (let i = this.stack.length - 1; i >= 0; i--) {
+      const s = this.stack[i];
+      if (s.has(name) && s.get(name)!.has(identifier)) {
+        return scopeIds.get(s);
+      }
     }
 
     if (this.global.has(name)) {
@@ -198,7 +215,7 @@ export default class ScopeManager {
   }
 
   addDeclareHandler(handler: DeclareHandler): () => void {
-    const scopeId = scopeIds.get(this.stack[0])!;
+    const scopeId = scopeIds.get(this.stack[this.stack.length - 1])!;
     this.handlers.get(scopeId)!.push(handler);
     return () => {
       const handlers = this.handlers.get(scopeId)!.filter(h => h !== handler);

packages/transform/src/evaluation/module.mts

@@ -35,6 +35,27 @@ let cache: Record<string, Module> = {};
 
 const NOOP = () => {};
 
+// Reusable VM context — avoids expensive vm.createContext() per evaluate() call.
+// Per-module bindings are swapped before each script.runInContext().
+const sharedSandbox = {
+  clearImmediate: NOOP,
+  clearInterval: NOOP,
+  clearTimeout: NOOP,
+  setImmediate: NOOP,
+  setInterval: NOOP,
+  setTimeout: NOOP,
+  fetch: NOOP,
+  global,
+  process: mockProcess,
+  // Per-module bindings (mutated before each run, prefixed with __ to avoid
+  // colliding with the function parameter names passed to the wrapper IIFE)
+  __module: null as Module | null,
+  __require: null as ((id: string) => unknown) | null,
+  __filename: '',
+  __dirname: '',
+};
+const sharedContext = vm.createContext(sharedSandbox);
+
 /** Checks if a value is an Error-like object (works across VM contexts where `instanceof Error` fails). */
 function isError(e: unknown): e is Error {
   return e != null && typeof e === 'object' && 'message' in e && 'stack' in e;
@@ -216,29 +237,19 @@ export class Module {
       this.debug('evaluate', `${this.filename} (only ${(only || []).join(', ')}):\n${code}`);
     }
 
-    const script = new vm.Script(`(function (exports) { ${code}\n})(exports);`, {
-      filename: this.filename,
-    });
+    const script = new vm.Script(
+      `(function (exports, module, require, __filename, __dirname) { ${code}\n})(__module.exports, __module, __require, __filename, __dirname);`,
+      { filename: this.filename },
+    );
+
+    // Swap per-module bindings on the shared context
+    sharedSandbox.__module = this;
+    sharedSandbox.__require = this.require;
+    sharedSandbox.__filename = this.filename;
+    sharedSandbox.__dirname = path.dirname(this.filename);
 
     try {
-      script.runInContext(
-        vm.createContext({
-          clearImmediate: NOOP,
-          clearInterval: NOOP,
-          clearTimeout: NOOP,
-          setImmediate: NOOP,
-          setInterval: NOOP,
-          setTimeout: NOOP,
-          fetch: NOOP,
-          global,
-          process: mockProcess,
-          module: this,
-          exports: this.exports,
-          require: this.require,
-          __filename: this.filename,
-          __dirname: path.dirname(this.filename),
-        }),
-      );
+      script.runInContext(sharedContext);
     } catch (vmError: unknown) {
       // Errors thrown inside vm.runInContext() use the VM context's Error constructor,
       // so they fail `instanceof Error` in the host context (e.g. webpack wraps them as NonErrorEmittedError).

packages/webpack-plugin/src/GriffelPlugin.mts

@@ -1,6 +1,8 @@
 import { defaultCompareMediaQueries, type GriffelRenderer } from '@griffel/core';
 import type { Compilation, Chunk, Compiler, Module, sources } from 'webpack';
 
+import * as path from 'node:path';
+
 import { PLUGIN_NAME, GriffelCssLoaderContextKey, type SupplementedLoaderContext } from './constants.mjs';
 import { createResolverFactory, type TransformResolverFactory } from './resolver/createResolverFactory.mjs';
 import { parseCSSRules } from './utils/parseCSSRules.mjs';
@@ -56,23 +58,11 @@ function getAssetSourceContents(assetSource: sources.Source): string {
   return source.toString();
 }
 
-// https://github.com/webpack-contrib/mini-css-extract-plugin/blob/26334462e419026086856787d672b052cd916c62/src/index.js#L90
-type CSSModule = Module & {
-  content: Buffer;
-};
-
-function isCSSModule(module: Module): module is CSSModule {
-  return module.type === 'css/mini-extract';
-}
+const __dirname = path.dirname(new URL(import.meta.url).pathname);
+const virtualLoaderPath = path.resolve(__dirname, 'virtual-loader', 'index.cjs');
 
 function isGriffelCSSModule(module: Module): boolean {
-  if (isCSSModule(module)) {
-    if (Buffer.isBuffer(module.content)) {
-      return module.content.indexOf('/** @griffel:css-start') !== -1;
-    }
-  }
-
-  return false;
+  return module.type === 'css/mini-extract' && module.identifier().includes(virtualLoaderPath);
 }
 
 function moveCSSModulesToGriffelChunk(compilation: Compilation) {
@@ -121,6 +111,7 @@ export class GriffelPlugin {
       evaluationMode: 'ast' | 'vm';
     }
   > = {};
+  #processAssetsTime: bigint = 0n;
   readonly #perfIssues: Map<string, { type: 'cjs-module' | 'barrel-export-star'; dependencyFilename: string; sourceFilenames: Set<string> }> = new Map();
 
   constructor(options: GriffelCSSExtractionPluginOptions = {}) {
@@ -297,6 +288,8 @@ export class GriffelPlugin {
           stage: Compilation.PROCESS_ASSETS_STAGE_PRE_PROCESS,
         },
         assets => {
+          const start = this.#collectStats ? process.hrtime.bigint() : 0n;
+
           let cssAssetDetails;
 
           // @ Rspack compat
@@ -327,6 +320,10 @@ export class GriffelPlugin {
           const cssSource = sortCSSRules([cssRulesByBucket], this.#compareMediaQueries);
 
           compilation.updateAsset(cssAssetName, new compiler.webpack.sources.RawSource(remainingCSS + cssSource));
+
+          if (this.#collectStats) {
+            this.#processAssetsTime = process.hrtime.bigint() - start;
+          }
         },
       );
 
@@ -353,20 +350,21 @@ export class GriffelPlugin {
             const fileCount = entries.length;
             const avgTime = fileCount > 0 ? totalTime / BigInt(fileCount) : 0n;
 
-            console.log('\nGriffel CSS extraction stats:');
+            const astEntries = entries.filter(s => s[1].evaluationMode === 'ast');
+            const vmEntries = entries.filter(s => s[1].evaluationMode === 'vm');
+            const astTime = astEntries.reduce((acc, cur) => acc + cur[1].time, 0n);
+            const vmTime = vmEntries.reduce((acc, cur) => acc + cur[1].time, 0n);
+            const astHitPct = ((astEntries.length / fileCount) * 100).toFixed(1) + '%';
 
-            console.log('------------------------------------');
-            console.log('Total time spent in Griffel loader:', logTime(totalTime));
-            console.log('Files processed:', fileCount);
-            console.log('Average time per file:', logTime(avgTime));
-            console.log(
-              'AST evaluation hit: ',
-              ((entries.filter(s => s[1].evaluationMode === 'ast').length / fileCount) * 100).toFixed(2) + '%',
-            );
-            console.log('------------------------------------');
+            console.log(`\n[Griffel] ${fileCount} files processed`);
+            console.log(`[Griffel] Loader: ${logTime(totalTime)} (AST ${logTime(astTime)} | VM ${logTime(vmTime)}), avg ${logTime(avgTime)}/file, AST eval hit ${astHitPct}`);
+            console.log(`[Griffel] Plugin: ${logTime(this.#processAssetsTime)}`);
+            console.log('');
 
             for (const [filename, info] of entries) {
-              console.log(`  ${logTime(info.time)} - ${filename} (evaluation mode: ${info.evaluationMode})`);
+              const time = logTime(info.time).padStart(6);
+              const mode = info.evaluationMode === 'vm' ? 'vm ' : 'ast';
+              console.log(`  ${time} ${mode} ${filename}`);
             }
 
             console.log();
@@ -377,16 +375,13 @@ export class GriffelPlugin {
             const cjsCount = issues.filter(i => i.type === 'cjs-module').length;
             const barrelCount = issues.filter(i => i.type === 'barrel-export-star').length;
 
-            console.log('\nGriffel performance issues:');
-            console.log('------------------------------------');
-            console.log(`CJS modules (no tree-shaking): ${cjsCount}`);
-            console.log(`Barrel files with remaining export *: ${barrelCount}`);
-            console.log('------------------------------------');
+            console.log(`[Griffel] Perf issues: ${cjsCount} CJS (no tree-shaking), ${barrelCount} barrel (export *)`);
+            console.log('');
 
             for (const issue of issues) {
-              const tag = issue.type === 'cjs-module' ? 'cjs' : 'barrel';
+              const tag = issue.type === 'cjs-module' ? ' cjs' : 'barrel';
               const sources = Array.from(issue.sourceFilenames).join(', ');
-              console.log(`  [${tag}] ${issue.dependencyFilename} (source: ${sources})`);
+              console.log(`  ${tag} ${issue.dependencyFilename} (from: ${sources})`);
             }
 
             console.log();