Project

General

Profile

Feature #14781 ยป enumerator_from.rb

knu (Akinori MUSHA), 10/18/2018 02:12 PM

 
1
class Enumerator
2
  class From < Enumerator
3
    def initialize(seeds, drop: 0, allow_nil: false, &block)
4
      @block = block or
5
        raise ArgumentError, "block not given"
6
      @seeds = Array.try_convert(seeds) or
7
        raise ArgumentError, "seeds must be an array"
8
      drop.is_a?(Integer) && drop >= 0 or
9
        raise ArgumentError, "drop must be a non-negative integer"
10
      @drop = drop
11
      allow_nil.equal?(true) || allow_nil.equal?(false) or
12
        raise ArgumentError, "allow_nil must be boolean"
13
      @allow_nil = allow_nil
14
    end
15

    
16
    def each
17
      drop = @drop
18
      allow_nil = @allow_nil
19
      nterms = @seeds.size
20
      block = @block
21
      if drop < nterms
22
        (drop...nterms).each { |i|
23
          value = @seeds[i]
24
          return if value.nil? && !allow_nil
25
          yield value
26
        }
27
        drop = 0
28
      else
29
        drop -= nterms
30
      end
31
      preceding_terms = [*@seeds]
32
      loop {
33
        value = block.call(*preceding_terms)
34
        break if value.nil? && !allow_nil
35
        yield value if drop == 0
36
        preceding_terms.shift
37
        preceding_terms.push(value)
38
        drop -= 1 if drop > 0
39
      }
40
    end
41

    
42
    def size
43
      nil
44
    end
45
  end
46

    
47
  # Enumerator.from(seeds, after: 0, allow_nil: false) { |*preceding_terms|
48
  #   next_term
49
  # }
50
  def self.from(seeds, drop: 0, allow_nil: false, &block)
51
    From.new(seeds, drop: drop, allow_nil: allow_nil, &block)
52
  end
53
end
54

    
55
require 'test/unit'
56
require 'ostruct'
57

    
58
class TestEnumerator < Test::Unit::TestCase
59
  def test_s_from
60
    assert_kind_of(Enumerator, Enumerator.from([1], &:succ))
61

    
62
    assert_equal([0, 0, 0, 0, 0, 0, 0, 0, 0, 0], Enumerator.from([]) { 0 }.take(10))
63

    
64
    assert_equal([1, 2, 3, 4, 5, 6, 7, 8, 9, 10], Enumerator.from([1], &:succ).take(10))
65

    
66
    assert_equal([0, 1, 1, 2, 3, 5, 8, 13, 21, 34], Enumerator.from([0, 1]) { |i, j| i + j }.take(10))
67
    assert_equal([1, 1, 2, 3, 5, 8, 13, 21, 34, 55], Enumerator.from([0, 1], drop: 1) { |i, j| i + j }.take(10))
68
    assert_equal([1, 2, 3, 5, 8, 13, 21, 34, 55, 89], Enumerator.from([0, 1], drop: 2) { |i, j| i + j }.take(10))
69
    assert_equal([2, 3, 5, 8, 13, 21, 34, 55, 89, 144], Enumerator.from([0, 1], drop: 3) { |i, j| i + j }.take(10))
70

    
71
    root = OpenStruct.new(name: "Root", parent: nil)
72
    node = ("A".."E").inject(root) { |parent, name|
73
      OpenStruct.new(name: name, parent: parent)
74
    }
75

    
76
    assert_equal(["E", "D", "C", "B", "A", "Root"], Enumerator.from([node], &:parent).map(&:name))
77
    assert_equal(["D", "C", "B", "A", "Root"], Enumerator.from([node], drop: 1, &:parent).map(&:name))
78
    assert_equal(["C", "B", "A", "Root"], Enumerator.from([node], drop: 2, &:parent).map(&:name))
79
    assert_equal(["B", "A", "Root"], Enumerator.from([node], drop: 3, &:parent).map(&:name))
80
    assert_equal(["E", "D", "C", "B", "A", "Root", nil], Enumerator.from([node], allow_nil: true) { |n|
81
      raise StopIteration if n.nil?
82
      n.parent
83
    }.map { |n| n&.name })
84
    assert_equal([], Enumerator.from([nil], &:parent).map(&:name))
85

    
86
    assert_equal([nil, nil, nil, nil, nil], Enumerator.from([nil], allow_nil: true, &:itself).take(5))
87
    assert_equal([1, nil, 1, nil, 1], Enumerator.from([1, nil], allow_nil: true) { |i, j| i }.take(5))
88
  end
89
end