diff options
| -rw-r--r-- | lib/puppet/simple_graph.rb | 20 | ||||
| -rwxr-xr-x | spec/unit/simple_graph.rb | 1 |
2 files changed, 11 insertions, 10 deletions
diff --git a/lib/puppet/simple_graph.rb b/lib/puppet/simple_graph.rb index bc81a6a65..b1aaf6363 100644 --- a/lib/puppet/simple_graph.rb +++ b/lib/puppet/simple_graph.rb @@ -27,7 +27,9 @@ class Puppet::SimpleGraph direction = options[:direction] || :out options[:type] ||= :vertices - return @adjacencies[direction].values.flatten if options[:type] == :edges + if options[:type] == :edges + return send(direction.to_s + "_edges") + end return @adjacencies[direction].keys.reject { |vertex| @adjacencies[direction][vertex].empty? } end @@ -39,7 +41,7 @@ class Puppet::SimpleGraph # Return all known edges. def edges - [:in, :out].collect { |dir| @adjacencies[dir].values }.flatten + in_edges + out_edges end # Test whether we share an edge with a given vertex. @@ -57,7 +59,7 @@ class Puppet::SimpleGraph # as expensive as just getting the edge list, so I've decided # to only add this method. define_method("%s_edges" % direction) do - @adjacencies[direction].values.flatten + @adjacencies[direction].values.inject([]) { |total, adjacent| total += adjacent.to_a; total } end end @@ -93,14 +95,14 @@ class Puppet::SimpleGraph # Look up the adjacencies for a given vertex. def vertex_adjacencies(direction, vertex) - @adjacencies[direction][vertex] ||= [] + @adjacencies[direction][vertex] ||= Set.new @adjacencies[direction][vertex] end end def initialize @vertices = {} - @edges = [] + @edges = Set.new end # Clear our graph. @@ -224,6 +226,7 @@ class Puppet::SimpleGraph def remove_vertex!(vertex) return nil unless vertex?(vertex) @vertices[vertex].edges.each { |edge| remove_edge!(edge) } + @edges.subtract(@vertices[vertex].edges) @vertices[vertex].clear @vertices.delete(vertex) end @@ -273,7 +276,7 @@ class Puppet::SimpleGraph end def edges - @edges.dup + @edges.to_a end # Remove an edge from our graph. @@ -281,10 +284,7 @@ class Puppet::SimpleGraph @vertices[edge.source].remove_edge(:out, edge) @vertices[edge.target].remove_edge(:in, edge) - # Here we are looking for an exact edge, so we don't want to use ==, because - # it's too darn expensive (in testing, deleting 3000 edges went from 6 seconds to - # 0.05 seconds with this change). - @edges.each_with_index { |test_edge, index| @edges.delete_at(index) and break if edge.equal?(test_edge) } + @edges.delete(edge) nil end diff --git a/spec/unit/simple_graph.rb b/spec/unit/simple_graph.rb index a5984c989..2e7bad658 100755 --- a/spec/unit/simple_graph.rb +++ b/spec/unit/simple_graph.rb @@ -141,6 +141,7 @@ describe Puppet::SimpleGraph do @graph.add_edge(one) @graph.add_edge(two) edges = @graph.edges + edges.should be_instance_of(Array) edges.length.should == 2 edges.should include(one) edges.should include(two) |