Merge branch 'ticket/next/4590' into next

Conflicts:
	lib/puppet/util/monkey_patches.rb
            -- two unrelated additions had been made, kept them both.

6 files changed, 481 insertions, 232 deletions

diff --git a/lib/puppet/simple_graph.rb b/lib/puppet/simple_graph.rb
index 55b39fadf..6d153b46d 100644
--- a/lib/puppet/simple_graph.rb
+++ b/lib/puppet/simple_graph.rb
@@ -7,123 +7,45 @@ require 'set'
 
 # A hopefully-faster graph class to replace the use of GRATR.
 class Puppet::SimpleGraph
-  # An internal class for handling a vertex's edges.
-  class VertexWrapper
-    attr_accessor :in, :out, :vertex
-
-    # Remove all references to everything.
-    def clear
-      @adjacencies[:in].clear
-      @adjacencies[:out].clear
-      @vertex = nil
-    end
-
-    def initialize(vertex)
-      @vertex = vertex
-      @adjacencies = {:in => {}, :out => {}}
-    end
-
-    # Find adjacent vertices or edges.
-    def adjacent(options)
-      direction = options[:direction] || :out
-      options[:type] ||= :vertices
-
-      return send(direction.to_s + "_edges") if options[:type] == :edges
-
-      @adjacencies[direction].keys.reject { |vertex| @adjacencies[direction][vertex].empty? }
-    end
-
-    # Add an edge to our list.
-    def add_edge(direction, edge)
-      opposite_adjacencies(direction, edge) << edge
-    end
-
-    # Return all known edges.
-    def edges
-      in_edges + out_edges
-    end
-
-    # Test whether we share an edge with a given vertex.
-    def has_edge?(direction, vertex)
-      return(vertex_adjacencies(direction, vertex).length > 0 ? true : false)
-    end
-
-    # Create methods for returning the degree and edges.
-    [:in, :out].each do |direction|
-      # LAK:NOTE If you decide to create methods for directly
-      # testing the degree, you'll have to get the values and flatten
-      # the results -- you might have duplicate edges, which can give
-      # a false impression of what the degree is.  That's just
-      # as expensive as just getting the edge list, so I've decided
-      # to only add this method.
-      define_method("#{direction}_edges") do
-        @adjacencies[direction].values.inject([]) { |total, adjacent| total += adjacent.to_a; total }
-      end
-    end
-
-    # The other vertex in the edge.
-    def other_vertex(direction, edge)
-      case direction
-      when :in; edge.source
-      else
-        edge.target
-      end
-    end
-
-    # Remove an edge from our list.  Assumes that we've already checked
-    # that the edge is valid.
-    def remove_edge(direction, edge)
-      opposite_adjacencies(direction, edge).delete(edge)
-    end
-
-    def to_s
-      vertex.to_s
-    end
-
-    private
-
-    # These methods exist so we don't need a Hash with a default proc.
-
-    # Look up the adjacencies for a vertex at the other end of an
-    # edge.
-    def opposite_adjacencies(direction, edge)
-      opposite_vertex = other_vertex(direction, edge)
-      vertex_adjacencies(direction, opposite_vertex)
-    end
-
-    # Look up the adjacencies for a given vertex.
-    def vertex_adjacencies(direction, vertex)
-      @adjacencies[direction][vertex] ||= Set.new
-      @adjacencies[direction][vertex]
-    end
-  end
-
+  #
+  # All public methods of this class must maintain (assume ^ ensure) the following invariants, where "=~=" means
+  # equiv. up to order:
+  #
+  #   @in_to.keys =~= @out_to.keys =~= all vertices
+  #   @in_to.values.collect { |x| x.values }.flatten =~= @out_from.values.collect { |x| x.values }.flatten =~= all edges
+  #   @in_to[v1][v2] =~= @out_from[v2][v1] =~= all edges from v1 to v2
+  #   @in_to   [v].keys =~= vertices with edges leading to   v
+  #   @out_from[v].keys =~= vertices with edges leading from v
+  #   no operation may shed reference loops (for gc)
+  #   recursive operation must scale with the depth of the spanning trees, or better (e.g. no recursion over the set
+  #       of all vertices, etc.)
+  #
+  # This class is intended to be used with DAGs.  However, if the
+  # graph has a cycle, it will not cause non-termination of any of the
+  # algorithms.  The topsort method detects and reports cycles.
+  #
   def initialize
-    @vertices = {}
-    @edges = []
+    @in_to = {}
+    @out_from = {}
+    @upstream_from = {}
+    @downstream_from = {}
   end
 
   # Clear our graph.
   def clear
