Ruby enumerator/enumerable implementations, minus lazy

enumerator.rb

# FOR TESTING, REMOVE FOR ACTUAL USE
Object.send(:remove_const, :Enumerator)
Object.send(:remove_const, :Enumerable)

module Enumerable
  # Used for undefined arguments, because "nil" can be significant.
  UNDEFINED = BasicObject.new

  def all?
    return all? {|obj| obj} unless block_given?

    each do |item|
      return false if !yield item
    end
    true
  end

  def any?
    return any? {|obj| obj} unless block_given?

    each do |item|
      return true if yield item
    end
    false
  end

  def chunk
    prevretval = nil
    chunk = nil
    Enumerator.new do |y|
      each do |item|
        retval = yield item
        case retval
        when nil, :_separator
          if chunk
            y << chunk
            chunk = nil
          end
        when :_alone
          if chunk
            y << chunk
            chunk = nil
          end
          y << [retval, [item]]
        else
          if retval.is_a? Symbol && retval[0] == '_'
            raise RuntimeError, 'symbols beginning with an underscore are reserved'
          end
          if chunk
            if retval == prevretval
              chunk[1] << item
            else
              y << chunk
              chunk = [retval, [item]]
            end
          else
            chunk = [retval, [item]]
          end
        end
      end

      y << chunk if chunk
    end
  end

  def chunk_while
    chunk = [first]
    Enumerator.new do |y|
      each_cons(2) do |elt_before, elt_after|
        if yield elt_before, elt_after
          chunk << elt_after
        else
          y << chunk 
          chunk = [elt_after]
        end
      end
      y << chunk
    end
  end

  def collect
    return to_enum :collect unless block_given?

    collection = []
    each do |elem|
      collection << yield(elem)
    end
    collection
  end
  alias_method :map, :collect

  def collect_concat
    return to_enum :collect_concat unless block_given?

    collection = []
    each do |elem|
      result = yield elem
      if result.is_a? Array
        collection.concat result
      else
        collection << result
      end
    end
    collection
  end

  alias_method :flat_map, :collect_concat

  def count(item=UNDEFINED)
    count = 0
    if !item.equal?(UNDEFINED)
      each do |elt|
        count += 1 if item == elt
      end
    elsif block_given?
      each do |elt|
        count += 1 if yield elt
      end
    else
      each do |elt|
        count += 1
      end
    end
    count
  end

  def cycle(n=nil)
    n = Integer(n + 0) if n != nil

    return to_enum(:cycle, n) unless block_given?

    if n
      arr = to_a if n > 0
      (1..n).each do
        arr.each do |elem|
          yield elem
        end
      end
    else
      arr = to_a
      loop do
        arr.each do |elem|
          yield elem
        end
      end
    end
  end

  def detect(ifnone=nil)
    return to_enum(:detect, ifnone) unless block_given?

    each do |elem|
      return elem if yield elem
    end
    if ifnone
      ifnone.call
    else
      nil
    end
  end

  alias_method :find, :detect

  def drop(n)
    index = 0
    output = []
    each do |elem|
      if index >= n
        output << elem
      else
        index += 1
      end
    end
    output
  end

  def drop_while
    return to_enum :drop_while unless block_given?

    found = false
    output = []
    each do |elem|
      found = yield elem unless found
      output << elem if found
    end
    output
  end

  def each_cons(n)
    return to_enum(:each_cons, n) unless block_given?

    cons = []
    each do |elem|
      cons << elem
      if cons.size > n
        cons.unshift
      end
      if cons.size == n
        yield cons
      end
    end
  end

  def each_entry
    return to_enum :each_entry unless block_given?

    each do |*elem|
      yield case elem.size
      when 0
        nil
      when 1
        elem[0]
      else
        elem
      end
    end
  end

  def each_slice(n)
    return to_enum(:each_slice, n) unless block_given?

    slice = []
    each do |elem|
      slice << elem
      if slice.size == n
        yield slice
        slice = []
      end
    end
    yield slice unless slice.empty?
  end

  def each_with_index
    return to_enum :each_with_index unless block_given?

    index = 0
    each do |elem|
      yield elem, index
      index += 1
    end
  end
  def each_with_object(obj)
    return to_enum(:each_with_object, obj) unless block_given?

    each do |elem|
      yield elem, obj
    end
    obj
  end
  def entries
    arr = []
    each do |elem|
      arr << elem
    end
    arr
  end
  alias_method :to_a, :entries

  def find_all
    return to_enum :find_all unless block_given?

    found = []
    each do |elem|
      found << elem if yield elem
    end
    found
  end
  alias_method :select, :find_all

  def find_index(value=UNDEFINED)
    if !value.equal?(UNDEFINED)
      each_with_index do |elt, index|
        return index if elt == value
      end
      nil
    elsif block_given?
      each_with_index do |elt, index|
        return index if yield elt
      end
      nil
    else
      to_enum :find_index
    end
  end

  def first(n=UNDEFINED)
    if n == undefined
      each do |elt|
        return elt
      end
    else
      output = []
      each do |elt|
        output << elt
        return output if output.size == n
      end
      output
    end
  end

  def grep(pattern)
    output = []
    each do |elt|
      if pattern === elt
        output << if block_given?
          yield elt
        else
          elt
        end
      end
    end
    output
  end

  # Identical to grep, but with the if changed to unless
  def grep_v(pattern)
    output = []
    each do |elt|
      unless pattern === elt
        output << if block_given?
          yield elt
        else
          elt
        end
      end
    end
    output
  end

  def group_by
    return to_enum :group_by unless block_given?

    output = {}

    each do |elt|
      key = yield elt
      output[key] = [] unless output.key? key
      output[key] << elt 
    end
    output
  end

  def include?(obj)
    each do |elt|
      return true if obj == elt
    end
    false
  end

  def inject(a=UNDEFINED, b=UNDEFINED)
    if a.equal? UNDEFINED
      # No arguments
      memo = UNDEFINED
      each do |elt|
        if memo.equal? UNDEFINED
          memo = elt
        else
          memo = yield memo, elt
        end
      end
    elsif !b.equal?(UNDEFINED)
      # Two arguments
      memo = a
      sym = b
      each do |elt|
        memo = memo.send(sym, elt)
      end
    elsif block_given?
      memo = a
      each do |elt|
        memo = yield memo, elt
      end
    else
      # One argument with no block
      memo = UNDEFINED
      sym = a
      each do |elt|
        if memo.equal? UNDEFINED
          memo = elt
        else
          memo = memo.send(sym, elt)
        end
      end
    end
    if memo.equal? UNDEFINED
      nil
    else
      memo
    end
  end

  alias_method :reduce, :inject

  def max(n=nil, &block)
    if n.nil?
      max(1, &block)[0]
    else
      block ||= Proc.new do |a, b|
        a <=> b
      end

      maxes = []
      each do |elt|
        if index = maxes.bsearch_index{|i| block.call(elt, i) > 0}
          maxes.insert(index, elt)
          maxes.pop if maxes.size > n
        elsif maxes.size < n
          maxes << elt
        end
      end
      maxes
    end
  end

  def max_by(n=nil, &block)
    return to_enum(:max_by, n) unless block_given?

    if n.nil?
      max_by(1, &block)[0]
    else
      # Array of [sortkey, value]
      maxes = []
      each do |elt|
        eltkey = block.call(elt)
        if index = maxes.bsearch_index{|i| eltkey > i[0]}
          maxes.insert(index, [eltkey, elt])
          maxes.pop if maxes.size > n
        elsif maxes.size < n
          maxes << [eltkey, elt]
        end
      end
      # Strip out the sort keys
      maxes.map(&:last)
    end
  end

  def min(n=nil, &block)
    block ||= Proc.new do |a, b|
      a <=> b
    end
    # Just do a reverse max
    max(n) do |a, b|
      block.call(b, a)
    end
  end

  def min_by(n=nil, &block)
    return to_enum(:min_by, n) unless block_given?

    if n.nil?
      min_by(1, &block)[0]
    else
      # Array of [sortkey, value]
      maxes = []
      each do |elt|
        eltkey = block.call(elt)
        if index = maxes.bsearch_index{|i| eltkey < i[0]}
          maxes.insert(index, [eltkey, elt])
          maxes.pop if maxes.size > n
        elsif maxes.size < n
          maxes << [eltkey, elt]
        end
      end
      # Strip out the sort keys
      maxes.map(&:last)
    end
  end

  def minmax(&block)
    # We do a specialized implementation instead of calling both min and max to
    # only iterate once
    block ||= Proc.new do |a, b|
      a <=> b
    end
    min = first
    max = min
    each do |elt|
      if block.call(elt, max) > 0
        max = elt
      end
      if block.call(elt, min) < 0
        min = elt
      end
    end
    [min, max]
  end

  def minmax_by
    return to_enum :minmax_by unless block_given?
    # [sortkey, value]
    minmax_first = first
    min = [yield(minmax_first), minmax_first]
    max = min

    each do |elt|
      sortkey = yield elt
      if sortkey > max[0]
        max = [sortkey, elt]
      end
      if sortkey < max[0]
        min = [sortkey, elt]
      end
    end
    [min[1], max[1]]
  end

  def none?(&block)
    !any?(&block)
  end

  def one?
    return one? {|obj| obj} unless block_given?

    found = false
    each do |item|
      if yield item
        if found
          return false
        else
          found = true
        end
      end
    end
    found
  end

  def partition
    return to_enum :partition unless block_given?

    true_array = []
    false_array = []
    each do |elt|
      if yield elt
        true_array << elt
      else
        false_array << elt
      end
    end
    [true_array, false_array]
  end

  def reject
    return to_enum :reject unless block_given?

    found = []
    each do |elem|
      found << elem if !yield elem
    end
    found
  end

  def reverse_each
    return to_enum :reverse_each unless block_given?

    to_a.reverse do |elt|
      yield elt
    end
  end

  def slice_after(pattern=nil, &block)
    if pattern and block_given?
      raise ArgumentError, 'both pattern and block are given'
    end

    Enumerator.new do |y|
      if pattern
        block = Proc.new do |elt|
          pattern === elt
        end
      end
      chunk = []
      each do |elt|
        chunk << elt
        # End the slice
        if block.call(elt)
          y << chunk
          chunk = []
        end
      end
      y << chunk unless chunk.empty?
    end
  end

  def slice_before(pattern=nil, &block)
    if pattern and block_given?
      raise ArgumentError, 'both pattern and block are given'
    end

    Enumerator.new do |y|
      if pattern
        block = Proc.new do |elt|
          pattern === elt
        end
      end
      chunk = []
      each do |elt|
        # elt begins a new slice
        if block.call(elt)
          y << chunk unless chunk.empty?
          chunk = []
        end
        chunk << elt
      end
      y << chunk unless chunk.empty?
    end
  end

  def slice_when
    chunk = [first]
    Enumerator.new do |y|
      each_cons(2) do |elt_before, elt_after|
        if yield elt_before, elt_after
          y << chunk 
          chunk = [elt_after]
        else
          chunk << elt_after
        end
      end
      y << chunk
    end
  end

  def sort(&block)
    block ||= Proc.new do |a, b|
      a <=> b
    end
    sorted = []
    each do |elt|
      if index = maxes.bsearch_index{|i| block.call(elt, i) < 0}
        sorted.insert(index, elt)
      else
        sorted << elt
      end
    end
    sorted
  end

  def sort_by
    return to_enum :sort_by unless block_given?

    # [sort_key, value]
    sorted = []
    each do |elt|
      eltkey = block.call(elt)
      if index = sorted.bsearch_index{|i| eltkey < i[0]}
        sorted.insert(index, [eltkey, elt])
      else
        sorted << [eltkey, elt]
      end
    end
    sorted.map(&:last)
  end

  def take(n)
    first(n)
  end

  def take_while
    return to_enum :take_while unless block_given?

    output = []
    each do |elt|
      if yield elt
        output << elt
      else
        return output
      end
    end
    output
  end

  def to_h
    hash = {}
    each do |key, value|
      hash[key] = value
    end
  end

  def zip(*args)
    output = [] unless block_given?
    # Get enumerators for successive "next" calls.  Can't just call each without
    # an argument, because we don't know if that will reliably return an
    # enumerator
    enums = args.map do |arg|
      Enumerator.new do |y|
        arg.each do |elt|
          y << elt
        end
      end
    end
    each do |elt|
      arr = [elt]
      # Call next once on each enum
      enums.each do |enum|
        arr << begin
          enum.next
        rescue StopIteration
          nil
        end
      end
      if block_given?
        yield arr
      else
        output << arr
      end
    end
    output
  end
