5 files changed, 57 insertions, 15 deletions

diff --git a/lib/psych/visitors/to_ruby.rb b/lib/psych/visitors/to_ruby.rb
index aed56ed..cadb8cf 100644
--- a/lib/psych/visitors/to_ruby.rb
+++ b/lib/psych/visitors/to_ruby.rb
@@ -67,8 +67,6 @@ module Psych
       end
 
       def visit_Psych_Nodes_Sequence o
-        list = []
-        @st[o.anchor] = list if o.anchor
         case o.tag
         when '!omap', 'tag:yaml.org,2002:omap'
           map = Psych::Omap.new
@@ -78,6 +76,8 @@ module Psych
           }
           map
         else
+          list = []
+          @st[o.anchor] = list if o.anchor
           o.children.each { |c| list.push accept c }
           list
         end
@@ -95,10 +95,13 @@ module Psych
           string
         when /!ruby\/struct:?(.*)?$/
           klass   = resolve_class($1)
-          members = o.children.map { |c| accept c }
 
           if klass
             s = klass.allocate
+            @st[o.anchor] = s if o.anchor
+
+            members = o.children.map { |c| accept c }
+
             struct_members = s.members.map { |x| x.to_sym }
             members.each_slice(2) { |k,v|
               if struct_members.include? k.to_sym
@@ -109,6 +112,7 @@ module Psych
             }
             s
           else
+            members = o.children.map { |c| accept c }
             h = Hash[*members]
             Struct.new(*h.map { |k,v| k.to_sym }).new(*h.map { |k,v| v })
           end
diff --git a/lib/psych/visitors/yaml_tree.rb b/lib/psych/visitors/yaml_tree.rb
index 995a890..e4e6947 100644
--- a/lib/psych/visitors/yaml_tree.rb
+++ b/lib/psych/visitors/yaml_tree.rb
@@ -15,17 +15,23 @@ module Psych
         target.class.ancestors.each do |klass|
           next unless klass.name
           method_name = :"visit_#{klass.name.split('::').join('_')}"
-          return send(method_name, target) if respond_to?(method_name)
+
+          if respond_to?(method_name)
+
+            # return any aliases we find
+            if node = @st[target.object_id]
+              node.anchor = target.object_id.to_s
+              return append Nodes::Alias.new target.object_id.to_s
+            end
+
+            return send(method_name, target)
+          end
+
         end
         raise TypeError, "Can't dump #{target.class}"
       end
 
       def visit_Psych_Set o
-        if node = @st[o.object_id]
-          node.anchor = o.object_id.to_s
-          return append Nodes::Alias.new o.object_id.to_s
-        end
-
         map = Nodes::Mapping.new(nil, '!set', false)
         @st[o.object_id] = map
 
@@ -40,11 +46,6 @@ module Psych
       end
 
       def visit_Psych_Omap o
-        if node = @st[o.object_id]
-          node.anchor = o.object_id.to_s
-          return append Nodes::Alias.new o.object_id.to_s
-        end
-
         seq = Nodes::Sequence.new(nil, '!omap', false)
         @st[o.object_id] = seq
 
@@ -75,7 +76,10 @@ module Psych
       def visit_Struct o
         tag = ['!ruby/struct', o.class.name].compact.join(':')
 
-        @stack.push append Nodes::Mapping.new(nil, tag, false)
+        map = Nodes::Mapping.new(nil, tag, false)
+        @st[o.object_id] = map
+
+        @stack.push append map
         o.members.each do |member|
           accept member
           accept o[member]
diff --git a/test/psych/test_struct.rb b/test/psych/test_struct.rb
new file mode 100644
index 0000000..1fb8149
--- /dev/null
+++ b/test/psych/test_struct.rb
@@ -0,0 +1,24 @@
+require 'helper'
+
+module Psych
+  class TestStruct < Test::Unit::TestCase
+    class StructSubclass < Struct.new(:foo)
+      def initialize foo, bar
+        super(foo)
+        @bar = bar
+      end
+    end
+
+    def test_self_referential_struct
+      ss = StructSubclass.new(nil, 'foo')
+      ss.foo = ss
+
+      loaded = Psych.load(Psych.dump(ss))
+      assert_instance_of(StructSubclass, loaded.foo)
+
+      # FIXME: This seems to cause an infinite loop.  wtf.  Must report a bug
+      # in ruby.
+      # assert_equal(ss, loaded)
+    end
+  end
+end
diff --git a/test/yaml/test_array.rb b/test/yaml/test_array.rb
index 208c996..7e7318e 100644
--- a/test/yaml/test_array.rb
+++ b/test/yaml/test_array.rb
@@ -6,6 +6,11 @@ module YAML
       @list = [{ :a => 'b' }, 'foo']
     end
 
+    def test_self_referential
+      @list << @list
+      assert_equal @list, YAML.load(@list.to_yaml)
+    end
+
     def test_to_yaml
       assert_equal @list, YAML.load(@list.to_yaml)
     end
diff --git a/test/yaml/test_hash.rb b/test/yaml/test_hash.rb
index e8869ae..16a9ee3 100644
--- a/test/yaml/test_hash.rb
+++ b/test/yaml/test_hash.rb
@@ -6,6 +6,11 @@ module YAML
       @hash = { :a => 'b' }
     end
 
+    def test_self_referential
+      @hash['self'] = @hash
+      assert_equal @hash, YAML.load(YAML.dump(@hash))
+    end
+
     def test_to_yaml
       assert_equal @hash, YAML.load(@hash.to_yaml)
     end