-    @vertices.each { |vertex, wrapper| wrapper.clear }
-    @vertices.clear
-    @edges.clear
-  end
-
-  # Which resources a given resource depends upon.
-  def dependents(resource)
-    tree_from_vertex(resource).keys
+    @in_to.clear
+    @out_from.clear
+    @upstream_from.clear
+    @downstream_from.clear
   end
 
   # Which resources depend upon the given resource.
   def dependencies(resource)
-    # Cache the reversal graph, because it's somewhat expensive
-    # to create.
-    @reversal ||= reversal
-    # Strangely, it's significantly faster to search a reversed
-    # tree in the :out direction than to search a normal tree
-    # in the :in direction.
-    @reversal.tree_from_vertex(resource, :out).keys
+    vertex?(resource) ? upstream_from_vertex(resource).keys : []
+  end
+
+  def dependents(resource)
+    vertex?(resource) ? downstream_from_vertex(resource).keys : []
   end
 
   # Whether our graph is directed.  Always true.  Used to produce dot files.
@@ -133,8 +55,7 @@ class Puppet::SimpleGraph
 
   # Determine all of the leaf nodes below a given vertex.
   def leaves(vertex, direction = :out)
-    tree = tree_from_vertex(vertex, direction)
-    l = tree.keys.find_all { |c| adjacent(c, :direction => direction).empty? }
+    tree_from_vertex(vertex, direction).keys.find_all { |c| adjacent(c, :direction => direction).empty? }
   end
 
   # Collect all of the edges that the passed events match.  Returns
@@ -149,9 +70,7 @@ class Puppet::SimpleGraph
     # Get all of the edges that this vertex should forward events
     # to, which is the same thing as saying all edges directly below
     # This vertex in the graph.
-    adjacent(source, :direction => :out, :type => :edges).find_all do |edge|
-      edge.match?(event.name)
-    end
+    @out_from[source].values.flatten.find_all { |edge| edge.match?(event.name) }
   end
 
   # Return a reversed version of this graph.
@@ -159,20 +78,50 @@ class Puppet::SimpleGraph
     result = self.class.new
     vertices.each { |vertex| result.add_vertex(vertex) }
     edges.each do |edge|
-      newedge = edge.class.new(edge.target, edge.source, edge.label)
-      result.add_edge(newedge)
+      result.add_edge edge.class.new(edge.target, edge.source, edge.label)
     end
     result
   end
 
   # Return the size of the graph.
   def size
-    @vertices.length
+    vertices.size
   end
 
-  # Return the graph as an array.
   def to_a
-    @vertices.keys
+    vertices
+  end
+
+  # Provide a topological sort with cycle reporting
+  def topsort_with_cycles
+    degree = {}
+    zeros = []
+    result = []
+
+    # Collect each of our vertices, with the number of in-edges each has.
+    vertices.each do |v|
+      edges = @in_to[v].dup
+      zeros << v if edges.empty?
+      degree[v] = edges
+    end
+
+    # Iterate over each 0-degree vertex, decrementing the degree of
+    # each of its out-edges.
+    while v = zeros.pop
+      result << v
+      @out_from[v].each { |v2,es| 
+        degree[v2].delete(v)
+        zeros << v2 if degree[v2].empty?
+      }
+    end
+
+    # If we have any vertices left with non-zero in-degrees, then we've found a cycle.
+    if cycles = degree.values.reject { |ns| ns.empty? } and cycles.length > 0
+      message = cycles.collect { |edges| '('+edges.collect { |e| e.to_s }.join(", ")+')' }.join(", ")
+      raise Puppet::Error, "Found dependency cycles in the following relationships: #{message}; try using the '--graph' option and open the '.dot' files in OmniGraffle or GraphViz"
+    end
+
+    result
   end
 
   # Provide a topological sort.
@@ -182,26 +131,24 @@ class Puppet::SimpleGraph
     result = []
 
     # Collect each of our vertices, with the number of in-edges each has.
-    @vertices.each do |name, wrapper|
-      edges = wrapper.in_edges
-      zeros << wrapper if edges.length == 0
-      degree[wrapper.vertex] = edges
+    vertices.each do |v|
+      edges = @in_to[v]
+      zeros << v if edges.empty?
+      degree[v] = edges.length
     end
 
     # Iterate over each 0-degree vertex, decrementing the degree of
     # each of its out-edges.
-    while wrapper = zeros.pop
-      result << wrapper.vertex
-      wrapper.out_edges.each do |edge|
-        degree[edge.target].delete(edge)
-        zeros << @vertices[edge.target] if degree[edge.target].length == 0
-      end
+    while v = zeros.pop
+      result << v
+      @out_from[v].each { |v2,es| 
+        zeros << v2 if (degree[v2] -= 1) == 0
+      }
     end
 
     # If we have any vertices left with non-zero in-degrees, then we've found a cycle.
