Project

General

Profile

Feature #4890 ยป lazy.rb

yhara (Yutaka HARA), 06/16/2011 07:23 PM

 
1
# = Enumerable#lazy example implementation
2
#
3
# Enumerable#lazy returns an instance of Enumerable::Lazy.
4
# You can use it just like as normal Enumerable object,
5
# except these methods act as 'lazy':
6
#
7
#   - map       collect
8
#   - select    find_all
9
#   - reject
10
#   - grep
11
#   - drop
12
#   - drop_while
13
#   - take_while
14
#   - flat_map  collect_concat
15
#   - zip
16
#
17
# == Example
18
#
19
# This code prints the first 100 primes.
20
#
21
#   require 'prime'
22
#   INFINITY = 1.0 / 0
23
#   p (1..INFINITY).lazy.map{|n| n**2+1}.
24
#                     select{|m| m.prime?}.take(100)
25
#
26
# == Acknowledgements
27
#
28
#   Inspired by https://github.com/antimon2/enumerable_lz
29
#   http://jp.rubyist.net/magazine/?0034-Enumerable_lz (ja)
30

    
31
module Enumerable
32
  def lazy
33
    Lazy.new(self)
34
  end
35

    
36
  class Lazy < Enumerator
37
    def initialize(obj, &block)
38
      super(){|yielder|
39
        begin
40
          obj.each{|x|
41
            if block
42
              block.call(yielder, x)
43
            else
44
              yielder << x
45
            end
46
          }
47
        rescue StopIteration
48
        end
49
      }
50
    end
51

    
52
    def map(&block)
53
      Lazy.new(self){|yielder, val|
54
        yielder << block.call(val)
55
      }
56
    end
57
    alias collect map
58

    
59
    def select(&block)
60
      Lazy.new(self){|yielder, val|
61
        if block.call(val)
62
          yielder << val
63
        end
64
      }
65
    end
66
    alias find_all select
67

    
68
    def reject(&block)
69
      Lazy.new(self){|yielder, val|
70
        if not block.call(val)
71
          yielder << val
72
        end
73
      }
74
    end
75

    
76
    def grep(pattern)
77
      Lazy.new(self){|yielder, val|
78
        if pattern === val
79
          yielder << val
80
        end
81
      }
82
    end
83

    
84
    def drop(n)
85
      dropped = 0
86
      Lazy.new(self){|yielder, val|
87
        if dropped < n
88
          dropped += 1
89
        else
90
          yielder << val
91
        end
92
      }
93
    end
94

    
95
    def drop_while(&block)
96
      dropping = true
97
      Lazy.new(self){|yielder, val|
98
        if dropping
99
          if not block.call(val)
100
            yielder << val
101
            dropping = false
102
          end
103
        else
104
          yielder << val
105
        end
106
      }
107
    end
108
    
109
    # def take(n)
110
    # def first(n=1)
111
    #
112
    # These methods are intentionally omitted, so that
113
    # we can print the result like
114
    #
115
    #  p obj.lazy.map{...}.take(10)
116
    #  # => prints the values, since take(10) returns an Array
117
    #  # instead of an instance of Enumerable::Lazy
118
    #
119
    # This means that map or select after take(n) 
120
    # (eg. obj.lazy.take(10).map{...})
121
    # is not be the lazy-version.
122
    # Since take(n) returns limited number of elements,
123
    # this will not be a problem.
124

    
125
    def take_while(&block)
126
      Lazy.new(self){|yielder, val|
127
        if block.call(val)
128
          yielder << val
129
        else
130
          raise StopIteration
131
        end
132
      }
133
    end
134

    
135
    def flat_map(&block)
136
      Lazy.new(self){|yielder, val|
137
        ary = block.call(val)
138
        # TODO: check ary is an Array
139
        ary.each{|x|
140
          yielder << x
141
        }
142
      }
143
    end
144
    alias collect_concat flat_map
145

    
146
    def zip(*args, &block)
147
      enums = [self] + args
148
      Lazy.new(self){|yielder, val|
149
        ary = enums.map{|e| e.next}
150
        if block
151
          yielder << block.call(ary)
152
        else
153
          yielder << ary
154
        end
155
      }
156
    end
157

    
158
    # def chunk
159
    # def slice_before
160
    #
161
    # There methods are already implemented with Enumerator.
162

    
163
  end
164
end
165

    
166
# Example
167

    
168
# -- Print the first 100 primes
169
#require 'prime'
170
#p (1..1.0/0).lazy.select{|m|m.prime?}.first(100)
171

    
172
# -- Print the first 10 word from a text file
173
#File.open("english.txt"){|f|
174
#  p f.lines.lazy.flat_map{|line| line.split}.take(10)
175
#}
176

    
177
# -- Example of cycle and zip
178
#e1 = [1, 2, 3].cycle
179
#e2 = [:a, :b].cycle
180
#p e1.lazy.zip(e2).take(10)
181

    
182
# -- Example of chunk and take_while
183
#p Enumerator.new{|y|
184
#  loop do
185
#    y << rand(100)
186
#  end
187
#}.chunk{|n| n.even?}.
188
#  lazy.map{|even, ns| ns}.
189
#  take_while{|ns| ns.length <= 5}.to_a