From f4d606eabbdeaf64306d3c4b42bc777f5245a2ef Mon Sep 17 00:00:00 2001
From: Emad Eid <eeid@comcast.net>
Date: Tue, 10 Mar 2015 17:25:43 -0400
Subject: [PATCH] styling issues and support the filter as a function. IE
 requires it.

---
 src/ShadowDOM/wrappers/Document.js  | 18 +++++++++++----
 src/ShadowDOM/wrappers/Selection.js |  8 ++++---
 tests/ShadowDOM/js/TreeWalker.js    | 34 +++++++++++++++++++++++++----
 3 files changed, 49 insertions(+), 11 deletions(-)

diff --git a/src/ShadowDOM/wrappers/Document.js b/src/ShadowDOM/wrappers/Document.js
index e25e25a..7dc3875 100644
--- a/src/ShadowDOM/wrappers/Document.js
+++ b/src/ShadowDOM/wrappers/Document.js
@@ -121,11 +121,21 @@
 
   var originalCreateTreeWalker = document.createTreeWalker;
   var TreeWalkerWrapper = scope.wrappers.TreeWalker;
-  Document.prototype.createTreeWalker = function(root, whatToShow, filter, expandEntityReferences ) {
+  Document.prototype.createTreeWalker = function(root,whatToShow,
+                                                 filter,expandEntityReferences ) {
 
-    var newFilter;
-    if (filter && filter.acceptNode && typeof filter.acceptNode === 'function'){
-      newFilter = { acceptNode:function(node) { return filter.acceptNode(wrap(node)); }  }
+    var newFilter = null; // IE does not like undefined.
+
+    // Support filter as a function or object with function defined as acceptNode.
+    // IE supports filter as a function only. Chrome and FF support both formats.
+    if (filter){
+      if (filter.acceptNode && typeof filter.acceptNode === 'function'){
+        newFilter = {
+          acceptNode:function(node) { return filter.acceptNode(wrap(node)); }
+        };
+      }else if (typeof filter === 'function'){
+        newFilter = function(node) { return filter(wrap(node)); }
+      }
     }
 
     return new TreeWalkerWrapper(originalCreateTreeWalker.call(unwrap(this), unwrap(root),
diff --git a/src/ShadowDOM/wrappers/Selection.js b/src/ShadowDOM/wrappers/Selection.js
index 556045f..3714f8d 100644
--- a/src/ShadowDOM/wrappers/Selection.js
+++ b/src/ShadowDOM/wrappers/Selection.js
@@ -53,10 +53,12 @@
     }
   };
 
-  // Not all browsers support Selection.extend. Some versions of IE do not support extend, code that checks
-  // if extend exists in the Selection would fail the test if we define extend on the wrapper and it does not exist in
+  // Not all browsers support Selection.extend. IE does not support extend.
+  // https://msdn.microsoft.com/en-us/library/ie/ms535869%28v=vs.85%29.aspx
+  // Code that checks if extend exists in the Selection would
+  // fail the test if we define extend on the wrapper and it does not exist in
   // the browser Selection object.
-  if (window.getSelection().extend){
+  if (OriginalSelection.prototype.extend) {
     Selection.prototype.extend = function(node, offset) {
       unsafeUnwrap(this).extend(unwrapIfNeeded(node), offset);
     };
diff --git a/tests/ShadowDOM/js/TreeWalker.js b/tests/ShadowDOM/js/TreeWalker.js
index ca3c10d..a924ee5 100644
--- a/tests/ShadowDOM/js/TreeWalker.js
+++ b/tests/ShadowDOM/js/TreeWalker.js
@@ -73,13 +73,20 @@ suite('TreeWalker', function() {
     assert(isWrapper(treeWalker.nextNode()));
     assert(isWrapper(treeWalker.parentNode()));
     assert(isWrapper(treeWalker.firstChild()));
-    assert(isWrapper(treeWalker.lastChild()) || isWrapper(treeWalker.lastChild())===null);
-    assert(isWrapper(treeWalker.previousSibling()) || isWrapper(treeWalker.previousSibling())===null);
+    assert(isWrapper(treeWalker.lastChild()) ||
+                          isWrapper(treeWalker.lastChild())===null);
+    assert(isWrapper(treeWalker.previousSibling()) ||
+                          isWrapper(treeWalker.previousSibling())===null);
     assert(isWrapper(treeWalker.previousNode()));
 
   });
 
-  test('createTreeWalker with filter', function() {
+  test('createTreeWalker with filter as object with acceptNode function', function() {
+
+    //  NodeFilter.acceptNode as filter does not work in IE.
+    // https://dom.spec.whatwg.org/#nodefilter
+    if (/Trident/.test(navigator.userAgent))
+      return;
 
     var treeWalker = document.createTreeWalker(childDiv, NodeFilter.SHOW_ELEMENT, {
       acceptNode:function(node){
@@ -92,7 +99,26 @@ suite('TreeWalker', function() {
     });
     
     assert.isNotNull(treeWalker.filter);
-    assert.instanceOf(treeWalker.filter,NodeFilter);
+    // in FF and IE treeWalker.filter is just a js object
+    //assert.instanceOf(treeWalker.filter, NodeFilter);
+
+    // we should have one node only.
+    assert.isNotNull(treeWalker.nextNode());
+    assert.isNull(treeWalker.nextNode());
+
+  });
+
+  test('createTreeWalker with filter as a function (works under IE and others)', function() {
+
+    var treeWalker = document.createTreeWalker(childDiv, NodeFilter.SHOW_ELEMENT, function(node){
+        assert(isWrapper(node));
+        if (node.hasAttribute("title") && node.getAttribute("title")==='a'){
+          return NodeFilter.FILTER_ACCEPT;
+        }
+        return false;
+      });
+
+    assert.isNotNull(treeWalker.filter);
 
     // we should have one node only.
     assert.isNotNull(treeWalker.nextNode());