-    if cycles = degree.find_all { |vertex, edges| edges.length > 0 } and cycles.length > 0
-      message = cycles.collect { |vertex, edges| edges.collect { |e| e.to_s }.join(", ") }.join(", ")
-      raise Puppet::Error, "Found dependency cycles in the following relationships: #{message}; try using the '--graph' option and open the '.dot' files in OmniGraffle or GraphViz"
+    if cycles = degree.values.reject { |ns| ns == 0  } and cycles.length > 0
+      topsort_with_cycles
     end
 
     result
@@ -209,103 +156,80 @@ class Puppet::SimpleGraph
 
   # Add a new vertex to the graph.
   def add_vertex(vertex)
-    @reversal = nil
-    return false if vertex?(vertex)
-    setup_vertex(vertex)
-    true # don't return the VertexWrapper instance.
+    @in_to[vertex]    ||= {}
+    @out_from[vertex] ||= {}
   end
 
   # Remove a vertex from the graph.
-  def remove_vertex!(vertex)
-    return nil unless vertex?(vertex)
-    @vertices[vertex].edges.each { |edge| remove_edge!(edge) }
-    @edges -= @vertices[vertex].edges
-    @vertices[vertex].clear
-    @vertices.delete(vertex)
+  def remove_vertex!(v)
+    return unless vertex?(v)
+    @upstream_from.clear
+    @downstream_from.clear
+    (@in_to[v].values+@out_from[v].values).flatten.each { |e| remove_edge!(e) }
+    @in_to.delete(v)
+    @out_from.delete(v)
   end
 
   # Test whether a given vertex is in the graph.
-  def vertex?(vertex)
-    @vertices.include?(vertex)
+  def vertex?(v)
+    @in_to.include?(v)
   end
 
   # Return a list of all vertices.
   def vertices
-    @vertices.keys
+    @in_to.keys
   end
 
   # Add a new edge.  The graph user has to create the edge instance,
   # since they have to specify what kind of edge it is.
-  def add_edge(source, target = nil, label = nil)
-    @reversal = nil
-    if target
-      edge = Puppet::Relationship.new(source, target, label)
-    else
-      edge = source
-    end
-    [edge.source, edge.target].each { |vertex| setup_vertex(vertex) unless vertex?(vertex) }
-    @vertices[edge.source].add_edge :out, edge
-    @vertices[edge.target].add_edge :in, edge
-    @edges << edge
-    true
+  def add_edge(e,*a)
+    return add_relationship(e,*a) unless a.empty?
+    @upstream_from.clear
+    @downstream_from.clear
+    add_vertex(e.source)
+    add_vertex(e.target)
+    @in_to[   e.target][e.source] ||= []; @in_to[   e.target][e.source] |= [e]
+    @out_from[e.source][e.target] ||= []; @out_from[e.source][e.target] |= [e]
   end
 
-  # Find a matching edge.  Note that this only finds the first edge,
-  # not all of them or whatever.
-  def edge(source, target)
-    @edges.each_with_index { |test_edge, index| return test_edge if test_edge.source == source and test_edge.target == target }
+  def add_relationship(source, target, label = nil)
+    add_edge Puppet::Relationship.new(source, target, label)
   end
 
-  def edge_label(source, target)
-    return nil unless edge = edge(source, target)
-    edge.label
+  # Find all matching edges.
+  def edges_between(source, target)
+    (@out_from[source] || {})[target] || []
   end
 
   # Is there an edge between the two vertices?
   def edge?(source, target)
-    return false unless vertex?(source) and vertex?(target)
-
-    @vertices[source].has_edge?(:out, target)
+    vertex?(source) and vertex?(target) and @out_from[source][target]
   end
 
   def edges
-    @edges.dup
+    @in_to.values.collect { |x| x.values }.flatten
   end
 
-  # Remove an edge from our graph.
-  def remove_edge!(edge)
-    @vertices[edge.source].remove_edge(:out, edge)
-    @vertices[edge.target].remove_edge(:in, edge)
-
-    @edges.delete(edge)
-    nil
+  def each_edge
+    @in_to.each { |t,ns| ns.each { |s,es| es.each { |e| yield e }}}
   end
 
