perf(napi/oxlint): optimize compiling visitor

Author: overlookmotel

napi/oxlint2/src-js/index.js

@@ -1,14 +1,14 @@
 import { createRequire } from 'node:module';
 import { lint } from './bindings.js';
 import { DATA_POINTER_POS_32, SOURCE_LEN_POS_32 } from './generated/constants.cjs';
+import { addVisitorToCompiled, compiledVisitor, finalizeCompiledVisitor, initCompiledVisitor } from './visitor.js';
 
-// Import lazy visitor from `oxc-parser`.
+// Import methods and objects from `oxc-parser`.
 // Use `require` not `import` as `oxc-parser` uses `require` internally,
 // and need to make sure get same instance of modules as it uses internally,
 // otherwise `TOKEN` here won't be same `TOKEN` as used within `oxc-parser`.
 const require = createRequire(import.meta.url);
 const { TOKEN } = require('../../parser/raw-transfer/lazy-common.js'),
-  { Visitor, getVisitorsArr } = require('../../parser/raw-transfer/visitor.js'),
   walkProgram = require('../../parser/generated/lazy/walk.js');
 
 // --------------------
@@ -115,12 +115,17 @@ function lintFile([filePath, bufferId, buffer, ruleIds]) {
   }
 
   // Get visitors for this file from all rules
-  const visitors = ruleIds.map(
-    ruleId => registeredRules[ruleId].create(new Context(ruleId, filePath)),
-  );
-  const visitor = new Visitor(
-    visitors.length === 1 ? visitors[0] : combineVisitors(visitors),
-  );
+  initCompiledVisitor();
+  for (const ruleId of ruleIds) {
+    const visitor = registeredRules[ruleId].create(new Context(ruleId, filePath));
+    addVisitorToCompiled(visitor);
+  }
+  const needsVisit = finalizeCompiledVisitor();
+
+  // Skip visiting AST if no visitors visit any nodes.
+  // Some rules seen in the wild return an empty visitor object from `create` if some initial check fails
+  // e.g. file extension is not one the rule acts on.
+  if (!needsVisit) return '[]';
 
   // Visit AST
   const programPos = buffer.uint32[DATA_POINTER_POS_32],
@@ -130,38 +135,14 @@ function lintFile([filePath, bufferId, buffer, ruleIds]) {
   const sourceIsAscii = sourceText.length === sourceByteLen;
   const ast = { buffer, sourceText, sourceByteLen, sourceIsAscii, nodes: new Map(), token: TOKEN };
 
-  walkProgram(programPos, ast, getVisitorsArr(visitor));
+  walkProgram(programPos, ast, compiledVisitor);
 
   // Send diagnostics back to Rust
   const ret = JSON.stringify(diagnostics);
   diagnostics.length = 0;
   return ret;
 }
 
