treeToArray: Use a filter function argument instead of includeRoot

Author: Joris-van-der-Wel

lib/SymbolTree.js

@@ -7,6 +7,10 @@
 
 const SymbolTreeNode = require('./SymbolTreeNode');
 
+function returnTrue() {
+        return true;
+}
+
 class SymbolTree {
 
         /**
@@ -262,20 +266,27 @@ class SymbolTree {
          * @memberOf module:symbol-tree#
          * @param {Object} root
          * @param {Object[]} [array=[]]
-         * @param {Boolean} [includeRoot=true] Include the given `root` object in the array?
+         * @param {Function} [filter] Function to test each object before it is added to the array.
+         *                            Invoked with arguments (object). Should return `true` if an object
+         *                            is to be included.
+         * @param {*} [thisArg] Value to use as `this` when executing `filter`.
          * @return {Object[]}
          */
-        treeToArray(root, array, includeRoot) {
+        treeToArray(root, array, filter, thisArg) {
                 if (!array) {
                         array = [];
                 }
 
-                let object = includeRoot || includeRoot === undefined
-                        ? root
-                        : this.following(root, root);
+                if (!filter) {
+                        filter = returnTrue;
+                }
+
+                let object = root;
 
                 while (object) {
-                        array.push(object);
+                        if (filter.call(thisArg, object)) {
+                                array.push(object);
+                        }
                         object = this.following(object, root);
                 }
 

test/SymbolTree.js

@@ -654,11 +654,45 @@ test('treeToArray', function(t) {
         tree.insertAfter(b, a);
 
         t.deepEqual([a, aa, ab, aba, abaa], tree.treeToArray(a));
-        t.deepEqual([aa, ab, aba, abaa], tree.treeToArray(a, null, false));
 
         const arr = ['a', 5];
         tree.treeToArray(a, arr);
         t.deepEqual(['a', 5, a, aa, ab, aba, abaa], arr);
 
         t.end();
 });
+
+test('treeToArray with filter', function(t) {
+        const tree = new SymbolTree();
+        const a = {};
+        const aa = {};
+        const ab = {};
+        const aba = {};
+        const abaa = {};
+        const b = {};
+
+        tree.insertLast(aa, a);
+        tree.insertLast(ab, a);
+        tree.insertLast(aba, ab);
+        tree.insertLast(abaa, aba);
+        tree.insertAfter(b, a);
+
+        const filter = function(object) {
+                t.equal(this, undefined);
+
+                return object !== a && object !== aba;
+        };
+
+        t.deepEqual([aa, ab, abaa], tree.treeToArray(a, null, filter));
+
+        const thisArg = {foo: 'bar'};
+        const filterThis = function(object) {
+                t.equal(this, thisArg);
+
+                return object !== a && object !== aba;
+        };
+
+        t.deepEqual([aa, ab, abaa], tree.treeToArray(a, null, filterThis, thisArg));
+
+        t.end();
+});