-  # Find adjacent edges.
-  def adjacent(vertex, options = {})
-    return [] unless wrapper = @vertices[vertex]
-    wrapper.adjacent(options)
+  # Remove an edge from our graph.
+  def remove_edge!(e)
+    if edge?(e.source,e.target)
+      @upstream_from.clear
+      @downstream_from.clear
+      @in_to   [e.target].delete e.source if (@in_to   [e.target][e.source] -= [e]).empty?
+      @out_from[e.source].delete e.target if (@out_from[e.source][e.target] -= [e]).empty?
+    end
   end
 
-  private
-
-  # An internal method that skips the validation, so we don't have
-  # duplicate validation calls.
-  def setup_vertex(vertex)
-    @vertices[vertex] = VertexWrapper.new(vertex)
+  # Find adjacent edges.
+  def adjacent(v, options = {})
+    return [] unless ns = (options[:direction] == :in) ? @in_to[v] : @out_from[v]
+    (options[:type] == :edges) ? ns.values.flatten : ns.keys
   end
-
-  public
-
-#    # For some reason, unconnected vertices do not show up in
-#    # this graph.
-#    def to_jpg(path, name)
-#        gv = vertices
-#        Dir.chdir(path) do
-#            induced_subgraph(gv).write_to_graphic_file('jpg', name)
-#        end
-#    end
-
+  
   # Take container information from another graph and use it
   # to replace any container vertices with their respective leaves.
   # This creates direct relationships where there were previously
@@ -381,6 +305,26 @@ class Puppet::SimpleGraph
     predecessor
   end
 
+  def downstream_from_vertex(v)
+    return @downstream_from[v] if @downstream_from[v]
+    result = @downstream_from[v] = {}
+    @out_from[v].keys.each do |node|
+      result[node] = 1
+      result.update(downstream_from_vertex(node))
+    end
+    result
+  end
+
+  def upstream_from_vertex(v)
+    return @upstream_from[v] if @upstream_from[v]
+    result = @upstream_from[v] = {}
+    @in_to[v].keys.each do |node|
+      result[node] = 1
+      result.update(upstream_from_vertex(node))
+    end
+    result
+  end
+
   # LAK:FIXME This is just a paste of the GRATR code with slight modifications.
 
   # Return a DOT::DOTDigraph for directed graphs or a DOT::DOTSubgraph for an
@@ -422,18 +366,6 @@ class Puppet::SimpleGraph
     system('dotty', dotfile)
   end
 
-  # Use +dot+ to create a graphical representation of the graph.  Returns the
-  # filename of the graphics file.
-  def write_to_graphic_file (fmt='png', dotfile='graph')
-    src = dotfile + '.dot'
-    dot = dotfile + '.' + fmt
-
-    File.open(src, 'w') {|f| f << self.to_dot << "\n"}
-
-    system( "dot -T#{fmt} #{src} -o #{dot}" )
-    dot
-  end
-
   # Produce the graph files if requested.
   def write_graph(name)
     return unless Puppet[:graph]
@@ -445,4 +377,73 @@ class Puppet::SimpleGraph
       f.puts to_dot("name" => name.to_s.capitalize)
     }
   end
+
+  # This flag may be set to true to use the new YAML serialzation
+  # format (where @vertices is a simple list of vertices rather than a
+  # list of VertexWrapper objects).  Deserialization supports both
+  # formats regardless of the setting of this flag.
+  class << self
+    attr_accessor :use_new_yaml_format
+  end
+  self.use_new_yaml_format = false
+
+  # Stub class to allow graphs to be represented in YAML using the old
+  # (version 2.6) format.
+  class VertexWrapper
+    attr_reader :vertex, :adjacencies
+    def initialize(vertex, adjacencies)
+      @vertex = vertex
+      @adjacencies = adjacencies
+    end
+  end
+
+  # instance_variable_get is used by Object.to_zaml to get instance
+  # variables.  Override it so that we can simulate the presence of
+  # instance variables @edges and @vertices for serialization.
+  def instance_variable_get(v)
+    case v.to_s
+    when '@edges' then
+      edges
+    when '@vertices' then
+      if self.class.use_new_yaml_format
+        vertices
+      else
+        result = {}
+        vertices.each do |vertex|
+          adjacencies = {}
+          [:in, :out].each do |direction|
+            adjacencies[direction] = {}
+            adjacent(vertex, :direction => direction, :type => :edges).each do |edge|
+              other_vertex = direction == :in ? edge.source : edge.target
+              (adjacencies[direction][other_vertex] ||= Set.new).add(edge)
+            end
+          end
+          result[vertex] = Puppet::SimpleGraph::VertexWrapper.new(vertex, adjacencies)
+        end
+        result
+      end
+    else
+      super(v)
+    end
+  end
+
+  def to_yaml_properties
+    other_vars = instance_variables.reject { |v| %w{@in_to @out_from @upstream_from @downstream_from}.include?(v) }
+    (other_vars + %w{@vertices @edges}).sort.uniq
+  end
+
+  def yaml_initialize(tag, var)
+    initialize()
+    vertices = var.delete('vertices')
+    edges = var.delete('edges')
+    if vertices.is_a?(Hash)
+      # Support old (2.6) format
+      vertices = vertices.keys
+    end
+    vertices.each { |v| add_vertex(v) }
+    edges.each { |e| add_edge(e) }
+    var.each do |varname, value|
+      instance_variable_set("@#{varname}", value)
+    end
+  end
 end
