Iterator#flatten implementation, it can specify except type

iterator-flatten.cr

module Iterator
  private class Flatten(I, T)
    include Iterator(T)

    @iterator : I
    @to_rewind : Array(I)
    @generators : Array(I)

    def initialize(@iterator)
      @generators = [@iterator]
      @to_rewind = [] of I
    end

    def next
      value = @generators.last.next
      case value
      when T
        value
      when Iterator
        @generators.push value
        self.next
      when Iterable
        @generators.push value.each
        self.next
      when Stop
        if @generators.size == 1
          stop
        else
          @to_rewind.push @generators.pop
          self.next
        end
      else
        raise "BUG: unreachable"
      end
    end

    def rewind
      @generators.each &.rewind
      @generators = [@iterator]
      @to_rewind.each &.rewind
      @to_rewind.clear
    end

    def self.iterator_type(iter)
      case iter
      when Iterator
        iter || iterator_type(iter.next)
      when Iterable
        iterator_type iter.each
      else
        raise ""
      end
    end

    def self.element_type(element, except : E.class) forall E
      case element
      when E
        element
      when Iterator::Stop
        raise ""
      when Iterator
        element_type(element.next, except)
      when Iterable
        element_type(element.each, except)
      else
        element
      end
    end
  end

  def flatten(except : E.class = NoReturn) forall E
    Flatten(typeof(Flatten.iterator_type(self)), typeof(Flatten.element_type(self, except))).new self
  end
end