Add fast path for querySelector(All)

Source/core/dom/SelectorQuery.cpp

Index: Source/core/dom/SelectorQuery.cpp
diff --git a/Source/core/dom/SelectorQuery.cpp b/Source/core/dom/SelectorQuery.cpp
index dd0440210c96f53dea60270f778f0b4469d59efc..caf128a6f99daf95531e1508cb5cd2a3d332c2f4 100644
--- a/Source/core/dom/SelectorQuery.cpp
+++ b/Source/core/dom/SelectorQuery.cpp
@@ -37,6 +37,93 @@
 
 namespace WebCore {
 
+class SimpleNodeList {
+public:
+    virtual bool isEmpty() const = 0;
+    virtual Node* next() = 0;
+};
+
+class SingleNodeList : public SimpleNodeList {
+public:
+    SingleNodeList(Node* rootNode) : m_currentNode(rootNode) { }
+
+    bool isEmpty() const { return !m_currentNode; }
+
+    Node* next()
+    {
+        Node* current = m_currentNode;
+        m_currentNode = 0;
+        return current;
+    }
+
+private:
+    Node* m_currentNode;
+};
+
+class ClassRootNodeList : public SimpleNodeList {
+public:
+    ClassRootNodeList(Node* rootNode, const AtomicString& className)
+        : m_className(className)
+        , m_rootNode(rootNode)
+        , m_currentElement(nextInternal(ElementTraversal::firstWithin(rootNode))) { }
+
+    bool isEmpty() const { return !m_currentElement; }
+
+    Node* next()
+    {
+        Node* current = m_currentElement;
+        ASSERT(current);
+        m_currentElement = nextInternal(ElementTraversal::nextSkippingChildren(m_currentElement, m_rootNode));
+        return current;
+    }
+
+private:
+    Element* nextInternal(Element* element)
+    {
+        for (; element; element = ElementTraversal::next(element, m_rootNode)) {
+            if (element->hasClass() && element->classNames().contains(m_className))
+                return element;
+        }
+        return 0;
+    }
+
+    const AtomicString& m_className;
+    Node* m_rootNode;
+    Element* m_currentElement;
+};
+
+class ClassElementList : public SimpleNodeList {
+public:
+    ClassElementList(Node* rootNode, const AtomicString& className)
+        : m_className(className)
+        , m_rootNode(rootNode)
+        , m_currentElement(nextInternal(ElementTraversal::firstWithin(rootNode))) { }
+
+    bool isEmpty() const { return !m_currentElement; }
+
+    Node* next()
+    {
+        Node* current = m_currentElement;
+        ASSERT(current);
+        m_currentElement = nextInternal(ElementTraversal::next(m_currentElement, m_rootNode));
+        return current;
+    }
+
+private:
+    Element* nextInternal(Element* element)
+    {
+        for (; element; element = ElementTraversal::next(element, m_rootNode)) {
+            if (element->hasClass() && element->classNames().contains(m_className))
+                return element;
+        }
+        return 0;
+    }
+
+    const AtomicString& m_className;
+    Node* m_rootNode;
+    Element* m_currentElement;
+};
+
 void SelectorDataList::initialize(const CSSSelectorList& selectorList)
 {
     ASSERT(m_selectors.isEmpty());
@@ -83,19 +170,13 @@ bool SelectorDataList::matches(Element* targetElement) const
 PassRefPtr<NodeList> SelectorDataList::queryAll(Node* rootNode) const
 {
     Vector<RefPtr<Node> > result;
-    execute<false>(rootNode, result);
+    executeQueryAll(rootNode, result);
     return StaticNodeList::adopt(result);
 }
 
 PassRefPtr<Element> SelectorDataList::queryFirst(Node* rootNode) const
 {
-    Vector<RefPtr<Node> > result;
-    execute<true>(rootNode, result);
-    if (result.isEmpty())
-        return 0;
-    ASSERT(result.size() == 1);
-    ASSERT(result.first()->isElementNode());
-    return toElement(result.first().get());
+    return executeQueryFirst(rootNode);
 }
 
 static inline bool isTreeScopeRoot(Node* node)
@@ -104,88 +185,276 @@ static inline bool isTreeScopeRoot(Node* node)
     return node->isDocumentNode() || node->isShadowRoot();
 }
 
-// If the first pair value is true, the returned Node is the single Element that may match the selector query.
+void SelectorDataList::collectElementsByClassName(Node* rootNode, const AtomicString& className, Vector<RefPtr<Node> >& traversalRoots) const
+{
+    for (Element* element = ElementTraversal::firstWithin(rootNode); element; element = ElementTraversal::next(element, rootNode)) {
+        if (element->hasClass() && element->classNames().contains(className))
+            traversalRoots.append(element);
+    }
+}
+
+void SelectorDataList::collectElementsByTagName(Node* rootNode, const QualifiedName& tagName, Vector<RefPtr<Node> >& traversalRoots) const
+{
+    for (Element* element = ElementTraversal::firstWithin(rootNode); element; element = ElementTraversal::next(element, rootNode)) {
+        if (SelectorChecker::tagMatches(element, tagName))
+            traversalRoots.append(element);
+    }
+}
+
+Element* SelectorDataList::findElementByClassName(Node* rootNode, const AtomicString& className) const
+{
+    for (Element* element = ElementTraversal::firstWithin(rootNode); element; element = ElementTraversal::next(element, rootNode)) {
+        if (element->hasClass() && element->classNames().contains(className))
+            return element;
+    }
+    return 0;
+}
+
+Element* SelectorDataList::findElementByTagName(Node* rootNode, const QualifiedName& tagName) const
+{
+    for (Element* element = ElementTraversal::firstWithin(rootNode); element; element = ElementTraversal::next(element, rootNode)) {
+        if (SelectorChecker::tagMatches(element, tagName))
+            return element;
+    }
+    return 0;
+}
+
+inline bool SelectorDataList::canUseFastQuery(Node* rootNode) const
+{
+    return m_selectors.size() == 1 && rootNode->inDocument() && !rootNode->document()->inQuirksMode();
+}
+
+// If returns true, traversalRoots has the elements that may match the selector query.
 //
-// If the first value is false, the returned Node is the rootNode parameter or a descendant of rootNode representing
+// If returns false, traversalRoots has the rootNode parameter or descendants of rootNode representing
 // the subtree for which we can limit the querySelector traversal.
 //
-// The returned Node may be 0, regardless of the returned bool value, if this method finds that the selectors won't
+// The travseralRoots may be empty, regardless of the returned bool value, if this method finds that the selectors won't
 // match any element.
-std::pair<bool, Node*> SelectorDataList::findTraverseRoot(Node* rootNode) const
+PassOwnPtr<SimpleNodeList> SelectorDataList::findTraverseRoots(Node* rootNode, bool& matchTraverseRoots) const
 {
     // We need to return the matches in document order. To use id lookup while there is possiblity of multiple matches
     // we would need to sort the results. For now, just traverse the document in that case.
-    if (m_selectors.size() != 1)
-        return std::make_pair(false, rootNode);
-    if (!rootNode->inDocument())
-        return std::make_pair(false, rootNode);
-    if (rootNode->document()->inQuirksMode())
-        return std::make_pair(false, rootNode);
+    ASSERT(rootNode);
+    ASSERT(m_selectors.size() == 1);
+    ASSERT(m_selectors[0].selector);
 
-    bool matchSingleNode = true;
+    bool isRightmostSelector = true;
     bool startFromParent = false;
+
     for (const CSSSelector* selector = m_selectors[0].selector; selector; selector = selector->tagHistory()) {
         if (selector->m_match == CSSSelector::Id && !rootNode->document()->containsMultipleElementsWithId(selector->value())) {
             Element* element = rootNode->treeScope()->getElementById(selector->value());
             if (element && (isTreeScopeRoot(rootNode) || element->isDescendantOf(rootNode)))
                 rootNode = element;
-            else if (!element || matchSingleNode)
+            else if (!element || isRightmostSelector)
                 rootNode = 0;
-            if (matchSingleNode)
-                return std::make_pair(true, rootNode);
+            if (isRightmostSelector) {
+                matchTraverseRoots = true;
+                return adoptPtr(new SingleNodeList(rootNode));
+            }
             if (startFromParent && rootNode)
                 rootNode = rootNode->parentNode();
-            return std::make_pair(false, rootNode);
+
+            matchTraverseRoots = false;
+            return adoptPtr(new SingleNodeList(rootNode));
         }
+
+        // If we have both CSSSelector::Id and CSSSelector::Class at the same time, we should use Id
+        // to find traverse root.
+        if (!startFromParent && selector->m_match == CSSSelector::Class) {
+            if (isRightmostSelector) {
+                matchTraverseRoots = true;
+                return adoptPtr(new ClassElementList(rootNode, selector->value()));
+            }
+            matchTraverseRoots = false;
+            return adoptPtr(new ClassRootNodeList(rootNode, selector->value()));
+        }
+
         if (selector->relation() == CSSSelector::SubSelector)
             continue;
-        matchSingleNode = false;
+        isRightmostSelector = false;
         if (selector->relation() == CSSSelector::DirectAdjacent || selector->relation() == CSSSelector::IndirectAdjacent)
             startFromParent = true;
         else
             startFromParent = false;
     }
-    return std::make_pair(false, rootNode);
+
+    matchTraverseRoots = false;
+    return adoptPtr(new SingleNodeList(rootNode));
+}
+
+void SelectorDataList::executeSlowQueryAll(Node* rootNode, Vector<RefPtr<Node> >& matchedElements) const
+{
+    for (Element* element = ElementTraversal::firstWithin(rootNode); element; element = ElementTraversal::next(element, rootNode)) {
+        for (unsigned i = 0; i < m_selectors.size(); ++i) {
+            if (selectorMatches(m_selectors[i], element, rootNode)) {
+                matchedElements.append(element);
+                break;
+            }
+        }
+    }
 }
 
-template <bool firstMatchOnly>
-void SelectorDataList::execute(Node* rootNode, Vector<RefPtr<Node> >& matchedElements) const
+void SelectorDataList::executeQueryAll(Node* rootNode, Vector<RefPtr<Node> >& matchedElements) const
 {
-    std::pair<bool, Node*> traverseRoot = findTraverseRoot(rootNode);
-    if (!traverseRoot.second)
+    if (!canUseFastQuery(rootNode))
+        return executeSlowQueryAll(rootNode, matchedElements);
+
+    ASSERT(m_selectors.size() == 1);
+    ASSERT(m_selectors[0].selector);
+
+    const CSSSelector* firstSelector = m_selectors[0].selector;
+
+    if (!firstSelector->tagHistory()) {
+        // Fast path for querySelectorAll('#id'), querySelectorAl('.foo'), and querySelectorAll('div').
+        switch (firstSelector->m_match) {
+        case CSSSelector::Id:
+            {
+                if (rootNode->document()->containsMultipleElementsWithId(firstSelector->value()))
+                    break;
+
+                // Just the same as getElementById.
+                Element* element = rootNode->treeScope()->getElementById(firstSelector->value());
+                if (element && (isTreeScopeRoot(rootNode) || element->isDescendantOf(rootNode)))
+                    matchedElements.append(element);
+                return;
+            }
+        case CSSSelector::Class:
+            return collectElementsByClassName(rootNode, firstSelector->value(), matchedElements);
+        case CSSSelector::Tag:
+            return collectElementsByTagName(rootNode, firstSelector->tagQName(), matchedElements);
+        default:
+            break; // If we need another fast path, add here.
+        }
+    }
+
+    bool matchTraverseRoots;
+    OwnPtr<SimpleNodeList> traverseRoots = findTraverseRoots(rootNode, matchTraverseRoots);
+    if (traverseRoots->isEmpty())
         return;
-    Node* traverseRootNode = traverseRoot.second;
-    if (traverseRoot.first) {
-        ASSERT(m_selectors.size() == 1);
-        ASSERT(traverseRootNode->isElementNode());
-        Element* element = toElement(traverseRootNode);
-        if (selectorMatches(m_selectors[0], element, rootNode))
-            matchedElements.append(element);
+
+    const SelectorData& selector = m_selectors[0];
+    if (matchTraverseRoots) {
+        while (!traverseRoots->isEmpty()) {
+            Node* node = traverseRoots->next();
+            Element* element = toElement(node);
+            if (selectorMatches(selector, element, rootNode))
+                matchedElements.append(element);
+        }
         return;
     }
 
-    unsigned selectorCount = m_selectors.size();
-    if (selectorCount == 1) {
-        const SelectorData& selector = m_selectors[0];
-        for (Element* element = ElementTraversal::firstWithin(rootNode); element; element = ElementTraversal::next(element, rootNode)) {
-            if (selectorMatches(selector, element, rootNode)) {
+    while (!traverseRoots->isEmpty()) {
+        Node* traverseRoot = traverseRoots->next();
+        for (Element* element = ElementTraversal::firstWithin(traverseRoot); element; element = ElementTraversal::next(element, traverseRoot)) {
+            if (selectorMatches(selector, element, rootNode))
                 matchedElements.append(element);
-                if (firstMatchOnly)
-                    return;
+        }
+    }
+}
+
+// If matchTraverseRoot is true, the returned Node is the single Element that may match the selector query.
+//
+// If matchTraverseRoot is false, the returned Node is the rootNode parameter or a descendant of rootNode representing
+// the subtree for which we can limit the querySelector traversal.
+//
+// The returned Node may be 0, regardless of matchTraverseRoot, if this method finds that the selectors won't
+// match any element.
+Node* SelectorDataList::findTraverseRoot(Node* rootNode, bool& matchTraverseRoot) const
+{
+    // We need to return the matches in document order. To use id lookup while there is possiblity of multiple matches
+    // we would need to sort the results. For now, just traverse the document in that case.
+    ASSERT(rootNode);
+    ASSERT(m_selectors.size() == 1);
+    ASSERT(m_selectors[0].selector);
+
+    bool matchSingleNode = true;
+    bool startFromParent = false;
+    for (const CSSSelector* selector = m_selectors[0].selector; selector; selector = selector->tagHistory()) {
+        if (selector->m_match == CSSSelector::Id && !rootNode->document()->containsMultipleElementsWithId(selector->value())) {
+            Element* element = rootNode->treeScope()->getElementById(selector->value());
+            if (element && (isTreeScopeRoot(rootNode) || element->isDescendantOf(rootNode)))
+                rootNode = element;
+            else if (!element || matchSingleNode)
+                rootNode = 0;
+            if (matchSingleNode) {
+                matchTraverseRoot = true;
+                return rootNode;
             }
+            if (startFromParent && rootNode)
+                rootNode = rootNode->parentNode();
+            matchTraverseRoot = false;
+            return rootNode;
         }
-        return;
+        if (selector->relation() == CSSSelector::SubSelector)
+            continue;
+        matchSingleNode = false;
+        if (selector->relation() == CSSSelector::DirectAdjacent || selector->relation() == CSSSelector::IndirectAdjacent)
+            startFromParent = true;
+        else
+            startFromParent = false;
     }
+    matchTraverseRoot = false;
+    return rootNode;
+}
+
+Element* SelectorDataList::executeSlowQueryFirst(Node* rootNode) const
+{
     for (Element* element = ElementTraversal::firstWithin(rootNode); element; element = ElementTraversal::next(element, rootNode)) {
-        for (unsigned i = 0; i < selectorCount; ++i) {
-            if (selectorMatches(m_selectors[i], element, rootNode)) {
-                matchedElements.append(element);
-                if (firstMatchOnly)
-                    return;
-                break;
+        for (unsigned i = 0; i < m_selectors.size(); ++i) {
+            if (selectorMatches(m_selectors[i], element, rootNode))
+                return element;
+        }
+    }
+    return 0;
+}
+
+Element* SelectorDataList::executeQueryFirst(Node* rootNode) const
+{
+    if (!canUseFastQuery(rootNode))
+        return executeSlowQueryFirst(rootNode);
+
+
+    const CSSSelector* selector = m_selectors[0].selector;
+    ASSERT(selector);
+
+    if (!selector->tagHistory()) {
+        // Fast path for querySelector('#id'), querySelector('.foo'), and querySelector('div').
+        // Many web developers uses querySelector with these simple selectors.
+        switch (selector->m_match) {
+        case CSSSelector::Id:
+            {
+                if (rootNode->document()->containsMultipleElementsWithId(selector->value()))
+                    break;
+                Element* element = rootNode->treeScope()->getElementById(selector->value());
+                return element && (isTreeScopeRoot(rootNode) || element->isDescendantOf(rootNode)) ? element : 0;
             }
+        case CSSSelector::Class:
+            return findElementByClassName(rootNode, selector->value());
+        case CSSSelector::Tag:
+            return findElementByTagName(rootNode, selector->tagQName());
+        default:
+            break; // If we need another fast path, add here.
         }
     }
+
+    bool matchTraverseRoot;
+    Node* traverseRootNode = findTraverseRoot(rootNode, matchTraverseRoot);
+    if (!traverseRootNode)
+        return 0;
+    if (matchTraverseRoot) {
+        ASSERT(m_selectors.size() == 1);
+        ASSERT(traverseRootNode->isElementNode());
+        Element* element = toElement(traverseRootNode);
+        return selectorMatches(m_selectors[0], element, rootNode) ? element : 0;
+    }
+
+    for (Element* element = ElementTraversal::firstWithin(traverseRootNode); element; element = ElementTraversal::next(element, traverseRootNode)) {
+        if (selectorMatches(m_selectors[0], element, rootNode))
+            return element;
+    }
+    return 0;
 }
 
 SelectorQuery::SelectorQuery(const CSSSelectorList& selectorList)