require 'test/unit'

require 'rockit/tree/visitor'

class UTestTreeVisitor < Test::Unit::TestCase
  def setup
    @a = Rockit::Tree::Base.new_tree_class(:A)
    @b = Rockit::Tree::Base.new_tree_class(:B, [:c1])
    @c = Rockit::Tree::Base.new_tree_class(:C, [:c1, :c2])
  end

  class TreeVisitor
    include Visitor

    def initialize
      @visited = []
    end

    def [](index); @visited[index]; end

    def visit_A(obj)
      @visited << obj
    end

    def visit_B(obj)
      @visited << obj
    end

    def visit_C(obj)
      @visited << obj
    end
  end

  class TreeAndFixnumVisitor < TreeVisitor
    def visit_Fixnum(obj)
      @visited << obj
    end
  end

  def test_01_visitor
    tree = @c[tb1 = @b[ta = @a[]], tb2 = @b[t1 = 1]]
    visitor = TreeVisitor.new
    tree.accept_visitor(visitor)
    # The visitor pattern by default does a depth-first visit of the nodes
    # in the tree. Note that the 1 is not in the visited list since the visitor
    # has no method to visit Fixnums.
    assert_equal(tree, visitor[0])
    assert_equal(tb1, visitor[1])
    assert_equal(ta, visitor[2])
    assert_equal(tb2, visitor[3])

    # This visitor has a method to visit Fixnums so it should have the same
    # obj's in the visited array as the previous one but then followed by
    # the 1 that is the child of the tb2 tree above.
    visitor2 = TreeAndFixnumVisitor.new
    tree.accept_visitor(visitor2)
    assert_equal(tree, visitor2[0])
    assert_equal(tb1, visitor2[1])
    assert_equal(ta, visitor2[2])
    assert_equal(tb2, visitor2[3])
    assert_equal(t1, visitor2[4])
  end
end