-/**
- * Combine multiple visitor objects into a single visitor object.
- * @param {Array<Object>} visitors - Array of visitor objects
- * @returns {Object} - Combined visitor object
- */
-function combineVisitors(visitors) {
-  const combinedVisitor = {};
-  for (const visitor of visitors) {
-    for (const nodeType of Object.keys(visitor)) {
-      if (!(nodeType in combinedVisitor)) {
-        combinedVisitor[nodeType] = function(node) {
-          for (const v of visitors) {
-            if (typeof v[nodeType] === 'function') {
-              v[nodeType](node);
-            }
-          }
-        };
-      }
-    }
-  }
-
-  return combinedVisitor;
-}
-
 /**
  * Context class.
  *

napi/oxlint2/src-js/visitor.js

@@ -0,0 +1,311 @@
+// Functions to compile 1 or more visitor objects into a single compiled visitor.
+//
+// # Visitor objects
+//
+// Visitor objects which are generated by rules' `create` functions have keys being either:
+// * Name of an AST type. or
+// * Name of an AST type postfixed with `:exit`.
+//
+// Each property value must be a function that handles that AST node.
+//
+// e.g.:
+//
+// ```
+// {
+//   BinaryExpression(node) {
+//     // Do stuff on enter
+//   },
+//   'BinaryExpression:exit'(node) {
+//     // Do stuff on exit
+//   },
+// }
+// ```
+//
+// # Compiled visitor
+//
+// Compiled visitor is an array with `NODE_TYPES_COUNT` length, keyed by the ID of the node type.
+// `NODE_TYPE_IDS_MAP` maps from type name to ID.
+//
+// Each element of compiled array is one of:
+// * No visitor for this type = `null`.
+// * Visitor for leaf node = visit function.
+// * Visitor for non-leaf node = object of form `{ enter, exit }`,
+//   where each property is either a visitor function or `null`.
+//
+// e.g.:
+//
+// ```
+// [
+//   // Leaf nodes
+//   function(node) { /* do stuff */ },
+//   // ...
+//
+//   // Non-leaf nodes
+//   {
+//     enter: function(node) { /* do stuff */ },
+//     exit: null,
+//   },
+//   // ...
+// ]
+// ```
+//
+// # Object reuse
+//
+// No more than 1 compiled visitor exists at any time, so we reuse a single array `compiledVisitor`,
+// rather than creating a new array for each file being linted.
+//
+// To compile visitors, call:
+// * `initCompiledVisitor` once.
+// * `addVisitorToCompiled` with each visitor object.
+// * `finalizeCompiledVisitor` once.
+//
+// After this sequence of calls, `compiledVisitor` is ready to be used to walk the AST.
+
+import types from '../../parser/generated/lazy/types.js';
+
+const { LEAF_NODE_TYPES_COUNT, NODE_TYPE_IDS_MAP, NODE_TYPES_COUNT } = types;
+
+const { isArray } = Array;
+
+// Compiled visitor used for visiting each file.
+// Same array is reused for each file.
+//
+// Initialized with `.push()` to ensure V8 treats the array as "packed" (linear array),
+// not "holey" (hash map). This is critical, as looking up elements in this array is a very hot path
+// during AST visitation, and holey arrays are much slower.
+// https://v8.dev/blog/elements-kinds
+export const compiledVisitor = [];
+
+for (let i = NODE_TYPES_COUNT; i !== 0; i--) {
+  compiledVisitor.push(null);
+}
+
+// Arrays containing type IDs of types which have multiple visit functions defined for them.
+//
+// Filled with `0` initially up to maximum size they could ever need to be so:
+// 1. These arrays never need to grow.
+// 2. V8 treats these arrays as "PACKED_SMI_ELEMENTS".
+const mergedLeafVisitorTypeIds = [],
+  mergedEnterVisitorTypeIds = [],
+  mergedExitVisitorTypeIds = [];
+
+for (let i = LEAF_NODE_TYPES_COUNT; i !== 0; i--) {
+  mergedLeafVisitorTypeIds.push(0);
+}
+
+for (let i = NODE_TYPES_COUNT - LEAF_NODE_TYPES_COUNT; i !== 0; i--) {
+  mergedEnterVisitorTypeIds.push(0);
+  mergedExitVisitorTypeIds.push(0);
+}
+
+mergedLeafVisitorTypeIds.length = 0;
+mergedEnterVisitorTypeIds.length = 0;
+mergedExitVisitorTypeIds.length = 0;
+
+// `true` if `addVisitor` has been called with a visitor which visits at least one AST type
+let hasActiveVisitors = false;
+
+/**
+ * Initialize compiled visitor, ready for calls to `addVisitor`.
+ * @returns {undefined}
+ */
+export function initCompiledVisitor() {
+  // Reset `compiledVisitor` array after previous compilation
+  for (let i = 0; i < NODE_TYPES_COUNT; i++) {
+    compiledVisitor[i] = null;
+  }
+}
+
+/**
+ * Add a visitor to compiled visitor.
+ *
+ * @param {Object} visitor - Visitor object
+ * @returns {undefined}
+ */
+export function addVisitorToCompiled(visitor) {
+  if (visitor === null || typeof visitor !== 'object') {
+    throw new TypeError('Visitor returned from `create` method must be an object');
+  }
+
+  // Exit if is empty visitor
+  const keys = Object.keys(visitor);
+  if (keys.length === 0) return;
+
+  hasActiveVisitors = true;
+
+  // Populate visitors array from provided object
+  for (let name of keys) {
+    const visitFn = visitor[name];
+    if (typeof visitFn !== 'function') {
+      throw new TypeError(`'${name}' property of visitor object is not a function`);
+    }
+
+    const isExit = name.endsWith(':exit');
+    if (isExit) name = name.slice(0, -5);
+
+    const typeId = NODE_TYPE_IDS_MAP.get(name);
+    if (typeId === void 0) throw new Error(`Unknown node type '${name}' in visitor object`);
+
+    const existing = compiledVisitor[typeId];
+    if (typeId < LEAF_NODE_TYPES_COUNT) {
+      // Leaf node - store just 1 function, not enter+exit pair
+      if (existing === null) {
+        compiledVisitor[typeId] = visitFn;
+      } else if (isArray(existing)) {
+        if (isExit) {
+          existing.push(visitFn);
+        } else {
+          // Insert before last in array in case last was enter visit function from the current rule,
+          // to ensure enter is called before exit.
+          // It could also be either an enter or exit visitor function for another rule, but the order
+          // rules are called in doesn't matter. We only need to make sure that a rule's exit visitor
+          // isn't called before enter visitor *for that same rule*.
+          existing.splice(existing.length - 1, 0, visitFn);
+        }
+      } else {
+        // Same as above, enter visitor is put to front of list to make sure enter is called before exit
+        compiledVisitor[typeId] = isExit ? [existing, visitFn] : [visitFn, existing];
+        mergedLeafVisitorTypeIds.push(typeId);
+      }
+    } else {
+      // Not leaf node - store enter+exit pair
+      if (existing === null) {
+        compiledVisitor[typeId] = isExit
+          ? { enter: null, exit: visitFn }
+          : { enter: visitFn, exit: null };
+      } else if (isExit) {
+        let { exit } = existing;
+        if (exit === null) {
+          existing.exit = visitFn;
+        } else if (isArray(exit)) {
+          exit.push(visitFn);
+        } else {
+          existing.exit = [exit, visitFn];
+          mergedExitVisitorTypeIds.push(typeId);
+        }
+      } else {
+        let { enter } = existing;
+        if (enter === null) {
+          existing.enter = visitFn;
+        } else if (isArray(enter)) {
+          enter.push(visitFn);
+        } else {
+          existing.enter = [enter, visitFn];
+          mergedEnterVisitorTypeIds.push(typeId);
+        }
+      }
+    }
+  }
+}
+
+/**
+ * Finalize compiled visitor.
+ *
+ * After calling this function, `compiledVisitor` is ready to be used to walk the AST.
+ *
+ * @returns {boolean} - `true` if compiled visitor visits at least 1 AST type
+ */
+export function finalizeCompiledVisitor() {
+  if (hasActiveVisitors === false) return false;
+
+  // Merge visit functions for node types which have multiple visitors from different rules,
+  // or enter+exit functions for leaf nodes
+  for (const typeId of mergedLeafVisitorTypeIds) {
+    compiledVisitor[typeId] = mergeVisitFns(compiledVisitor[typeId]);
+  }
+  for (const typeId of mergedEnterVisitorTypeIds) {
+    const enterExit = compiledVisitor[typeId];
+    enterExit.enter = mergeVisitFns(enterExit.enter);
+  }
+  for (const typeId of mergedExitVisitorTypeIds) {
+    const enterExit = compiledVisitor[typeId];
+    enterExit.exit = mergeVisitFns(enterExit.exit);
+  }
+
+  // Reset state, ready for next time
+  mergedLeafVisitorTypeIds.length = 0;
+  mergedEnterVisitorTypeIds.length = 0;
+  mergedExitVisitorTypeIds.length = 0;
+
+  hasActiveVisitors = false;
+
+  return true;
+}
+
+/**
+ * Merge array of visit functions into a single function, which calls each of input functions in turn.
+ *
+ * The merged function is statically defined and does not contain a loop, to hopefully allow
+ * JS engine to heavily optimize it.
+ *
+ * `mergers` contains pre-defined functions to merge up to 5 visit functions.
+ * Merger functions for merging more than 5 visit functions are created dynamically on demand.
+ *
+ * @param {Array<function>} visitFns - Array of visit functions
+ * @returns {function} - Function which calls all of `visitFns` in turn.
+ */
+function mergeVisitFns(visitFns) {
+  const numVisitFns = visitFns.length;
+
+  // Get or create merger for merging `numVisitFns` functions
+  let merger;
+  if (mergers.length <= numVisitFns) {
+    while (mergers.length < numVisitFns) {
+      mergers.push(null);
+    }
+    merger = createMerger(numVisitFns);
+    mergers.push(merger);
+  } else {
+    merger = mergers[numVisitFns];
+    if (merger === null) merger = mergers[numVisitFns] = createMerger(numVisitFns);
+  }
+
+  // Merge functions
+  return merger(...visitFns);
+}
+
+/**
+ * Create a merger function that merges `fnCount` functions.
+ *
+ * @param {number} fnCount - Number of functions to be merged
+ * @returns {function} - Function to merge `fnCount` functions
+ */
+function createMerger(fnCount) {
+  const args = [];
+  let body = 'return node=>{';
+  for (let i = 1; i <= fnCount; i++) {
+    args.push(`visit${i}`);
+    body += `visit${i}(node);`;
+  }
+  body += '}';
+  args.push(body);
+  return new Function(...args);
+}
+
+// Pre-defined mergers for merging up to 5 functions
+const mergers = [
+  null, // No merger for 0 functions
+  null, // No merger for 1 function
+  (visit1, visit2) => node => {
+    visit1(node);
+    visit2(node);
+  },
+  (visit1, visit2, visit3) => node => {
+    visit1(node);
+    visit2(node);
+    visit3(node);
+  },
+  (visit1, visit2, visit3, visit4) => node => {
+    visit1(node);
+    visit2(node);
+    visit3(node);
+    visit4(node);
+  },
+  (visit1, visit2, visit3, visit4, visit5) => node => {
+    visit1(node);
+    visit2(node);
+    visit3(node);
+    visit4(node);
+    visit5(node);
+  },
+];