nested_set.ex

defmodule SetMember do
  defstruct val: 0, lt: 0, rt: 0, children: []
end

defmodule Test do

  # function header for defaults
  def assign_counts(set_member, count\\1)

  # no child just assign counts and return last used count
  def assign_counts(%SetMember{children: []} = member, count) do
    {%SetMember{member | lt: count, rt: count + 1}, count + 1}
  end


  def assign_counts(%SetMember{} = member, count) do
    # map reduce over children, accumalating count and updated children
    # see: http://elixir-lang.org/docs/v1.0/elixir/Enum.html#map_reduce/3
    {updated_children, last_used_count} =
      Enum.map_reduce(member.children, count,
        fn(child, acc) ->
          Test.assign_counts(child, acc + 1)
        end)
    # set the counts and updated children, and return
    {%SetMember{member| lt: count, rt: last_used_count + 1, children: updated_children}, last_used_count + 1 }
  end

  def run do
    # Test cases
    case1
    case2
    case3
    case4
    case5
  end

  defp case1 do
    IO.puts "With no child"
    set_member = %SetMember{children: [], val: 1}
    {updated_member, last_count} = assign_counts(set_member)
    IO.puts inspect(updated_member)
  end
  defp case2 do
    IO.puts "With one child"
    set_member = %SetMember{children: [], val: 1}
    set_member2 = %SetMember{children: [set_member], val: 2}
    {updated_member, last_count} = assign_counts(set_member2)
    IO.puts inspect(updated_member)
  end
  defp case3 do
    IO.puts "With sub-child"
    set_member2 = %SetMember{children: [], val: 2}
    set_member3 = %SetMember{children: [], val: 3}
    set_member = %SetMember{children: [set_member2, set_member3], val: 1}
    {updated_member, last_count} = assign_counts(set_member)
    IO.puts inspect(updated_member)
  end
  defp case4 do
    IO.puts "With sub-sub-child"
    set_member = %SetMember{children: [], val: 1}
    set_member2 = %SetMember{children: [set_member], val: 2}
    set_member3 = %SetMember{children: [set_member2], val: 3}
    {updated_member, last_count} = assign_counts(set_member3)
    IO.puts inspect(updated_member)
  end
  defp case5 do
    IO.puts "With multiple child and sub-sub-child"
    set_member1 = %SetMember{children: [], val: 1}
    set_member2 = %SetMember{children: [], val: 2}
    set_member3 = %SetMember{children: [], val: 3}
    set_member5 = %SetMember{children: [set_member2], val: 5}
    set_member4 = %SetMember{children: [set_member1, set_member3, set_member5], val: 4}
    {updated_member, last_count} = assign_counts(set_member4)
    IO.puts inspect(updated_member)
  end

end