1 files changed, 23 insertions, 12 deletions
diff --git a/components/script/dom/treewalker.rs b/components/script/dom/treewalker.rs
index de5219a608f..89ef5527e4e 100644
--- a/components/script/dom/treewalker.rs
+++ b/components/script/dom/treewalker.rs
@@ -8,12 +8,13 @@ use dom::bindings::codegen::Bindings::NodeFilterBinding::NodeFilter;
 use dom::bindings::codegen::Bindings::NodeFilterBinding::NodeFilterConstants;
 use dom::bindings::codegen::Bindings::TreeWalkerBinding;
 use dom::bindings::codegen::Bindings::TreeWalkerBinding::TreeWalkerMethods;
-use dom::bindings::error::Fallible;
+use dom::bindings::error::{Error, Fallible};
 use dom::bindings::reflector::{Reflector, reflect_dom_object};
 use dom::bindings::root::{Dom, DomRoot, MutDom};
 use dom::document::Document;
 use dom::node::Node;
 use dom_struct::dom_struct;
+use std::cell::Cell;
 use std::rc::Rc;
 
 // https://dom.spec.whatwg.org/#interface-treewalker
@@ -24,7 +25,8 @@ pub struct TreeWalker {
     current_node: MutDom<Node>,
     what_to_show: u32,
     #[ignore_malloc_size_of = "function pointers and Rc<T> are hard"]
-    filter: Filter
+    filter: Filter,
+    active: Cell<bool>,
 }
 
 impl TreeWalker {
@@ -36,7 +38,8 @@ impl TreeWalker {
             root_node: Dom::from_ref(root_node),
             current_node: MutDom::new(root_node),
             what_to_show: what_to_show,
-            filter: filter
+            filter: filter,
+            active: Cell::new(false),
         }
     }
 
@@ -415,22 +418,30 @@ impl TreeWalker {
 
     // https://dom.spec.whatwg.org/#concept-node-filter
     fn accept_node(&self, node: &Node) -> Fallible<u16> {
-        // "To filter node run these steps:"
-        // "1. Let n be node's nodeType attribute value minus 1."
+        // Step 1.
+        if self.active.get() {
+            return Err(Error::InvalidState);
+        }
+        // Step 2.
         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 3.
         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."
         match self.filter {
+            // Step 4.
             Filter::None => Ok(NodeFilterConstants::FILTER_ACCEPT),
             Filter::Native(f) => Ok((f)(node)),
-            Filter::Dom(ref callback) => callback.AcceptNode_(self, node, Rethrow)
+            Filter::Dom(ref callback) => {
+                // Step 5.
+                self.active.set(true);
+                // Step 6.
+                let result = callback.AcceptNode_(self, node, Rethrow);
+                // Step 7.
+                self.active.set(false);
+                // Step 8.
+                result
+            },
         }
     }