Update NodeIterator code to apply review comments.

4 files changed, 237 insertions, 190 deletions

diff --git a/components/script/dom/document.rs b/components/script/dom/document.rs
index b8ed26bbc90..ad7af5562f9 100644
--- a/components/script/dom/document.rs
+++ b/components/script/dom/document.rs
@@ -1364,7 +1364,7 @@ impl<'a> DocumentMethods for JSRef<'a, Document> {
         Range::new_with_doc(self)
     }
 
-    // https://dom.spec.whatwg.org/#dom-document-createnodeiterator
+    // https://dom.spec.whatwg.org/#dom-document-createnodeiteratorroot-whattoshow-filter
     fn CreateNodeIterator(self, root: JSRef<Node>, whatToShow: u32, filter: Option<NodeFilter>)
                         -> Temporary<NodeIterator> {
         NodeIterator::new(self, root, whatToShow, filter)
diff --git a/components/script/dom/node.rs b/components/script/dom/node.rs
index 50f4db71f8f..a2a9e159f01 100644
--- a/components/script/dom/node.rs
+++ b/components/script/dom/node.rs
@@ -426,6 +426,9 @@ pub trait NodeHelpers {
     fn child_elements(self) -> ChildElementIterator;
     fn following_siblings(self) -> NodeSiblingIterator;
     fn preceding_siblings(self) -> ReverseSiblingIterator;
+    fn following_nodes(self, root: JSRef<Node>) -> FollowingNodeIterator;
+    fn preceding_nodes(self, root: JSRef<Node>) -> PrecedingNodeIterator;
+    fn descending_last_children(self) -> LastChildIterator;
     fn is_in_doc(self) -> bool;
     fn is_inclusive_ancestor_of(self, parent: JSRef<Node>) -> bool;
     fn is_parent_of(self, child: JSRef<Node>) -> bool;
@@ -785,6 +788,26 @@ impl<'a> NodeHelpers for JSRef<'a, Node> {
         }
     }
 
+    fn following_nodes(self, root: JSRef<Node>) -> FollowingNodeIterator {
+        FollowingNodeIterator {
+            current: Some(Temporary::from_rooted(self)),
+            root: Temporary::from_rooted(root),
+        }
+    }
+
+    fn preceding_nodes(self, root: JSRef<Node>) -> PrecedingNodeIterator {
+        PrecedingNodeIterator {
+            current: Some(Temporary::from_rooted(self)),
+            root: Temporary::from_rooted(root),
+        }
+    }
+
+    fn descending_last_children(self) -> LastChildIterator {
+        LastChildIterator {
+            current: self.GetLastChild(),
+        }
+    }
+
     fn is_parent_of(self, child: JSRef<Node>) -> bool {
         match child.parent_node.get().root() {
             Some(ref parent) => parent.r() == self,
@@ -922,7 +945,6 @@ impl<'a> NodeHelpers for JSRef<'a, Node> {
         Ok(NodeList::new_simple_list(window.r(), iter))
     }
 
-
     fn ancestors(self) -> AncestorIterator {
         AncestorIterator {
             current: self.GetParentNode()
@@ -1236,6 +1258,114 @@ impl Iterator for ReverseSiblingIterator {
     }
 }
 
+pub struct FollowingNodeIterator {
+    current: Option<Temporary<Node>>,
+    root: Temporary<Node>,
+}
+
+impl Iterator for FollowingNodeIterator {
+    type Item = Temporary<Node>;
+
+    // https://dom.spec.whatwg.org/#concept-tree-following
+    fn next(&mut self) -> Option<Temporary<Node>> {
+        let current = match self.current.take() {
+            None => return None,
+            Some(current) => current,
+        };
+
+        let node = current.root();
+        if let Some(first_child) = node.r().GetFirstChild() {
+            self.current = Some(first_child);
+            return node.r().GetFirstChild()
+        }
+
+        if self.root == current {
+            self.current = None;
+            return None;
+        }
+
+        if let Some(next_sibling) = node.r().GetNextSibling() {
+            self.current = Some(next_sibling);
+            return node.r().GetNextSibling()
+        }
+
+        for ancestor in node.r().inclusive_ancestors() {
+            if self.root == ancestor {
+                break;
+            }
+            if let Some(next_sibling) = ancestor.root().r().GetNextSibling() {
+                self.current = Some(next_sibling);
+                return ancestor.root().r().GetNextSibling()
+            }
+        }
+        self.current = None;
+        return None
+    }
+}
+
+pub struct PrecedingNodeIterator {
+    current: Option<Temporary<Node>>,
+    root: Temporary<Node>,
+}
+
+impl Iterator for PrecedingNodeIterator {
+    type Item = Temporary<Node>;
+
+    // https://dom.spec.whatwg.org/#concept-tree-preceding
+    fn next(&mut self) -> Option<Temporary<Node>> {
+        let current = match self.current.take() {
+            None => return None,
+            Some(current) => current,
+        };
+
+        if self.root == current {
+            self.current = None;
+            return None
+        }
+
+        let node = current.root();
+        if let Some(previous_sibling) = node.r().GetPreviousSibling() {
+            if self.root == previous_sibling {
+                self.current = None;
+                return None
+            }
+
+            if let Some(last_child) = previous_sibling.root().r().descending_last_children().last() {
+                self.current = Some(last_child);
+                return previous_sibling.root().r().descending_last_children().last()
+            }
+
+            self.current = Some(previous_sibling);
+            return node.r().GetPreviousSibling()
+        };
+
+        if let Some(parent_node) = node.r().GetParentNode() {
+            self.current = Some(parent_node);
+            return node.r().GetParentNode()
+        }
+
+        self.current = None;
+        return None
+    }
+}
+
+pub struct LastChildIterator {
+    current: Option<Temporary<Node>>,
+}
+
+impl Iterator for LastChildIterator {
+    type Item = Temporary<Node>;
+
+    fn next(&mut self) -> Option<Temporary<Node>> {
+        let current = match self.current.take() {
+            None => return None,
+            Some(current) => current,
+        }.root();
+        self.current = current.r().GetLastChild();
+        Some(Temporary::from_rooted(current.r()))
+    }
+}
+
 pub struct AncestorIterator {
     current: Option<Temporary<Node>>,
 }
@@ -1298,7 +1428,6 @@ impl Iterator for TreeIterator {
     }
 }
 
-
 /// Specifies whether children must be recursively cloned or not.
 #[derive(Copy, Clone, PartialEq)]
 pub enum CloneChildrenFlag {
diff --git a/components/script/dom/nodeiterator.rs b/components/script/dom/nodeiterator.rs
index 531e98c724c..42c51ee8c55 100644
--- a/components/script/dom/nodeiterator.rs
+++ b/components/script/dom/nodeiterator.rs
@@ -10,10 +10,10 @@ use dom::bindings::codegen::Bindings::NodeFilterBinding::NodeFilter;
 use dom::bindings::codegen::Bindings::NodeFilterBinding::NodeFilterConstants;
 use dom::bindings::error::Fallible;
 use dom::bindings::global::GlobalRef;
-use dom::bindings::js::{JS, JSRef, MutNullableHeap, OptionalRootable, Temporary, Rootable};
+use dom::bindings::js::{JS, JSRef, MutHeap, OptionalRootable, Temporary, Rootable, RootedReference};
 use dom::bindings::utils::{Reflector, reflect_dom_object};
 use dom::document::{Document, DocumentHelpers};
-use dom::node::Node;
+use dom::node::{Node, NodeHelpers};
 
 use std::cell::Cell;
 
@@ -21,7 +21,7 @@ use std::cell::Cell;
 pub struct NodeIterator {
     reflector_: Reflector,
     root_node: JS<Node>,
-    reference_node: MutNullableHeap<JS<Node>>,
+    reference_node: MutHeap<JS<Node>>,
     pointer_before_reference_node: Cell<bool>,
     what_to_show: u32,
     filter: Filter,
@@ -29,12 +29,12 @@ pub struct NodeIterator {
 
 impl NodeIterator {
     fn new_inherited(root_node: JSRef<Node>,
-                         what_to_show: u32,
-                         filter: Filter) -> NodeIterator {
+                     what_to_show: u32,
+                     filter: Filter) -> NodeIterator {
         NodeIterator {
             reflector_: Reflector::new(),
             root_node: JS::from_rooted(root_node),
-            reference_node: MutNullableHeap::new(Some(JS::from_rooted(root_node))),
+            reference_node: MutHeap::new(JS::from_rooted(root_node)),
             pointer_before_reference_node: Cell::new(true),
             what_to_show: what_to_show,
             filter: filter
@@ -57,7 +57,7 @@ impl NodeIterator {
                node_filter: Option<NodeFilter>) -> Temporary<NodeIterator> {
         let filter = match node_filter {
             None => Filter::None,
-            Some(jsfilter) => Filter::JS(jsfilter)
+            Some(jsfilter) => Filter::Callback(jsfilter)
         };
         NodeIterator::new_with_filter(document, root_node, what_to_show, filter)
     }
@@ -78,28 +78,111 @@ impl<'a> NodeIteratorMethods for JSRef<'a, NodeIterator> {
     fn GetFilter(self) -> Option<NodeFilter> {
         match self.filter {
             Filter::None => None,
-            Filter::JS(nf) => Some(nf),
+            Filter::Callback(nf) => Some(nf),
             Filter::Native(_) => panic!("Cannot convert native node filter to DOM NodeFilter")
         }
     }
 
     // https://dom.spec.whatwg.org/#dom-nodeiterator-referencenode
-    fn GetReferenceNode(self) -> Option<Temporary<Node>> {
-        self.reference_node.get().map(Temporary::from_rooted)
+    fn ReferenceNode(self) -> Temporary<Node> {
+        Temporary::from_rooted(self.reference_node.get())
     }
 
+    // https://dom.spec.whatwg.org/#dom-nodeiterator-pointerbeforereferencenode
     fn PointerBeforeReferenceNode(self) -> bool {
         self.pointer_before_reference_node.get()
     }
 
+    // https://dom.spec.whatwg.org/#dom-nodeiterator-nextnode
+    fn NextNode(self) -> Fallible<Option<Temporary<Node>>> {
+        // https://dom.spec.whatwg.org/#concept-NodeIterator-traverse
+        // Step 1.
+        let node = self.reference_node.get().root();
+
+        // Step 2.
+        let mut before_node = self.pointer_before_reference_node.get();
+
+        // Step 3-1.
+        if before_node {
+            before_node = false;
+
+            // Step 3-2.
+            let result = try!(self.accept_node(node.r()));
+
+            // Step 3-3.
+            if result == NodeFilterConstants::FILTER_ACCEPT {
+                // Step 4.
+                self.reference_node.set(JS::from_rooted(node.r()));
+                self.pointer_before_reference_node.set(before_node);
+
+                return Ok(Some(Temporary::from_rooted(node.r())));
+            }
+        }
+
+        // Step 3-1.
+        for following_node in node.r().following_nodes(self.root_node.root().r()) {
+            let following_node = following_node.root();
+
+            // Step 3-2.
+            let result = try!(self.accept_node(following_node.r()));
+
+            // Step 3-3.
+            if result == NodeFilterConstants::FILTER_ACCEPT {
+                // Step 4.
+                self.reference_node.set(JS::from_rooted(following_node.r()));
+                self.pointer_before_reference_node.set(before_node);
+
+                return Ok(Some(Temporary::from_rooted(following_node.r())));
+            }
+        }
+
+        return Ok(None);
+    }
+
     // https://dom.spec.whatwg.org/#dom-nodeiterator-previousnode
     fn PreviousNode(self) -> Fallible<Option<Temporary<Node>>> {
-        self.prev_node()
-    }
+        // https://dom.spec.whatwg.org/#concept-NodeIterator-traverse
+        // Step 1.
+        let node = self.reference_node.get().root();
 
-    // https://dom.spec.whatwg.org/#dom-nodeiterator-nextnode
-    fn NextNode(self) -> Fallible<Option<Temporary<Node>>> {
-        self.next_node()
+        // Step 2.
+        let mut before_node = self.pointer_before_reference_node.get();
+
+        // Step 3-1.
+        if !before_node {
+            before_node = true;
+
+            // Step 3-2.
+            let result = try!(self.accept_node(node.r()));
+
+            // Step 3-3.
+            if result == NodeFilterConstants::FILTER_ACCEPT {
+                // Step 4.
+                self.reference_node.set(JS::from_rooted(node.r()));
+                self.pointer_before_reference_node.set(before_node);
+
+                return Ok(Some(Temporary::from_rooted(node.r())));
+            }
+        }
+
+        // Step 3-1.
+        for preceding_node in node.r().preceding_nodes(self.root_node.root().r()) {
+            let preceding_node = preceding_node.root();
+
+            // Step 3-2.
+            let result = try!(self.accept_node(preceding_node.r()));
+
+            // Step 3-3.
+            if result == NodeFilterConstants::FILTER_ACCEPT {
+                // Step 4.
+                self.reference_node.set(JS::from_rooted(preceding_node.r()));
+                self.pointer_before_reference_node.set(before_node);
+
+                return Ok(Some(Temporary::from_rooted(preceding_node.r())));
+            }
+        }
+
+        return Ok(None);
     }
 
     // https://dom.spec.whatwg.org/#dom-nodeiterator-detach
@@ -109,103 +192,24 @@ impl<'a> NodeIteratorMethods for JSRef<'a, NodeIterator> {
 }
 
 trait PrivateNodeIteratorHelpers {
-    fn is_root_node(self, node: JSRef<Node>) -> bool;
-    fn first_following_node_not_following_root(self, node: JSRef<Node>) -> Option<Temporary<Node>>;
-    fn first_preceding_node_not_preceding_root(self, node: JSRef<Node>) -> Option<Temporary<Node>>;
     fn accept_node(self, node: JSRef<Node>) -> Fallible<u16>;
+    fn is_root_node(self, node: JSRef<Node>) -> bool;
 }
 
 impl<'a> PrivateNodeIteratorHelpers for JSRef<'a, NodeIterator> {
-
-    // https://dom.spec.whatwg.org/#concept-tree-following
-    fn first_following_node_not_following_root(self, node: JSRef<Node>)
-                                               -> Option<Temporary<Node>> {
-        // "An object A is following an object B if A and B are in the same tree
-        //  and A comes after B in tree order."
-        match node.GetFirstChild() {
-            Some (child) => { return Some(Temporary::from_rooted(child.root().get_unsound_ref_forever())) },
-            None => {}
-        }
-        match node.GetNextSibling() {
-            None => {
-                let mut candidate = node;
-                while !self.is_root_node(candidate) && candidate.GetNextSibling().is_none() {
-                    match candidate.GetParentNode() {
-                        None =>
-                            // This can happen if the user set the current node to somewhere
-                            // outside of the tree rooted at the original root.
-                            return None,
-                        Some(n) => { candidate = n.root().get_unsound_ref_forever(); }
-                    }
-                }
-                if self.is_root_node(candidate) {
-                    None
-                } else {
-                    candidate.GetNextSibling()
-                }
-            },
-            it => {
-                if self.is_root_node(node) {
-                    return None
-                } else {
-                    it
-                }
-            }
-        }
-    }
-
-    // https://dom.spec.whatwg.org/#concept-tree-preceding
-    fn first_preceding_node_not_preceding_root(self, node: JSRef<Node>)
-                                               -> Option<Temporary<Node>> {
-        // "An object A is preceding an object B if A and B are in the same tree
-        //  and A comes before B in tree order."
-        match node.GetPreviousSibling() {
-            None => {
-                let mut candidate = node;
-                while !self.is_root_node(candidate) && candidate.GetPreviousSibling().is_none() {
-                    match candidate.GetParentNode() {
-                        None =>
-                            // This can happen if the user set the current node to somewhere
-                            // outside of the tree rooted at the original root.
-                            return None,
-                        Some(n) => { candidate = n.root().get_unsound_ref_forever(); }
-                    }
-                }
-                if self.is_root_node(candidate) {
-                    None
-                } else {
-                    candidate.GetPreviousSibling()
-                }
-            },
-            it => {
-                let candidate = node;
-                if self.is_root_node(candidate) {
-                    return None
-                } else {
-                    it
-                }
-            }
-        }
-    }
-
     // https://dom.spec.whatwg.org/#concept-node-filter
     fn accept_node(self, node: JSRef<Node>) -> Fallible<u16> {
-        // "To filter node run these steps:"
-        // "1. Let n be node's nodeType attribute value minus 1."
+        // Step 1.
         let n = node.NodeType() - 1;
-        // "2. If the nth bit (where 0 is the least significant bit) of whatToShow is not set,
-        //     return FILTER_SKIP."
+        // Step 2.
         if (self.what_to_show & (1 << n)) == 0 {
             return Ok(NodeFilterConstants::FILTER_SKIP)
         }
-        // "3. If filter is null, return FILTER_ACCEPT."
-        // "4. Let result be the return value of invoking filter."
-        // "5. If an exception was thrown, re-throw the exception."
-        // "6. Return result."
+        // Step 3-5.
         match self.filter {
             Filter::None => Ok(NodeFilterConstants::FILTER_ACCEPT),
             Filter::Native(f) => Ok((f)(node)),
-            Filter::JS(callback) => callback.AcceptNode_(self, node, Rethrow)
+            Filter::Callback(callback) => callback.AcceptNode_(self, node, Rethrow)
         }
     }
 
@@ -214,96 +218,10 @@ impl<'a> PrivateNodeIteratorHelpers for JSRef<'a, NodeIterator> {
     }
 }
 
-pub trait NodeIteratorHelpers {
-    fn next_node(self) -> Fallible<Option<Temporary<Node>>>;
-    fn prev_node(self) -> Fallible<Option<Temporary<Node>>>;
-    fn traverse(self, direction: Direction) -> Fallible<Option<Temporary<Node>>>;
-}
-
-impl<'a> NodeIteratorHelpers for JSRef<'a, NodeIterator> {
-    // https://dom.spec.whatwg.org/#dom-nodeiterator-nextnode
-    fn next_node(self) -> Fallible<Option<Temporary<Node>>> {
-        self.traverse(Direction::Next)
-    }
-
-    // https://dom.spec.whatwg.org/#dom-nodeiterator-previousnode
-    fn prev_node(self) -> Fallible<Option<Temporary<Node>>> {
-        self.traverse(Direction::Previous)
-    }
-
-    fn traverse(self, direction: Direction) -> Fallible<Option<Temporary<Node>>> {
-        // "1. Let node be the value of the referenceNode attribute."
-        let mut node = self.reference_node.get();
-        // "2. Let before node be the value of the pointerBeforeReferenceNode attribute."
-        let mut before_node = self.pointer_before_reference_node.get();
-        // "3. Run these substeps:
-        loop {
-            match direction {
-                // "1. If direction is next"
-                Direction::Next => {
-                    // "If before node is false, let node be the first node following node in the iterator collection.
-                    //  If there is no such node return null. If before node is true, set it to false."
-                    if !before_node {
-                        match node {
-                            None => return Ok(None),
-                            Some(n) => {
-                                match self.first_following_node_not_following_root(n.root().get_unsound_ref_forever()) {
-                                    None => return Ok(None),
-                                    Some(n) => node = Some(JS::from_rooted(n))
-                                }
-                            }
-                        }
-                    }
-                    else {
-                        before_node = false;
-                    }
-                }
-                // "If direction is previous"
-                Direction::Previous => {
-                    // "If before node is true, let node be the first node preceding node in the iterator collection.
-                    //  If there is no such node return null. If before node is false, set it to true.
-                    if before_node {
-                        match node {
-                            None => return Ok(None),
-                            Some(n) => {
-                                match self.first_preceding_node_not_preceding_root(n.root().get_unsound_ref_forever()) {
-                                    None => return Ok(None),
-                                    Some(n) => node = Some(JS::from_rooted(n))
-                                }
-                            }
-                        }
-                    }
-                    else {
-                        before_node = true;
-                    }
-                }
-            }
-            // "2. Filter node and let result be the return value."
-            let result = try!(self.accept_node(node.unwrap().root().get_unsound_ref_forever()));
-
-            // "3. If result is FILTER_ACCEPT, go to the next step in the overall set of steps. Otherwise, run these substeps again."
-            match result {
-                NodeFilterConstants::FILTER_ACCEPT => break,
-                _ => {}
-            }
-
-        }
-        // "4. Set the referenceNode attribute to node, set the pointerBeforeReferenceNode attribute to before node, and return node."
-        self.reference_node.set(Some(JS::from_rooted(node.unwrap())));
-        self.pointer_before_reference_node.set(before_node);
-
-        Ok(Some(Temporary::from_rooted(node.unwrap())))
-    }
-}
-
-pub enum Direction {
-    Next,
-    Previous
-}
 
 #[jstraceable]
 pub enum Filter {
     None,
     Native(fn (node: JSRef<Node>) -> u16),
-    JS(NodeFilter)
+    Callback(NodeFilter)
 }
diff --git a/components/script/dom/webidls/NodeIterator.webidl b/components/script/dom/webidls/NodeIterator.webidl
index bd5970747ca..73a28040715 100644
--- a/components/script/dom/webidls/NodeIterator.webidl
+++ b/components/script/dom/webidls/NodeIterator.webidl
@@ -15,7 +15,7 @@ interface NodeIterator {
   [Constant]
   readonly attribute Node root;
   [Pure]
-  readonly attribute Node? referenceNode;
+  readonly attribute Node referenceNode;
   [Pure]
   readonly attribute boolean pointerBeforeReferenceNode;
   [Constant]