diff --git a/lib/puppet/util/monkey_patches.rb b/lib/puppet/util/monkey_patches.rb
index bdce5ec1d..1c35ae523 100644
--- a/lib/puppet/util/monkey_patches.rb
+++ b/lib/puppet/util/monkey_patches.rb
@@ -48,6 +48,7 @@ if RUBY_VERSION == '1.8.7'
   end
 end
 
+
 class Object
   # ActiveSupport 2.3.x mixes in a dangerous method
   # that can cause rspec to fork bomb
@@ -56,3 +57,17 @@ class Object
     raise NotImplementedError, "Kernel.daemonize is too dangerous, please don't try to use it."
   end
 end
+
+# Workaround for yaml_initialize, which isn't supported before Ruby
+# 1.8.3.
+if RUBY_VERSION == '1.8.1' || RUBY_VERSION == '1.8.2'
+  YAML.add_ruby_type( /^object/ ) { |tag, val|
+    type, obj_class = YAML.read_type_class( tag, Object )
+    r = YAML.object_maker( obj_class, val )
+    if r.respond_to? :yaml_initialize
+      r.instance_eval { instance_variables.each { |name| remove_instance_variable name } }
+      r.yaml_initialize(tag, val)
+    end
+    r
+  }
+end
diff --git a/spec/unit/resource/catalog_spec.rb b/spec/unit/resource/catalog_spec.rb
index 2b6beb5e9..fbfe29ff3 100755
--- a/spec/unit/resource/catalog_spec.rb
+++ b/spec/unit/resource/catalog_spec.rb
@@ -374,7 +374,7 @@ describe Puppet::Resource::Catalog, "when compiling" do
       @original.add_edge(@r1,@r2)
       @original.filter do |r|
         r == @r1
-      end.edge(@r1,@r2).should be_empty
+      end.edge?(@r1,@r2).should_not be
     end
   end
 
@@ -933,8 +933,8 @@ describe Puppet::Resource::Catalog, "when converting to pson" do
     @catalog.add_edge(one, two)
     @catalog.add_edge(two, three)
 
-    @catalog.edge(one, two  ).expects(:to_pson_data_hash).returns "one_two_pson"
-    @catalog.edge(two, three).expects(:to_pson_data_hash).returns "two_three_pson"
+    @catalog.edges_between(one, two  )[0].expects(:to_pson_data_hash).returns "one_two_pson"
+    @catalog.edges_between(two, three)[0].expects(:to_pson_data_hash).returns "two_three_pson"
 
     PSON.parse(@catalog.to_pson,:create_additions => false)['data']['edges'].sort.should == %w{one_two_pson two_three_pson}.sort
   end
diff --git a/spec/unit/simple_graph_spec.rb b/spec/unit/simple_graph_spec.rb
index 2ca8888c5..fa0bcb06a 100755
--- a/spec/unit/simple_graph_spec.rb
+++ b/spec/unit/simple_graph_spec.rb
@@ -31,12 +31,6 @@ describe Puppet::SimpleGraph do
     proc { @graph.to_dot_graph }.should_not raise_error
   end
 
-  it "should always put its edges first when printing yaml" do
-    @graph = Puppet::SimpleGraph.new
-    @graph.add_edge(:one, :two)
-    @graph.to_yaml_properties[0].should == "@edges"
-  end
-
   describe "when managing vertices" do
     before do
       @graph = Puppet::SimpleGraph.new
@@ -117,16 +111,31 @@ describe Puppet::SimpleGraph do
       @graph.edge?(:one, :two).should be_true
     end
 
-    it "should provide a method for retrieving an edge label" do
-      edge = Puppet::Relationship.new(:one, :two, :callback => :awesome)
-      @graph.add_edge(edge)
-      @graph.edge_label(:one, :two).should == {:callback => :awesome}
-    end
+    describe "when retrieving edges between two nodes" do
+      it "should handle the case of nodes not in the graph" do
+        @graph.edges_between(:one, :two).should == []
+      end
 
