Refactor caching of nth checks

Author: wooorm

lib/parse.js

@@ -1,88 +1,33 @@
 /**
- * @typedef {import('./types.js').Selector} Selector
  * @typedef {import('./types.js').Selectors} Selectors
- * @typedef {import('./types.js').RuleSet} RuleSet
- * @typedef {import('./types.js').Rule} Rule
- * @typedef {import('./types.js').RulePseudo} RulePseudo
- * @typedef {import('./types.js').RulePseudoNth} RulePseudoNth
  */
 
 import {CssSelectorParser} from 'css-selector-parser'
-import fauxEsmNthCheck from 'nth-check'
-import {zwitch} from 'zwitch'
-
-/** @type {import('nth-check').default} */
-// @ts-expect-error
-const nthCheck = fauxEsmNthCheck.default || fauxEsmNthCheck
-
-const nth = new Set([
-  'nth-child',
-  'nth-last-child',
-  'nth-of-type',
-  'nth-last-of-type'
-])
 
 const parser = new CssSelectorParser()
 
 parser.registerAttrEqualityMods('~', '^', '$', '*')
 parser.registerSelectorPseudos('any', 'matches', 'not', 'has')
 parser.registerNestingOperators('>', '+', '~')
 
-/** @type {(query: Selectors | RuleSet | Rule | undefined) => void} */
-const compile = zwitch('type', {handlers: {selectors, ruleSet, rule}})
-
 /**
  * @param {string} selector
- * @returns {Selector}
+ * @returns {Selectors}
  */
 export function parse(selector) {
   if (typeof selector !== 'string') {
     throw new TypeError('Expected `string` as selector, not `' + selector + '`')
   }
 
-  const parsed = parser.parse(selector)
-  compile(parsed)
-  return parsed
-}
-
-/**
- * @param {Selectors} query
- * @returns {void}
- */
-function selectors(query) {
-  let index = -1
-
-  while (++index < query.selectors.length) {
-    compile(query.selectors[index])
+  const query = parser.parse(selector)
+  // Empty selectors object doesn’t match anything.
+  if (!query) {
+    return {type: 'selectors', selectors: []}
   }
-}
-
-/**
- * @param {RuleSet} query
- * @returns {void}
- */
-function ruleSet(query) {
-  rule(query.rule)
-}
-
-/**
- * @param {Rule} query
- * @returns {void}
- */
-function rule(query) {
-  const pseudos = query.pseudos || []
-  let index = -1
-
-  while (++index < pseudos.length) {
-    const pseudo = pseudos[index]
 
-    if (nth.has(pseudo.name)) {
-      // @ts-expect-error Patch a non-primitive type.
-      pseudo.value = nthCheck(pseudo.value)
-      // @ts-expect-error Patch a non-primitive type.
-      pseudo.valueType = 'function'
-    }
+  if (query.type === 'selectors') {
+    return query
   }
 
-  compile(query.rule)
+  return {type: 'selectors', selectors: [query]}
 }

lib/pseudo.js

@@ -1,7 +1,6 @@
 /**
  * @typedef {import('./types.js').Rule} Rule
  * @typedef {import('./types.js').RulePseudo} RulePseudo
- * @typedef {import('./types.js').RulePseudoNth} RulePseudoNth
  * @typedef {import('./types.js').RulePseudoSelector} RulePseudoSelector
  * @typedef {import('./types.js').Parent} Parent
  * @typedef {import('./types.js').Selector} Selector
@@ -10,10 +9,15 @@
  * @typedef {import('./types.js').Node} Node
  */
 
-import {zwitch} from 'zwitch'
+import fauxEsmNthCheck from 'nth-check'
 import {convert} from 'unist-util-is'
+import {zwitch} from 'zwitch'
 import {parent} from './util.js'
 
+/** @type {import('nth-check').default} */
+// @ts-expect-error
+const nthCheck = fauxEsmNthCheck.default || fauxEsmNthCheck
+
 const is = convert()
 
 /** @type {(rule: Rule | RulePseudo, element: Node, index: number | undefined, parent: Parent | undefined, state: SelectState) => boolean} */
@@ -191,62 +195,66 @@ function onlyChild(query, _1, _2, _3, state) {
 }
 
 /**
- * @param {RulePseudoNth} query
+ * @param {RulePseudo} query
  * @param {Node} _1
  * @param {number | undefined} _2
  * @param {Parent | undefined} _3
  * @param {SelectState} state
  * @returns {boolean}
  */
 function nthChild(query, _1, _2, _3, state) {
+  const fn = getCachedNthCheck(query)
   assertDeep(state, query)
-  return typeof state.nodeIndex === 'number' && query.value(state.nodeIndex)
+  return typeof state.nodeIndex === 'number' && fn(state.nodeIndex)
 }
 
 /**
- * @param {RulePseudoNth} query
+ * @param {RulePseudo} query
  * @param {Node} _1
  * @param {number | undefined} _2
  * @param {Parent | undefined} _3
  * @param {SelectState} state
  * @returns {boolean}
  */
 function nthLastChild(query, _1, _2, _3, state) {
+  const fn = getCachedNthCheck(query)
   assertDeep(state, query)
   return (
     typeof state.nodeCount === 'number' &&
     typeof state.nodeIndex === 'number' &&
-    query.value(state.nodeCount - state.nodeIndex - 1)
+    fn(state.nodeCount - state.nodeIndex - 1)
   )
 }
 
 /**
- * @param {RulePseudoNth} query
+ * @param {RulePseudo} query
  * @param {Node} _1
  * @param {number | undefined} _2
  * @param {Parent | undefined} _3
  * @param {SelectState} state
  * @returns {boolean}
  */
 function nthOfType(query, _1, _2, _3, state) {
+  const fn = getCachedNthCheck(query)
   assertDeep(state, query)
-  return typeof state.typeIndex === 'number' && query.value(state.typeIndex)
+  return typeof state.typeIndex === 'number' && fn(state.typeIndex)
 }
 
 /**
- * @param {RulePseudoNth} query
+ * @param {RulePseudo} query
  * @param {Node} _1
  * @param {number | undefined} _2
  * @param {Parent | undefined} _3
  * @param {SelectState} state
  * @returns {boolean}
  */
 function nthLastOfType(query, _1, _2, _3, state) {
+  const fn = getCachedNthCheck(query)
   assertDeep(state, query)
   return (
     typeof state.typeIndex === 'number' &&
     typeof state.typeCount === 'number' &&
-    query.value(state.typeCount - 1 - state.typeIndex)
+    fn(state.typeCount - 1 - state.typeIndex)
   )
 }
 
@@ -314,7 +322,7 @@ function unknownPseudo(query) {
 
 /**
  * @param {SelectState} state
- * @param {RulePseudo | RulePseudoNth} query
+ * @param {RulePseudo} query
  */
 function assertDeep(state, query) {
   if (state.shallow) {
@@ -387,3 +395,22 @@ function appendScope(value) {
 
   return selector
 }
+
+/**
+ * @param {RulePseudo} query
+ * @returns {(value: number) => boolean}
+ */
+function getCachedNthCheck(query) {
+  /** @type {(value: number) => boolean} */
+  // @ts-expect-error: cache.
+  let fn = query._cachedFn
+
+  if (!fn) {
+    // @ts-expect-error: always string.
+    fn = nthCheck(query.value)
+    // @ts-expect-error: cache.
+    query._cachedFn = fn
+  }
+
+  return fn
+}

lib/types.js

@@ -37,15 +37,6 @@
  * @property {Selectors | RuleSet} value
  *   Selector.
  *
- * @typedef RulePseudoNth
- *   Overwrite to compile nth-checks once.
- * @property {string} name
- *   Name of pseudo, such as `'nth-child'`.
- * @property {'function'} valueType
- *   Set to `'function'`, because `value` is a compiled check.
- * @property {(index: number) => boolean} value
- *   Compiled function from `nth-check`.
- *
  * @typedef SelectState
  *   Current state.
  * @property {(query: Selectors | RuleSet | Rule, node: Node | undefined, state: SelectState) => Array<Node>} any