require 'test/unit'

require 'rockit/tree/enter_leave_visitor'

class UTestTreeEnterLeaveVisitor < 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])
    @d = Rockit::Tree::Base.new_tree_class(:D)
    @e = Rockit::Tree::Base.new_tree_class(:E)
  end

  # Visits a node both on the way down in the tree and on the way up.
  # Still a depth-first search.
  class TreeELVisitor
    include EnterLeaveVisitor

    def initialize
      @visited = []
    end

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

    def enter_A(obj); @visited << obj; end
    def leave_A(obj); @visited << obj; end
    def enter_B(obj); @visited << obj; end
    def leave_B(obj); @visited << obj; end
    def enter_C(obj); @visited << obj; end
    def leave_C(obj); @visited << obj; end
    def visit_E(obj); @visited << obj; end
  end

  class TreeAndFixnumELVisitor < TreeELVisitor
    def enter_Fixnum(obj); @visited << obj; end
  end

  def test_01_visitor
    tree = @c[tb1 = @b[ta = @a[@d[te = @e[]]]], tb2 = @b[t1 = 1]]
    visitor = TreeELVisitor.new
    tree.accept_visitor(visitor)

    assert_equal(tree, visitor[0])
    assert_equal(tb1, visitor[1])
    assert_equal(ta, visitor[2])
    # D is never visited sine we have no enter, visit or leave method for it.
    # E is visited even though we only have a visit method but no enter or 
    # leave methods.
    assert_equal(te, visitor[3])
    assert_equal(ta, visitor[4])
    assert_equal(tb1, visitor[5])
    assert_equal(tb2, visitor[6])
    assert_equal(tb2, visitor[7])
    assert_equal(tree, visitor[8])

    visitor2 = TreeAndFixnumELVisitor.new
    tree.accept_visitor(visitor2)
    assert_equal(tree, visitor2[0])
    assert_equal(tb1, visitor2[1])
    assert_equal(ta, visitor2[2])
    assert_equal(te, visitor[3])
    assert_equal(ta, visitor2[4])
    assert_equal(tb1, visitor2[5])
    assert_equal(tb2, visitor2[6])
    assert_equal(1, visitor2[7])
    assert_equal(tb2, visitor2[8])
    assert_equal(tree, visitor2[9])
  end
end