-    it "should provide a method for retrieving an edge" do
-      edge = Puppet::Relationship.new(:one, :two)
-      @graph.add_edge(edge)
-      @graph.edge(:one, :two).should equal(edge)
+      it "should handle the case of nodes with no edges between them" do
+        @graph.add_vertex(:one)
+        @graph.add_vertex(:two)
+        @graph.edges_between(:one, :two).should == []
+      end
+
+      it "should handle the case of nodes connected by a single edge" do
+        edge = Puppet::Relationship.new(:one, :two)
+        @graph.add_edge(edge)
+        @graph.edges_between(:one, :two).length.should == 1
+        @graph.edges_between(:one, :two)[0].should equal(edge)
+      end
+
+      it "should handle the case of nodes connected by multiple edges" do
+        edge1 = Puppet::Relationship.new(:one, :two, :callback => :foo)
+        edge2 = Puppet::Relationship.new(:one, :two, :callback => :bar)
+        @graph.add_edge(edge1)
+        @graph.add_edge(edge2)
+        Set.new(@graph.edges_between(:one, :two)).should == Set.new([edge1, edge2])
+      end
     end
 
     it "should add the edge source as a vertex if it is not already" do
@@ -253,7 +262,7 @@ describe Puppet::SimpleGraph do
 
     it "should retain labels on edges" do
       @graph.add_edge(:one, :two, :callback => :awesome)
-      edge = @graph.reversal.edge(:two, :one)
+      edge = @graph.reversal.edges_between(:two, :one)[0]
       edge.label.should == {:callback => :awesome}
     end
   end
@@ -522,14 +531,14 @@ describe Puppet::SimpleGraph do
     it "should not add labels to edges that have none" do
       @depgraph.add_edge(@two, @three)
       splice
-      @depgraph.edge_label("c", "i").should == {}
+      @depgraph.edges_between("c", "i")[0].label.should == {}
     end
 
     it "should copy labels over edges that have none" do
       @depgraph.add_edge("c", @three, {:callback => :refresh})
       splice
       # And make sure the label got copied.
-      @depgraph.edge_label("c", "i").should == {:callback => :refresh}
+      @depgraph.edges_between("c", "i")[0].label.should == {:callback => :refresh}
     end
 
     it "should not replace a label with a nil label" do
@@ -537,7 +546,7 @@ describe Puppet::SimpleGraph do
       @depgraph.add_edge(@middle, @three)
       @depgraph.add_edge("c", @three, {:callback => :refresh})
       splice
-      @depgraph.edge_label("c", "i").should == {:callback => :refresh}
+      @depgraph.edges_between("c", "i")[0].label.should == {:callback => :refresh}
     end
 
     it "should copy labels to all created edges" do
@@ -547,8 +556,207 @@ describe Puppet::SimpleGraph do
       @three.each do |child|
         edge = Puppet::Relationship.new("c", child)
         @depgraph.should be_edge(edge.source, edge.target)
