Implement QuerySelectorIterator (fixes #4082)

components/script/dom/node.rs

@@ -50,7 +50,7 @@ use devtools_traits::NodeInfo;
 use script_traits::UntrustedNodeAddress;
 use servo_util::geometry::Au;
 use servo_util::str::{DOMString, null_str_as_empty};
-use style::{parse_selector_list_from_str, matches};
+use style::{parse_selector_list_from_str, matches, SelectorList};
 
 use js::jsapi::{JSContext, JSObject, JSTracer, JSRuntime};
 use js::jsfriendapi;
@@ -377,6 +377,29 @@ impl<'a> PrivateNodeHelpers for JSRef<'a, Node> {
     }
 }
 
+pub struct QuerySelectorIterator<'a> {
+    selectors: SelectorList,
+    iterator: TreeIterator<'a>,
+}
+
+impl<'a> QuerySelectorIterator<'a> {
+    unsafe fn new(iter: TreeIterator<'a>, selectors: SelectorList) -> QuerySelectorIterator<'a> {
+        QuerySelectorIterator {
+            selectors: selectors,
+            iterator: iter,
+        }
+    }
+}
+
+impl<'a> Iterator<JSRef<'a, Node>> for QuerySelectorIterator<'a> {
+    fn next(&mut self) -> Option<JSRef<'a, Node>> {
+        let selectors = &self.selectors;
+        // TODO(cgaebel): Is it worth it to build a bloom filter here
+        // (instead of passing `None`)? Probably.
+        self.iterator.find(|node| node.is_element() && matches(selectors, node, &mut None))
+    }
+}
+
 pub trait NodeHelpers<'a> {
     fn ancestors(self) -> AncestorIterator<'a>;
     fn children(self) -> NodeChildrenIterator<'a>;
@@ -449,6 +472,7 @@ pub trait NodeHelpers<'a> {
     fn get_content_boxes(self) -> Vec<Rect<Au>>;
 
     fn query_selector(self, selectors: DOMString) -> Fallible<Option<Temporary<Element>>>;
+    unsafe fn query_selector_iter(self, selectors: DOMString) -> Fallible<QuerySelectorIterator<'a>>;
     fn query_selector_all(self, selectors: DOMString) -> Fallible<Temporary<NodeList>>;
 
     fn remove_self(self);
@@ -719,8 +743,10 @@ impl<'a> NodeHelpers<'a> for JSRef<'a, Node> {
         Ok(None)
     }
 
-    // http://dom.spec.whatwg.org/#dom-parentnode-queryselectorall
+    /// Get an iterator over all nodes which match a set of selectors
-    fn query_selector_all(self, selectors: DOMString) -> Fallible<Temporary<NodeList>> {
+    /// Be careful not to do anything which may manipulate the DOM tree whilst iterating, otherwise 
+    /// the iterator may be invalidated
+    unsafe fn query_selector_iter(self, selectors: DOMString) -> Fallible<QuerySelectorIterator<'a>> {
         // Step 1.
         let nodes;
         let root = self.ancestors().last().unwrap_or(self.clone());
@@ -728,17 +754,25 @@ impl<'a> NodeHelpers<'a> for JSRef<'a, Node> {
             // Step 2.
             Err(()) => return Err(Syntax),
             // Step 3.
-            Ok(ref selectors) => {
+            Ok(selectors) => {
-                nodes = root.traverse_preorder().filter(
+                nodes = QuerySelectorIterator::new(root.traverse_preorder(), selectors);
-                    // TODO(cgaebel): Is it worth it to build a bloom filter here
-                    // (instead of passing `None`)? Probably.
-                    |node| node.is_element() && matches(selectors, node, &mut None)).collect()
             }
-        }
+        };
-        let window = window_from_node(self).root();
+        Ok(nodes)
-        Ok(NodeList::new_simple_list(*window, nodes))
     }
 
+    // http://dom.spec.whatwg.org/#dom-parentnode-queryselectorall
+    fn query_selector_all(self, selectors: DOMString) -> Fallible<Temporary<NodeList>> {
+        // Step 1.
+        unsafe {
+            self.query_selector_iter(selectors).map(|mut iter| {
+                let window = window_from_node(self).root();
+                NodeList::new_simple_list(*window, iter.collect())
+            })
+        }
+    }
+
+
     fn ancestors(self) -> AncestorIterator<'a> {
         AncestorIterator {
             current: self.parent_node.get().map(|node| (*node.root()).clone()),