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 | lz = Lazy.new(self, &block)
|
---|
47 | lz.obj = self
|
---|
48 | lz.meth = meth
|
---|
49 | lz.args = args
|
---|
50 | lz
|
---|
51 | end
|
---|
52 | alias enum_for to_enum
|
---|
53 |
|
---|
54 | def map(&block)
|
---|
55 | Lazy.new(self){|yielder, val|
|
---|
56 | yielder << block.call(val)
|
---|
57 | }
|
---|
58 | end
|
---|
59 | alias collect map
|
---|
60 |
|
---|
61 | def select(&block)
|
---|
62 | Lazy.new(self){|yielder, val|
|
---|
63 | if block.call(val)
|
---|
64 | yielder << val
|
---|
65 | end
|
---|
66 | }
|
---|
67 | end
|
---|
68 | alias find_all select
|
---|
69 |
|
---|
70 | def reject(&block)
|
---|
71 | Lazy.new(self){|yielder, val|
|
---|
72 | if not block.call(val)
|
---|
73 | yielder << val
|
---|
74 | end
|
---|
75 | }
|
---|
76 | end
|
---|
77 |
|
---|
78 | def grep(pattern)
|
---|
79 | Lazy.new(self){|yielder, val|
|
---|
80 | if pattern === val
|
---|
81 | yielder << val
|
---|
82 | end
|
---|
83 | }
|
---|
84 | end
|
---|
85 |
|
---|
86 | def drop(n)
|
---|
87 | dropped = 0
|
---|
88 | Lazy.new(self){|yielder, val|
|
---|
89 | if dropped < n
|
---|
90 | dropped += 1
|
---|
91 | else
|
---|
92 | yielder << val
|
---|
93 | end
|
---|
94 | }
|
---|
95 | end
|
---|
96 |
|
---|
97 | def drop_while(&block)
|
---|
98 | dropping = true
|
---|
99 | Lazy.new(self){|yielder, val|
|
---|
100 | if dropping
|
---|
101 | if not block.call(val)
|
---|
102 | yielder << val
|
---|
103 | dropping = false
|
---|
104 | end
|
---|
105 | else
|
---|
106 | yielder << val
|
---|
107 | end
|
---|
108 | }
|
---|
109 | end
|
---|
110 |
|
---|
111 | def take(n)
|
---|
112 | if n == 0
|
---|
113 | return Lazy.new(self){raise StopIteration}
|
---|
114 | end
|
---|
115 | taken = 0
|
---|
116 | Lazy.new(self){|yielder, val|
|
---|
117 | yielder << val
|
---|
118 | taken += 1
|
---|
119 | if taken >= n
|
---|
120 | raise StopIteration
|
---|
121 | end
|
---|
122 | }
|
---|
123 | end
|
---|
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 | alias force to_a
|
---|
159 | end
|
---|
160 | end
|
---|