-        @depgraph.edge_label(edge.source, edge.target).should == {:callback => :refresh}
+        @depgraph.edges_between(edge.source, edge.target)[0].label.should == {:callback => :refresh}
+      end
+    end
+  end
+
+  it "should serialize to YAML using the old format by default" do
+    Puppet::SimpleGraph.use_new_yaml_format.should == false
+  end
+
+  describe "(yaml tests)" do
+    def empty_graph(graph)
+    end
+
+    def one_vertex_graph(graph)
+      graph.add_vertex(:a)
+    end
+
+    def graph_without_edges(graph)
+      [:a, :b, :c].each { |x| graph.add_vertex(x) }
+    end
+
+    def one_edge_graph(graph)
+      graph.add_edge(:a, :b)
+    end
+
+    def many_edge_graph(graph)
+      graph.add_edge(:a, :b)
+      graph.add_edge(:a, :c)
+      graph.add_edge(:b, :d)
+      graph.add_edge(:c, :d)
+    end
+
+    def labeled_edge_graph(graph)
+      graph.add_edge(:a, :b, :callback => :foo, :event => :bar)
+    end
+
+    def overlapping_edge_graph(graph)
+      graph.add_edge(:a, :b, :callback => :foo, :event => :bar)
+      graph.add_edge(:a, :b, :callback => :biz, :event => :baz)
+    end
+
+    def self.all_test_graphs
+      [:empty_graph, :one_vertex_graph, :graph_without_edges, :one_edge_graph, :many_edge_graph, :labeled_edge_graph,
+       :overlapping_edge_graph]
+    end
+
+    def object_ids(enumerable)
+      # Return a sorted list of the object id's of the elements of an
+      # enumerable.
+      enumerable.collect { |x| x.object_id }.sort
+    end
+
+    def graph_to_yaml(graph, which_format)
+      previous_use_new_yaml_format = Puppet::SimpleGraph.use_new_yaml_format
+      Puppet::SimpleGraph.use_new_yaml_format = (which_format == :new)
+      ZAML.dump(graph)
+    ensure
+      Puppet::SimpleGraph.use_new_yaml_format = previous_use_new_yaml_format
+    end
+
+    # Test serialization of graph to YAML.
+    [:old, :new].each do |which_format|
+      all_test_graphs.each do |graph_to_test|
+        it "should be able to serialize #{graph_to_test} to YAML (#{which_format} format)" do
+          graph = Puppet::SimpleGraph.new
+          send(graph_to_test, graph)
+          yaml_form = graph_to_yaml(graph, which_format)
+
+          # Hack the YAML so that objects in the Puppet namespace get
+          # changed to YAML::DomainType objects.  This lets us inspect
+          # the serialized objects easily without invoking any
+          # yaml_initialize hooks.
+          yaml_form.gsub!('!ruby/object:Puppet::', '!hack/object:Puppet::')
+          serialized_object = YAML.load(yaml_form)
+
+          # Check that the object contains instance variables @edges and
+          # @vertices only.  @reversal is also permitted, but we don't
+          # check it, because it is going to be phased out.
+          serialized_object.type_id.should == 'object:Puppet::SimpleGraph'
+          serialized_object.value.keys.reject { |x| x == 'reversal' }.sort.should == ['edges', 'vertices']
+
+          # Check edges by forming a set of tuples (source, target,
+          # callback, event) based on the graph and the YAML and make sure
+          # they match.
+          edges = serialized_object.value['edges']
+          edges.should be_a(Array)
+          expected_edge_tuples = graph.edges.collect { |edge| [edge.source, edge.target, edge.callback, edge.event] }
+          actual_edge_tuples = edges.collect do |edge|
+            edge.type_id.should == 'object:Puppet::Relationship'
+            %w{source target}.each { |x| edge.value.keys.should include(x) }
+            edge.value.keys.each { |x| ['source', 'target', 'callback', 'event'].should include(x) }
+            %w{source target callback event}.collect { |x| edge.value[x] }
+          end
+          Set.new(actual_edge_tuples).should == Set.new(expected_edge_tuples)
+          actual_edge_tuples.length.should == expected_edge_tuples.length
+
+          # Check vertices one by one.
+          vertices = serialized_object.value['vertices']
+          if which_format == :old
+            vertices.should be_a(Hash)
+            Set.new(vertices.keys).should == Set.new(graph.vertices)
+            vertices.each do |key, value|
+              value.type_id.should == 'object:Puppet::SimpleGraph::VertexWrapper'
+              value.value.keys.sort.should == %w{adjacencies vertex}
+              value.value['vertex'].should equal(key)
+              adjacencies = value.value['adjacencies']
+              adjacencies.should be_a(Hash)
+              Set.new(adjacencies.keys).should == Set.new([:in, :out])
+              [:in, :out].each do |direction|
+                adjacencies[direction].should be_a(Hash)
+                expected_adjacent_vertices = Set.new(graph.adjacent(key, :direction => direction, :type => :vertices))
+                Set.new(adjacencies[direction].keys).should == expected_adjacent_vertices
+                adjacencies[direction].each do |adj_key, adj_value|
+                  # Since we already checked edges, just check consistency
+                  # with edges.
+                  desired_source = direction == :in ? adj_key : key
+                  desired_target = direction == :in ? key : adj_key
+                  expected_edges = edges.select do |edge|
+                    edge.value['source'] == desired_source && edge.value['target'] == desired_target
+                  end
+                  adj_value.should be_a(Set)
+                  if object_ids(adj_value) != object_ids(expected_edges)
+                    raise "For vertex #{key.inspect}, direction #{direction.inspect}: expected adjacencies #{expected_edges.inspect} but got #{adj_value.inspect}"
+                  end
+                end
+              end
+            end
+          else
+            vertices.should be_a(Array)
+            Set.new(vertices).should == Set.new(graph.vertices)
+            vertices.length.should == graph.vertices.length
+          end
+        end
+      end
+
+      # Test deserialization of graph from YAML.  This presumes the
+      # correctness of serialization to YAML, which has already been
+      # tested.
+      all_test_graphs.each do |graph_to_test|
+        it "should be able to deserialize #{graph_to_test} from YAML (#{which_format} format)" do
+          reference_graph = Puppet::SimpleGraph.new
+          send(graph_to_test, reference_graph)
+          yaml_form = graph_to_yaml(reference_graph, which_format)
+          recovered_graph = YAML.load(yaml_form)
+
+          # Test that the recovered vertices match the vertices in the
+          # reference graph.
+          expected_vertices = reference_graph.vertices.to_a
+          recovered_vertices = recovered_graph.vertices.to_a
+          Set.new(recovered_vertices).should == Set.new(expected_vertices)
+          recovered_vertices.length.should == expected_vertices.length
+
+          # Test that the recovered edges match the edges in the
+          # reference graph.
+          expected_edge_tuples = reference_graph.edges.collect do |edge|
+            [edge.source, edge.target, edge.callback, edge.event]
+          end
+          recovered_edge_tuples = recovered_graph.edges.collect do |edge|
+            [edge.source, edge.target, edge.callback, edge.event]
+          end
+          Set.new(recovered_edge_tuples).should == Set.new(expected_edge_tuples)
+          recovered_edge_tuples.length.should == expected_edge_tuples.length
+
+          # We ought to test that the recovered graph is self-consistent
+          # too.  But we're not going to bother with that yet because
+          # the internal representation of the graph is about to change.
+        end
+      end
+
+      it "should be able to serialize a graph where the vertices contain backreferences to the graph (#{which_format} format)" do
+        reference_graph = Puppet::SimpleGraph.new
+        vertex = Object.new
+        vertex.instance_eval { @graph = reference_graph }
+        reference_graph.add_edge(vertex, :other_vertex)
+        yaml_form = graph_to_yaml(reference_graph, which_format)
+        recovered_graph = YAML.load(yaml_form)
+
+        recovered_graph.vertices.length.should == 2
+        recovered_vertex = recovered_graph.vertices.reject { |x| x.is_a?(Symbol) }[0]
+        recovered_vertex.instance_eval { @graph }.should equal(recovered_graph)
+        recovered_graph.edges.length.should == 1
+        recovered_edge = recovered_graph.edges[0]
+        recovered_edge.source.should equal(recovered_vertex)
+        recovered_edge.target.should == :other_vertex
+      end
+    end
+
+    it "should serialize properly when used as a base class" do
+      class Puppet::TestDerivedClass < Puppet::SimpleGraph
+        attr_accessor :foo
       end
