木構造を再帰的に走査するためのコレクションを作成する

やりたいこと

class Node
  attr_reader :name
  attr_accessor :children

  def initialize(name, children = [])
    @name = name
    @children = children
  end
end

a = Node.new('A')
a1 = Node.new('A-1')
a11 = Node.new('A-1-1')
a12 = Node.new('A-1-2')
a2 = Node.new('A-2')
a1.children += [a11, a12]
a.children += [a1, a2]

というコードで木構造を再現した場合に、各ノードを走査するための仕組みがほしい。

実現方法

class Nodes
  include Enumerable

  attr_reader :nodes

  def initialize(*nodes)
    @nodes = nodes
  end

  def each
    return enum_for(:each) unless block_given?

    nodes.each do |node|
      yield(node)
      self.class.new(*node.children).each do |child|
        yield(child)
      end
    end
  end
end

nodes = Nodes.new(a)
nodes.each.with_index(1) { |node, i| puts("#{i}: #{node.name}") }

1: A
2: A-1
3: A-1-1
4: A-1-2
5: A-2