1 | module Enumerable
|
---|
2 |
|
---|
3 | # = Enumerable#lazy implementation
|
---|
4 | #
|
---|
5 | # Enumerable#lazy returns an instance of Enumerable::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 | Lazy.new(self)
|
---|
20 | end
|
---|
21 |
|
---|
22 | # == Acknowledgements
|
---|
23 | #
|
---|
24 | # Based on https://github.com/yhara/enumerable-lazy
|
---|
25 | # Inspired by https://github.com/antimon2/enumerable_lz
|
---|
26 | # http://jp.rubyist.net/magazine/?0034-Enumerable_lz (ja)
|
---|
27 | class Lazy < Enumerator
|
---|
28 | def initialize(obj, &block)
|
---|
29 | super(){|yielder|
|
---|
30 | begin
|
---|
31 | obj.each{|x|
|
---|
32 | if block
|
---|
33 | block.call(yielder, x)
|
---|
34 | else
|
---|
35 | yielder << x
|
---|
36 | end
|
---|
37 | }
|
---|
38 | rescue StopIteration
|
---|
39 | end
|
---|
40 | }
|
---|
41 | end
|
---|
42 |
|
---|
43 | def map(&block)
|
---|
44 | Lazy.new(self){|yielder, val|
|
---|
45 | yielder << block.call(val)
|
---|
46 | }
|
---|
47 | end
|
---|
48 | alias collect map
|
---|
49 |
|
---|
50 | def select(&block)
|
---|
51 | Lazy.new(self){|yielder, val|
|
---|
52 | if block.call(val)
|
---|
53 | yielder << val
|
---|
54 | end
|
---|
55 | }
|
---|
56 | end
|
---|
57 | alias find_all select
|
---|
58 |
|
---|
59 | def reject(&block)
|
---|
60 | Lazy.new(self){|yielder, val|
|
---|
61 | if not block.call(val)
|
---|
62 | yielder << val
|
---|
63 | end
|
---|
64 | }
|
---|
65 | end
|
---|
66 |
|
---|
67 | def grep(pattern)
|
---|
68 | Lazy.new(self){|yielder, val|
|
---|
69 | if pattern === val
|
---|
70 | yielder << val
|
---|
71 | end
|
---|
72 | }
|
---|
73 | end
|
---|
74 |
|
---|
75 | def drop(n)
|
---|
76 | dropped = 0
|
---|
77 | Lazy.new(self){|yielder, val|
|
---|
78 | if dropped < n
|
---|
79 | dropped += 1
|
---|
80 | else
|
---|
81 | yielder << val
|
---|
82 | end
|
---|
83 | }
|
---|
84 | end
|
---|
85 |
|
---|
86 | def drop_while(&block)
|
---|
87 | dropping = true
|
---|
88 | Lazy.new(self){|yielder, val|
|
---|
89 | if dropping
|
---|
90 | if not block.call(val)
|
---|
91 | yielder << val
|
---|
92 | dropping = false
|
---|
93 | end
|
---|
94 | else
|
---|
95 | yielder << val
|
---|
96 | end
|
---|
97 | }
|
---|
98 | end
|
---|
99 |
|
---|
100 | def take(n)
|
---|
101 | if n == 0
|
---|
102 | return Lazy.new(self){raise StopIteration}
|
---|
103 | end
|
---|
104 | taken = 0
|
---|
105 | Lazy.new(self){|yielder, val|
|
---|
106 | yielder << val
|
---|
107 | taken += 1
|
---|
108 | if taken >= n
|
---|
109 | raise StopIteration
|
---|
110 | end
|
---|
111 | }
|
---|
112 | end
|
---|
113 |
|
---|
114 | def take_while(&block)
|
---|
115 | Lazy.new(self){|yielder, val|
|
---|
116 | if block.call(val)
|
---|
117 | yielder << val
|
---|
118 | else
|
---|
119 | raise StopIteration
|
---|
120 | end
|
---|
121 | }
|
---|
122 | end
|
---|
123 |
|
---|
124 | def flat_map(&block)
|
---|
125 | Lazy.new(self){|yielder, val|
|
---|
126 | ary = block.call(val)
|
---|
127 | # TODO: check ary is an Array
|
---|
128 | ary.each{|x|
|
---|
129 | yielder << x
|
---|
130 | }
|
---|
131 | }
|
---|
132 | end
|
---|
133 | alias collect_concat flat_map
|
---|
134 |
|
---|
135 | def zip(*args, &block)
|
---|
136 | enums = [self] + args
|
---|
137 | Lazy.new(self){|yielder, val|
|
---|
138 | ary = enums.map{|e| e.next}
|
---|
139 | if block
|
---|
140 | yielder << block.call(ary)
|
---|
141 | else
|
---|
142 | yielder << ary
|
---|
143 | end
|
---|
144 | }
|
---|
145 | end
|
---|
146 |
|
---|
147 | alias force to_a
|
---|
148 | end
|
---|
149 | end
|
---|