end

class Enumerator
  include Enumerable

  attr_accessor :size

  class Yielder
    attr_accessor :feed

    def initialize(appending_args=[], &block)
      @block = block
      @appending_args = appending_args
    end
    def yield(*args)
      if @block
        @block.call(*(args + @appending_args))
      else
        Fiber.yield (args + @appending_args)
      end

      # To interact with Enumerator#feed
      feed, @feed = @feed, nil
      feed
    end
    alias_method :<<, :yield
  end

  def initialize(*args, &block)
    # Redirect to the correct constructor
    if block_given?
      initialize_block(*args, &block)
    else
      initialize_meth(*args)
    end
  end

  def initialize_block(size = nil, &block)
    @size = size
    @block = block
  end

  def initialize_meth(obj, method=:each, *args)
    @block = Proc.new do |y|
      obj.method(method).(*args) do |*values|
        y.yield(*values)
      end
    end
  end

  def each(*appending_args, &block)
    if block_given?
      yielder = Yielder.new(appending_args, &block)
      @block.call(yielder)
      yielder
    else
      self
    end
  end

  def each_with_index(&block)
    with_index(&block)
  end

  def feed(val)
    @yielder.feed = val
  end

  def next
    val = next_values
    case val.size
    when 0
      nil
    when 1
      val[0]
    else
      val
    end
  end

  def next_values
    if @peek_values
      # If we've peeked, use the stored peek instead
      peek_values, @peek_values = @peek_values, nil
      peek_values
    else
      @yielder ||= Yielder.new
      @fiber ||= Fiber.new do
        @block.call(@yielder)
      end
      begin
        values = @fiber.resume
      rescue FiberError
        values = nil
      end
      # Values should always exist on a yield, because this will always return an
      # array
      if not values
        raise StopIteration, 'iteration reached an end'
      end
      values
    end
  end

  def peek
    val = peek_values
    case val.size
    when 0
      nil
    when 1
      val[0]
    else
      val
    end
  end

  def peek_values
    if @peek_values
      @peek_values
    else
      @peek_values = next_values
    end
  end

  def rewind
    @fiber = nil
    @yielder = nil
    @peek_values = nil
  end

  def with_index(offset = 0, &block)
    begin
      index = Integer(offset + 0)
    rescue
      raise TypeError, 'Could not convert offset to integer'
    end
    Enumerator.new do |y|
      each do |*value|
        y.yield(*value, index)
        index += 1
      end
    end.each(&block)
  end

  def with_object(obj, &block)
    Enumerator.new do |y|
      each do |*value|
        y.yield(*value, obj)
        index += 1
      end
    end.each(&block)
  end
  alias_method :each_with_object, :with_object
end