Add tests, benchmark + misc fixes to filterNodes

Author: cpettitt-linkedin

lib/graph.js

@@ -270,27 +270,50 @@ Graph.prototype.neighbors = function(v) {
   }
 };
 
-Graph.prototype.eachNode = function(func) {
-  for (var id in this._nodes) {
-    func(id, this._nodes[id]);
-  }
-};
-
 Graph.prototype.filterNodes = function(filter) {
-  var copy = new this.constructor();
-  copy.graph(this.graph());
-  this.eachNode(function(u, value) {
-    if (filter(u)) {
-      copy.setNode(u, value);
-    }
+  var copy = new this.constructor({
+    directed: this._isDirected,
+    multigraph: this._isMultigraph,
+    compound: this._isCompound
   });
-  this.eachEdge(function(id, v, w, value) {
-    if (copy.hasNode(v) && copy.hasNode(w)) {
-      copy.setEdge(v, w, value);
+
+  copy.setGraph(this.graph());
+
+  _.each(this._nodes, function(value, v) {
+    if (filter(v)) {
+      copy.setNode(v, value);
     }
-  });
+  }, this);
+
+  _.each(this._edgeObjs, function(e) {
+    if (copy.hasNode(e.v) && copy.hasNode(e.w)) {
+      copy.setEdge(e, this.edge(e));
+    }
+  }, this);
+
+  var self = this;
+  var parents = {};
+  function findParent(v) {
+    var parent = self.parent(v);
+    if (parent === undefined || copy.hasNode(parent)) {
+      parents[v] = parent;
+      return parent;
+    } else if (parent in parents) {
+      return parents[parent];
+    } else {
+      return findParent(parent);
+    }
+  }
+
+  if (this._isCompound) {
+    _.each(copy.nodes(), function(v) {
+      copy.setParent(v, findParent(v));
+    });
+  }
+
   return copy;
 };
+
 /* === Edge functions ========== */
 
 Graph.prototype.setDefaultEdgeLabel = function(newDefault) {
@@ -451,14 +474,6 @@ Graph.prototype.nodeEdges = function(v, w) {
   }
 };
 
-Graph.prototype.eachEdge = function(func) {
-  var id, edge;
-  for (id in this._edgeObjs) {
-    edge=this._edgeObjs[id];
-    func(id, edge.v, edge.w, this._edgeLabels[id]);
-  }
-};
-
 function incrementOrInitEntry(map, k) {
   if (_.has(map, k)) {
     map[k]++;

src/bench.js

@@ -88,6 +88,14 @@ NODE_SIZES.forEach(function(size) {
     g.sinks();
   });
 
+  runBenchmark("filterNodes all" + nameSuffix, function() {
+    g.filterNodes(function() { return true; });
+  });
+
+  runBenchmark("filterNodes none" + nameSuffix, function() {
+    g.filterNodes(function() { return false; });
+  });
+
   runBenchmark("setNode" + nameSuffix, function() {
     g.setNode("key", "label");
   });

test/graph-test.js

@@ -89,6 +89,93 @@ describe("Graph", function() {
     });
   });
 
+  describe("filterNodes", function() {
+    it("returns an identical graph when the filter selects everything", function() {
+      g.setGraph("graph label");
+      g.setNode("a", 123);
+      g.setPath(["a", "b", "c"]);
+      g.setEdge("a", "c", 456);
+      var g2 = g.filterNodes(function() { return true; });
+      expect(_.sortBy(g2.nodes())).eqls(["a", "b", "c"]);
+      expect(_.sortBy(g2.successors("a"))).eqls(["b", "c"]);
+      expect(_.sortBy(g2.successors("b"))).eqls(["c"]);
+      expect(g2.node("a")).eqls(123);
+      expect(g2.edge("a", "c")).eqls(456);
+      expect(g2.graph()).eqls("graph label");
+    });
+
+    it("returns an empty graph when the filter selects nothing", function() {
+      g.setPath(["a", "b", "c"]);
+      var g2 = g.filterNodes(function() { return false; });
+      expect(g2.nodes()).eqls([]);
+      expect(g2.edges()).eqls([]);
+    });
+
+    it("only includes nodes for which the filter returns true", function() {
+      g.setNodes(["a", "b"]);
+      var g2 = g.filterNodes(function(v) { return v === "a"; });
+      expect(g2.nodes()).eqls(["a"]);
+    });
+
+    it("removes edges that are connected to removed nodes", function() {
+      g.setEdge("a", "b");
+      var g2 = g.filterNodes(function(v) { return v === "a"; });
+      expect(_.sortBy(g2.nodes())).eqls(["a"]);
+      expect(g2.edges()).eqls([]);
+    });
+
+    it("preserves the directed option", function() {
+      g = new Graph({ directed: true });
+      expect(g.filterNodes(function() { return true; }).isDirected()).to.be.true;
+
+      g = new Graph({ directed: false });
+      expect(g.filterNodes(function() { return true; }).isDirected()).to.be.false;
+    });
+
+    it("preserves the multigraph option", function() {
+      g = new Graph({ multigraph: true });
+      expect(g.filterNodes(function() { return true; }).isMultigraph()).to.be.true;
+
+      g = new Graph({ multigraph: false });
+      expect(g.filterNodes(function() { return true; }).isMultigraph()).to.be.false;
+    });
+
+    it("preserves the compound option", function() {
+      g = new Graph({ compound: true });
+      expect(g.filterNodes(function() { return true; }).isCompound()).to.be.true;
+
+      g = new Graph({ compound: false });
+      expect(g.filterNodes(function() { return true; }).isCompound()).to.be.false;
+    });
+
+    it("includes subgraphs", function() {
+      g = new Graph({ compound: true });
+      g.setParent("a", "parent");
+
+      var g2 = g.filterNodes(function() { return true; });
+      expect(g2.parent("a")).eqls("parent");
+    });
+
+    it("includes multi-level subgraphs", function() {
+      g = new Graph({ compound: true });
+      g.setParent("a", "parent");
+      g.setParent("parent", "root");
+
+      var g2 = g.filterNodes(function() { return true; });
+      expect(g2.parent("a")).eqls("parent");
+      expect(g2.parent("parent")).eqls("root");
+    });
+
+    it("promotes a node to a higher subgraph if its parent is not included", function() {
+      g = new Graph({ compound: true });
+      g.setParent("a", "parent");
+      g.setParent("parent", "root");
+
+      var g2 = g.filterNodes(function(v) { return v !== "parent"; });
+      expect(g2.parent("a")).eqls("root");
+    });
+  });
+
   describe("setNodes", function() {
     it("creates multiple nodes", function() {
       g.setNodes(["a", "b", "c"]);