add :nth-of-type and :first-of-type

Author: eush77

README.md

@@ -63,11 +63,11 @@ That's it!
   - [x] `:root`
   - [x] `:nth-child(2n+1)`
   - [x] `:nth-last-child(2n+1)`
-  - [ ] `:nth-of-type(2n+1)`
+  - [x] `:nth-of-type(2n+1)`
   - [ ] `:nth-last-of-type(2n+1)`
   - [x] `:first-child`
   - [x] `:last-child`
-  - [ ] `:first-of-type`
+  - [x] `:first-of-type`
   - [ ] `:last-of-type`
   - [x] `:only-child`
   - [ ] `:only-of-type`

lib/ast-walkers.js

@@ -1,44 +1,94 @@
 'use strict';
 
+var TypeIndex = require('./type-index');
+
 var walkers = exports;
 
 
-walkers.topScan = function (node, nodeIndex, parent, iterator) {
+walkers.topScan = function (node, nodeIndex, parent, iterator, opts) {
+  if (parent) {
+    // We would like to avoid spinning an extra loop through the starting
+    // node's siblings just to count its typeIndex.
+    throw Error('topScan is supposed to be called from the root node');
+  }
+
   iterator(node, nodeIndex, parent);
   walkers.descendant.apply(this, arguments);
 };
 
 
-walkers.descendant = function (node, nodeIndex, parent, iterator) {
-  if (node.children) {
-    node.children.forEach(function (child, childIndex) {
-      iterator(child, childIndex, node);
-      walkers.descendant(child, childIndex, node, iterator);
-    });
+walkers.descendant = function (node, nodeIndex, parent, iterator, opts) {
+  if (!node.children || !node.children.length) {
+    return;
+  }
+
+  if ((opts = opts || {}).typeIndex) {
+    var typeIndex = TypeIndex();
   }
+
+  node.children.forEach(function (child, childIndex) {
+    iterator(child, childIndex, node,
+             opts.typeIndex ? { typeIndex: typeIndex(child) } : undefined);
+    walkers.descendant(child, childIndex, node, iterator, opts);
+  });
 };
 
 
-walkers.child = function (node, nodeIndex, parent, iterator) {
-  if (node.children) {
-    node.children.forEach(function (child, childIndex) {
-      iterator(child, childIndex, node);
-    });
+walkers.child = function (node, nodeIndex, parent, iterator, opts) {
+  if (!node.children || !node.children.length) {
+    return;
+  }
+
+  if ((opts = opts || {}).typeIndex) {
+    var typeIndex = TypeIndex();
   }
+
+  node.children.forEach(function (child, childIndex) {
+    iterator(child, childIndex, node,
+             opts.typeIndex ? { typeIndex: typeIndex(child) } : undefined);
+  });
 };
 
 
-walkers.adjacentSibling = function (node, nodeIndex, parent, iterator) {
-  if (parent && ++nodeIndex < parent.children.length) {
-    iterator(parent.children[nodeIndex], nodeIndex, parent);
+walkers.adjacentSibling = function (node, nodeIndex, parent, iterator, opts) {
+  if (!parent) {
+    return;
+  }
+
+  if ((opts = opts || {}).typeIndex) {
+    var typeIndex = TypeIndex();
+
+    // Prefill type indexes with preceding nodes.
+    for (var prevIndex = 0; prevIndex <= nodeIndex; ++prevIndex) {
+      typeIndex(parent.children[prevIndex]);
+    }
+  }
+
+  if (++nodeIndex < parent.children.length) {
+    node = parent.children[nodeIndex];
+    iterator(node, nodeIndex, parent,
+             opts.typeIndex ? { typeIndex: typeIndex(node) } : undefined);
   }
 };
 
 
-walkers.generalSibling = function (node, nodeIndex, parent, iterator) {
-  if (parent) {
-    while (++nodeIndex < parent.children.length) {
-      iterator(parent.children[nodeIndex], nodeIndex, parent);
+walkers.generalSibling = function (node, nodeIndex, parent, iterator, opts) {
+  if (!parent) {
+    return;
+  }
+
+  if ((opts = opts || {}).typeIndex) {
+    var typeIndex = TypeIndex();
+
+    // Prefill type indexes with preceding nodes.
+    for (var prevIndex = 0; prevIndex <= nodeIndex; ++prevIndex) {
+      typeIndex(parent.children[prevIndex]);
     }
   }
+
+  while (++nodeIndex < parent.children.length) {
+    node = parent.children[nodeIndex];
+    iterator(node, nodeIndex, parent,
+             opts.typeIndex ? { typeIndex: typeIndex(node) } : undefined);
+  }
 };

lib/match-node.js

@@ -4,10 +4,10 @@ module.exports = matchNode;
 
 
 // Match node against a simple selector.
-function matchNode (rule, node, nodeIndex, parent) {
+function matchNode (rule, node, nodeIndex, parent, props) {
   return matchType(rule, node) &&
     matchAttrs(rule, node) &&
-    matchPseudos(rule, node, nodeIndex, parent);
+    matchPseudos(rule, node, nodeIndex, parent, props);
 }
 
 
@@ -58,7 +58,7 @@ function matchAttrs (rule, node) {
 }
 
 
-function matchPseudos (rule, node, nodeIndex, parent) {
+function matchPseudos (rule, node, nodeIndex, parent, props) {
   return !rule.pseudos || rule.pseudos.every(function (pseudo) {
     switch (pseudo.name) {
       case 'root':
@@ -70,20 +70,26 @@ function matchPseudos (rule, node, nodeIndex, parent) {
       case 'nth-last-child':
         return parent && pseudo.value(parent.children.length - 1 - nodeIndex);
 
+      case 'nth-of-type':
+        return parent && pseudo.value(props.typeIndex);
+
       case 'first-child':
         return parent && nodeIndex == 0;
 
       case 'last-child':
         return parent && nodeIndex == parent.children.length - 1;
 
+      case 'first-of-type':
+        return parent && props.typeIndex == 0;
+
       case 'only-child':
         return parent && parent.children.length == 1;
 
       case 'empty':
         return node.children && !node.children.length;
 
       case 'not':
-        return !matchNode(pseudo.value.rule, node, nodeIndex, parent);
+        return !matchNode(pseudo.value.rule, node, nodeIndex, parent, props);
 
       default:
         throw Error('Undefined pseudo-class: ' + pseudo.name);

lib/select.js

@@ -36,12 +36,12 @@ select.rule = function (rule, ast) {
       '+': walkers.adjacentSibling,
       '~': walkers.generalSibling
     })[rule.nestingOperator](
-      node, nodeIndex, parent, match.bind(null, rule)
+      node, nodeIndex, parent, match.bind(null, rule), searchOpts(rule)
     );
   }
 
-  function match (rule, node, nodeIndex, parent) {
-    if (matchNode(rule, node, nodeIndex, parent)) {
+  function match (rule, node, nodeIndex, parent, props) {
+    if (matchNode(rule, node, nodeIndex, parent, props)) {
       if (rule.rule) {
         search(rule.rule, node, nodeIndex, parent);
       }
@@ -51,3 +51,17 @@ select.rule = function (rule, ast) {
     }
   }
 };
+
+
+function searchOpts (rule) {
+  var needsTypeIndex = rule.pseudos && rule.pseudos.some(function (pseudo) {
+    return pseudo.name == 'nth-of-type' ||
+      pseudo.name == 'nth-last-of-type' ||
+      pseudo.name == 'first-of-type' ||
+      pseudo.name == 'last-of-type';
+  });
+
+  if (needsTypeIndex) {
+    return { typeIndex: true };
+  }
+}

lib/selector.js

@@ -30,7 +30,10 @@ function compileNthChecks (ast) {
     case 'rule':
       if (ast.pseudos) {
         ast.pseudos.forEach(function (pseudo) {
-          if (pseudo.name == 'nth-child' || pseudo.name == 'nth-last-child') {
+          if (pseudo.name == 'nth-child' ||
+              pseudo.name == 'nth-last-child' ||
+              pseudo.name == 'nth-of-type' ||
+              pseudo.name == 'nth-last-of-type') {
             pseudo.value = nthCheck(pseudo.value);
             pseudo.valueType = 'function';
           }

lib/type-index.js

@@ -0,0 +1,16 @@
+'use strict';
+
+
+module.exports = function TypeIndex () {
+  var typeLists = Object.create(null);
+
+  return function (node) {
+    var type = node.type;
+
+    if (!typeLists[type]) {
+      typeLists[type] = [];
+    }
+
+    return typeLists[type].push(node) - 1;
+  };
+};

test/collector.js

@@ -5,7 +5,7 @@ var Collector = require('../lib/collector');
 var test = require('tape');
 
 
-test('collector', function (t) {
+test('Collector', function (t) {
   var collect = Collector();
   collect('foo');
   collect(['foo', 'bar', 'baz', 'bar']);

test/select.js

@@ -176,6 +176,38 @@ test('structural pseudo-classes', function (t) {
     t.end();
   });
 
+  t.test(':nth-of-type', function (t) {
+    t.deepEqual(select(ast, ':root > :nth-of-type(-2n+4)'), [
+      path(ast, [1]),
+      path(ast, [5]),
+      path(ast, [6]),
+      path(ast, [11]),
+      path(ast, [12]),
+      path(ast, [14])
+    ]);
+    t.deepEqual(select(ast, ':root > :nth-of-type(odd)'), [
+      path(ast, [0]),
+      path(ast, [2]),
+      path(ast, [3]),
+      path(ast, [4]),
+      path(ast, [7]),
+      path(ast, [8]),
+      path(ast, [9]),
+      path(ast, [10]),
+      path(ast, [13]),
+      path(ast, [15]),
+      path(ast, [16]),
+      path(ast, [17]),
+      path(ast, [18])
+    ]);
+    t.deepEqual(select(ast, 'list ~ :nth-of-type(2)'), [
+      path(ast, [5]),
+      path(ast, [6]),
+      path(ast, [14])
+    ]);
+    t.end();
+  });
+
   t.test(':first-child', function (t) {
     t.deepEqual(select(ast, ':first-child'), select(ast, ':nth-child(1)'));
     t.deepEqual(select(ast, ':root:first-child'), []);
@@ -198,6 +230,27 @@ test('structural pseudo-classes', function (t) {
     t.end();
   });
 
+  t.test(':first-of-type', function (t) {
+    t.deepEqual(select(ast, ':first-of-type'), select(ast, ':nth-of-type(1)'));
+    t.deepEqual(select(ast, ':root:first-of-type'), []);
+    t.deepEqual(select(ast, ':root > :first-of-type'), [
+      path(ast, [0]),
+      path(ast, [2]),
+      path(ast, [3]),
+      path(ast, [4]),
+      path(ast, [9]),
+      path(ast, [10]),
+      path(ast, [13]),
+      path(ast, [18])
+    ]);
+    t.deepEqual(select(ast, 'code ~ :first-of-type'), [
+      path(ast, [10]),
+      path(ast, [13]),
+      path(ast, [18])
+    ]);
+    t.end();
+  });
+
   t.test(':only-child', function (t) {
     t.deepEqual(select(ast, ':only-child'),
                 select(ast, ':first-child:last-child'));

test/type-index.js

@@ -0,0 +1,22 @@
+'use strict';
+
+var TypeIndex = require('../lib/type-index');
+
+var test = require('tape');
+
+
+test('TypeIndex', function (t) {
+  var typeIndex = TypeIndex();
+
+  t.equal(typeIndex({ type: 'foo' }), 0);
+  t.equal(typeIndex({ type: 'bar' }), 0);
+  t.equal(typeIndex({ type: 'foo' }), 1);
+  t.equal(typeIndex({ type: 'foo' }), 2);
+  t.equal(typeIndex({ type: 'bar' }), 1);
+  t.equal(typeIndex({ type: 'baz' }), 0);
+
+  typeIndex = TypeIndex();
+  t.equal(typeIndex({ type: 'foo' }), 0);
+
+  t.end();
+});