require 'test/unit'

require 'rockit/tree/base'
require 'samples/rockit/ruby/sample_ruby_ast_trees/mr_1'

class UTestTreeBase < 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

  def test_01_subclass_with_no_named_children
    t = @a[1]
    assert_kind_of(@a, t)
    assert_equal(1, t.children.length)
    assert_equal(1, t[0])

    t2 = @a["1", 2]
    assert_kind_of(@a, t2)
    assert_equal(2, t2.children.length)
    assert_equal("1", t2[0])
    assert_equal(2, t2[1])
  end

  def test_02_subclass_with_one_named_child
    t = @b[:a]
    assert_kind_of(@b, t)
    assert_equal(1, t.children.length)
    assert_equal(:a, t[0])
    assert_equal(:a, t.c1)
  end

  def test_03_subclass_with_two_named_children
    t = @c[3.0, 4]
    assert_kind_of(@c, t)
    assert_equal(2, t.children.length)
    assert_equal(3.0, t[0])
    assert_equal(3.0, t.c1)
    assert_equal(4, t[1])
    assert_equal(4, t.c2)

    t.c1 = 5
    assert_equal(5, t[0])
    assert_equal(5, t.c1)

    t.c2 = 6
    assert_equal(6, t[1])
    assert_equal(6, t.c2)
  end

  def test_04_complex_tree
    t = @a[@b[@c[1, 2]], @c[3, 4]]
    assert_kind_of(@a, t)
    assert_kind_of(@b, t[0])
    assert_kind_of(@c, t[1])
    assert_kind_of(@c, t[0][0])
    assert_equal(1, t[0][0][0])
    assert_equal(2, t[0][0][1])
    assert_equal(3, t[1][0])
    assert_equal(4, t[1][1])
  end

  def test_05_subclass_of_tree
    @d = @a.new_tree_class(:D, [:d1, :d2])

    t = @d["sub", /a/]
    assert_kind_of(@d, t)
    assert_kind_of(@a, t)
    assert_equal("sub", t[0])
    assert_equal("sub", t.d1)
    assert_equal(/a/, t[1])
    assert_equal(/a/, t.d2)
  end

  def test_06_upcase_symbols_to_patternvar
    t = @a[:X]
    nt = t.upcase_symbols_to_pattern_vars
    assert_kind_of(@a, nt)
    pvar = nt[0]
    assert_kind_of(Rockit::Tree::Base::PatternVar, pvar)
    assert_equal(:X, pvar.name)

    t2 = @c[@b[:X], @a[:Y, :Z]]
    nt2 = t2.upcase_symbols_to_pattern_vars
    assert_kind_of(@c, nt2)
    xvar = nt2[0][0]
    assert_kind_of(Rockit::Tree::Base::PatternVar, xvar)
    assert_equal(:X, pvar.name)
    yvar = nt2[1][0]
    assert_kind_of(Rockit::Tree::Base::PatternVar, yvar)
    assert_equal(:Y, yvar.name)
    zvar = nt2[1][1]
    assert_kind_of(Rockit::Tree::Base::PatternVar, zvar)
    assert_equal(:Z, zvar.name)
  end

  def assert_match_result(expected, tree, pattern)
    actual = pattern.upcase_symbols_to_pattern_vars.match_tree(tree)
    assert_equal(expected, actual)
  end

  def test_07_match_tree_simple
    pattern = @a[:X].upcase_symbols_to_pattern_vars
    tree = @a[1]
    res = pattern.match_tree(tree)
    assert_kind_of(Hash, res)
    assert_equal(1, res[:X])

    t2 = @a[t3 = @b[1]]
    assert_match_result({:X => t3}, t2, pattern)
  end

  def test_08_match_tree_complex
    pattern = @c[@b[:X], @a[:Y]].upcase_symbols_to_pattern_vars
    tree = @c[@b["23"], @a[/a/]]
    res = pattern.match_tree(tree)
    assert_kind_of(Hash, res)
    assert_equal("23", res[:X])
    assert_equal(/a/, res[:Y])

    pattern = @b[@c[:L, :R]].upcase_symbols_to_pattern_vars
    t2 = @b[t3 = @c[3, 4]]
    res = pattern.match_tree(t2)
    assert_equal({:L => 3, :R => 4}, res)
  end

  def test_09_match_tree_complex_with_subtrees
    pattern = @c[@b[:ARNE], @a[:STree]].upcase_symbols_to_pattern_vars
    stree = @c[3, 4]
    tree = @c[@b[1.4], @a[stree]]
    res = pattern.match_tree(tree)
    assert_kind_of(Hash, res)
    assert_equal(1.4, res[:ARNE])
    assert_equal(stree, res[:STree])
  end

  def test_10_regexp_style_matching
    stree = @a[]
    res = @b[@c[stree, 2]] =~ @b[@c[:First, :Second]]
    assert_kind_of(Hash, res)
    assert_equal(2, res[:Second])
    assert_equal(stree, res[:First])
  end

  def test_11_failing_match_tree
    assert_match_result(nil, @b[1], @a[])
    assert_match_result(nil, @a[1], @b[])
    assert_match_result(nil, @a[1], @a[2])
    assert_match_result(nil, @c[@a[], 1], @c[@a[], 2])
    assert_match_result(nil, @c[@a[], "1"], @c[@a[], "2"])
  end

  def assert_pp(expected, tree, width = 79)
    actual = ""
    PP.pp(tree, actual, width)
    assert_equal(expected, actual)
  end

  def test_12_pretty_print
    assert_pp("A[1]\n", @a[1])
    assert_pp("A[B[2]]\n", @a[@b[2]])
    assert_pp("C[B[3.0], /a/]\n", @c[@b[3.0], /a/])
    lstr = "." * 75
    assert_pp("C[\n B[3.0],\n #{lstr.inspect}]\n", @c[@b[3.0], lstr])
  end

  def collect_trees(tree)
    ary = []
    tree.each_tree {|t| ary << t}
    ary
  end

  def test_13_each_tree_simple
    t = @a[]
    assert_equal([t], collect_trees(t))
  end

  def test_14_each_tree_complex
    t = @c[t2 = @b[t3 = @a[1]], t4 = @b[2]]
    assert_equal([t, t2, t3, 1, t4, 2], collect_trees(t))
  end

  def collect_trees_of_class(tree, klass, recursive = false)
    ary = []
    tree.each_tree_of_class(klass, recursive) {|t| ary << t}
    ary
  end

  def test_15_each_tree_of_klass
    t = @c[t2 = @b[t3 = @a[t4 = @b[1]]], t5 = @b[t6 = @c[3, 4]]]
    assert_equal([t2, t5], collect_trees_of_class(t, @b))    
    assert_equal([t2, t4, t5], collect_trees_of_class(t, @b, true))    
    assert_equal([t3], collect_trees_of_class(t, @a))    
    assert_equal([t], collect_trees_of_class(t, @c))
    assert_equal([t, t6], collect_trees_of_class(t, @c, true))
  end

  def test_16_set_parents_and_parent
    t = @c[t2 = @b[t3 = @a[t4 = @b[1]]], t5 = @b[t6 = @c[3, 4]]]

    assert_equal(nil, t.parent)
    assert_equal(nil, t2.parent)
    assert_equal(nil, t3.parent)
    assert_equal(nil, t4.parent)
    assert_equal(nil, t5.parent)
    assert_equal(nil, t6.parent)

    t.set_parents

    assert_equal(nil, t.parent)
    assert_equal(t, t2.parent)
    assert_equal(t2, t3.parent)
    assert_equal(t3, t4.parent)
    assert_equal(t, t5.parent)
    assert_equal(t5, t6.parent)
  end

  def test_17_all_matching_trees
    t = @c[t2 = @b[t3 = @a[t4 = @b[1]]], t5 = @b[t6 = @c[3, 4]]]

    pattern = @a[:X]
    res = pattern.all_matching_trees(t)
    assert_kind_of(Array, res)
    assert_equal(1, res.length)
    mapping, subtree = res[0]
    assert_kind_of(Hash, mapping)
    assert_equal({:X => t4}, mapping)
    assert_kind_of(Rockit::Tree::Base, subtree)
    assert_equal(t3, subtree)

    pattern = @b[:X]
    res = pattern.all_matching_trees(t)
    assert_equal(3, res.length)
    assert_equal({:X => t3}, res[0][0])
    assert_equal(t2, res[0][1])
    assert_equal({:X => 1}, res[1][0])
    assert_equal(t4, res[1][1])
    assert_equal({:X => t6}, res[2][0])
    assert_equal(t5, res[2][1])

    pattern = @c[:Left, :Right]
    res = pattern.all_matching_trees(t)
    assert_equal(2, res.length)
    assert_equal({:Left => t2, :Right => t5}, res[0][0])
    assert_equal(t, res[0][1])
    assert_equal({:Left => 3, :Right => 4}, res[1][0])
    assert_equal(t6, res[1][1])

    pattern = @b[@a[:Z]]
    res = pattern.all_matching_trees(t)
    assert_equal(1, res.length)
    assert_equal({:Z => t4}, res[0][0])
    assert_equal(t2, res[0][1])

    pattern = @b[@c[:L, :R]]
    res = pattern.all_matching_trees(t)
    assert_equal(1, res.length)
    assert_equal({:L => 3, :R => 4}, res[0][0])
    assert_equal(t5, res[0][1])
  end

  def test_18_bug_with_nested_string_child
    t = @a["y", 1]
    res = @a["y", :X].upcase_symbols_to_pattern_vars.match_tree(t)
    assert_equal({:X => 1}, res)

    t = @a[@b["y"], 1]
    res = @a[@b["y"], :X].upcase_symbols_to_pattern_vars.match_tree(t)
    assert_equal({:X => 1}, res)

    res = @a[@b["y"], :X].all_matching_trees(t)
    assert_equal(1, res.length)
  end

  def test_18_complex_tree_binop_bug
    pattern = BinOpCall[:X, "-", :Y]
    res = pattern.all_matching_trees(MR1Tree)
    assert_equal(1, res.length)

    res = BinOpCall[:X, "/", :Y].all_matching_trees(MR1Tree)
    assert_equal(1, res.length)

    res = BinOpCall[:X, "+", :Y].all_matching_trees(MR1Tree)
    assert_equal(1, res.length)

    res = BinOpCall[:X, "*", :Y].all_matching_trees(MR1Tree)
    assert_equal(0, res.length)

    res = Assign[LVar["y"], :X].all_matching_trees(MR1Tree)
    assert_equal(4, res.length)
  end

  def test_19_should_traverse_into_arrays
    # We need this since it is quite common that a collection of statements
    # are simply given in an Array. It might be a design smell though!
    tree = @a[[@b[1]]]
    bs = tree.select_trees_of_class(@b)
    assert_equal(1, bs.length)
  end
end