1 | module Enumerable
|
---|
2 |
|
---|
3 | # = Enumerable#lazy implementation
|
---|
4 | #
|
---|
5 | # Enumerable#lazy returns an instance of Enumerator::Lazy.
|
---|
6 | # You can use it just like as normal Enumerable object,
|
---|
7 | # except these methods act as 'lazy':
|
---|
8 | #
|
---|
9 | # - map collect
|
---|
10 | # - select find_all
|
---|
11 | # - reject
|
---|
12 | # - grep
|
---|
13 | # - drop
|
---|
14 | # - drop_while
|
---|
15 | # - take_while
|
---|
16 | # - flat_map collect_concat
|
---|
17 | # - zip
|
---|
18 | def lazy
|
---|
19 | Enumerator::Lazy.new(self)
|
---|
20 | end
|
---|
21 | end
|
---|
22 |
|
---|
23 | class Enumerator
|
---|
24 | # == Acknowledgements
|
---|
25 | #
|
---|
26 | # Based on https://github.com/yhara/enumerable-lazy
|
---|
27 | # Inspired by https://github.com/antimon2/enumerable_lz
|
---|
28 | # http://jp.rubyist.net/magazine/?0034-Enumerable_lz (ja)
|
---|
29 | class Lazy < Enumerator
|
---|
30 | def initialize(obj, &block)
|
---|
31 | super(){|yielder|
|
---|
32 | begin
|
---|
33 | obj.each{|x|
|
---|
34 | if block
|
---|
35 | block.call(yielder, x)
|
---|
36 | else
|
---|
37 | yielder << x
|
---|
38 | end
|
---|
39 | }
|
---|
40 | rescue StopIteration
|
---|
41 | end
|
---|
42 | }
|
---|
43 | end
|
---|
44 |
|
---|
45 | def to_enum(meth=:each, *args, &block)
|
---|
46 | unless self.respond_to?(meth)
|
---|
47 | raise ArgumentError, "undefined method #{meth}"
|
---|
48 | end
|
---|
49 | lz = Lazy.new(self, &block)
|
---|
50 | lz.obj = self
|
---|
51 | lz.meth = meth
|
---|
52 | lz.args = args
|
---|
53 | lz
|
---|
54 | end
|
---|
55 | alias enum_for to_enum
|
---|
56 |
|
---|
57 | def map(&block)
|
---|
58 | Lazy.new(self){|yielder, val|
|
---|
59 | yielder << block.call(val)
|
---|
60 | }
|
---|
61 | end
|
---|
62 | alias collect map
|
---|
63 |
|
---|
64 | def select(&block)
|
---|
65 | Lazy.new(self){|yielder, val|
|
---|
66 | if block.call(val)
|
---|
67 | yielder << val
|
---|
68 | end
|
---|
69 | }
|
---|
70 | end
|
---|
71 | alias find_all select
|
---|
72 |
|
---|
73 | def reject(&block)
|
---|
74 | Lazy.new(self){|yielder, val|
|
---|
75 | unless block.call(val)
|
---|
76 | yielder << val
|
---|
77 | end
|
---|
78 | }
|
---|
79 | end
|
---|
80 |
|
---|
81 | def grep(pattern)
|
---|
82 | Lazy.new(self){|yielder, val|
|
---|
83 | if pattern === val
|
---|
84 | yielder << val
|
---|
85 | end
|
---|
86 | }
|
---|
87 | end
|
---|
88 |
|
---|
89 | def drop(n)
|
---|
90 | dropped = 0
|
---|
91 | Lazy.new(self){|yielder, val|
|
---|
92 | if dropped < n
|
---|
93 | dropped += 1
|
---|
94 | else
|
---|
95 | yielder << val
|
---|
96 | end
|
---|
97 | }
|
---|
98 | end
|
---|
99 |
|
---|
100 | def drop_while(&block)
|
---|
101 | dropping = true
|
---|
102 | Lazy.new(self){|yielder, val|
|
---|
103 | if dropping
|
---|
104 | if not block.call(val)
|
---|
105 | yielder << val
|
---|
106 | dropping = false
|
---|
107 | end
|
---|
108 | else
|
---|
109 | yielder << val
|
---|
110 | end
|
---|
111 | }
|
---|
112 | end
|
---|
113 |
|
---|
114 | def take(n)
|
---|
115 | if n == 0
|
---|
116 | return Lazy.new(self){raise StopIteration}
|
---|
117 | end
|
---|
118 | taken = 0
|
---|
119 | Lazy.new(self){|yielder, val|
|
---|
120 | yielder << val
|
---|
121 | taken += 1
|
---|
122 | if taken >= n
|
---|
123 | raise StopIteration
|
---|
124 | end
|
---|
125 | }
|
---|
126 | end
|
---|
127 |
|
---|
128 | def take_while(&block)
|
---|
129 | Lazy.new(self){|yielder, val|
|
---|
130 | if block.call(val)
|
---|
131 | yielder << val
|
---|
132 | else
|
---|
133 | raise StopIteration
|
---|
134 | end
|
---|
135 | }
|
---|
136 | end
|
---|
137 |
|
---|
138 | def flat_map(&block)
|
---|
139 | Lazy.new(self){|yielder, val|
|
---|
140 | ary = block.call(val)
|
---|
141 | # TODO: check ary is an Array
|
---|
142 | ary.each{|x|
|
---|
143 | yielder << x
|
---|
144 | }
|
---|
145 | }
|
---|
146 | end
|
---|
147 | alias collect_concat flat_map
|
---|
148 |
|
---|
149 | def zip(*args, &block)
|
---|
150 | enums = [self] + args
|
---|
151 | Lazy.new(self){|yielder, val|
|
---|
152 | ary = enums.map{|e| e.next}
|
---|
153 | if block
|
---|
154 | yielder << block.call(ary)
|
---|
155 | else
|
---|
156 | yielder << ary
|
---|
157 | end
|
---|
158 | }
|
---|
159 | end
|
---|
160 |
|
---|
161 | def uniq(&block)
|
---|
162 | hash = {}
|
---|
163 | Lazy.new(self){|yielder, val|
|
---|
164 | if block
|
---|
165 | v = block.call(val)
|
---|
166 | else
|
---|
167 | v = val
|
---|
168 | end
|
---|
169 | unless hash.include?(v)
|
---|
170 | yielder << val
|
---|
171 | hash[v] = val
|
---|
172 | end
|
---|
173 | }
|
---|
174 | end
|
---|
175 |
|
---|
176 | alias force to_a
|
---|
177 | end
|
---|
178 | end
|
---|