+      derived = Puppet::TestDerivedClass.new
+      derived.add_edge(:a, :b)
+      derived.foo = 1234
+      recovered_derived = YAML.load(YAML.dump(derived))
+      recovered_derived.class.should equal(Puppet::TestDerivedClass)
+      recovered_derived.edges.length.should == 1
+      recovered_derived.edges[0].source.should == :a
+      recovered_derived.edges[0].target.should == :b
+      recovered_derived.vertices.length.should == 2
+      recovered_derived.foo.should == 1234
     end
   end
 end
diff --git a/spec/unit/util/monkey_patches_spec.rb b/spec/unit/util/monkey_patches_spec.rb
index b0f61c808..049ed1044 100644
--- a/spec/unit/util/monkey_patches_spec.rb
+++ b/spec/unit/util/monkey_patches_spec.rb
@@ -5,3 +5,29 @@ Dir.chdir(File.dirname(__FILE__)) { (s = lambda { |f| File.exist?(f) ? require(f
 require 'puppet/util/monkey_patches'
 
 
+
+describe "yaml deserialization" do
+  it "should call yaml_initialize when deserializing objects that have that method defined" do
+    class Puppet::TestYamlInitializeClass
+      attr_reader :foo
+
+      def yaml_initialize(tag, var)
+        var.should == {'foo' => 100}
+        instance_variables.should == []
+        @foo = 200
+      end
+    end
+
+    obj = YAML.load("--- !ruby/object:Puppet::TestYamlInitializeClass\n  foo: 100")
+    obj.foo.should == 200
+  end
+
+  it "should not call yaml_initialize if not defined" do
+    class Puppet::TestYamlNonInitializeClass
+      attr_reader :foo
+    end
+
+    obj = YAML.load("--- !ruby/object:Puppet::TestYamlNonInitializeClass\n  foo: 100")
+    obj.foo.should == 100
+  end
+end
diff --git a/spec/unit/util/zaml_spec.rb b/spec/unit/util/zaml_spec.rb
index b223f89d4..358c6aa11 100755
--- a/spec/unit/util/zaml_spec.rb
+++ b/spec/unit/util/zaml_spec.rb
@@ -61,4 +61,3 @@ describe "Pure ruby yaml implementation" do
     x2[2].should equal(x2)
   end
 end
-