-
Notifications
You must be signed in to change notification settings - Fork 3
/
itertools.lisp
280 lines (251 loc) · 11.5 KB
/
itertools.lisp
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
(in-package :snakes)
(defgenerator icount (n &optional (step 1))
"Make a generator that returns evenly spaced values starting with n. Step can be fractional. Eg:
(counter 10) -> 10 11 12 13 14 ...
(counter 2.5 0.5) -> 2.5 3.0 3.5 ..."
(unless (numberp n)
(error "n must be a number"))
(unless (numberp step)
(error "step must be a number"))
(loop for i from n by step
do (yield i)))
(defgenerator cycle (list-or-gen)
"Repeatedly cycles through a list or generator, yielding its elements indefinitely. Eg:
(cycle '(a b c d)) -> a b c d a b c d a b ...
Note that this tool may consume considerable storage if the source iterable is long."
(let ((data
(cond
((generatorp list-or-gen)
(let ((accum nil))
(do-generator (g list-or-gen)
(yield g)
(push g accum))
(nreverse accum)))
((listp list-or-gen) list-or-gen)
(t (error "Gen must be a list or generator")))))
(loop do
(dolist (x data)
(yield x)))))
(defgenerator repeat (item &optional n)
"Returns a single item indefinitely, or up to n times."
(if n
(dotimes (i n)
(yield item))
(loop do (yield item))))
(defgenerator chain (&rest generators)
"Chains a number of generators together head to tail. Eg:
(chain (list->generator '(1 2 3) (list->generator '(6 7))) -> 1 2 3 6 7"
(dolist (g generators)
(yield-all g)))
(defgenerator izip (&rest generators)
"Steps through multiple generators in parallel, emitting each of their items as values. Stops with the end of the shortest source generator."
(do-generator-value-list (vars (apply #'multi-gen generators))
(apply #'yield (mapcar #'car vars))))
(defgenerator izip-longest (&rest generators-and-fill-value)
"Steps through multiple generators in parallel, emitting each of their items as values. Keeps going until the end of the longest source generator, padding the rest out with the value specified by :fill-value"
(multiple-value-bind (keywords gens)
(extract-keywords '(:fill-value) generators-and-fill-value)
(let ((fillspec (if keywords
(cdr (assoc :fill-value keywords))
nil)))
(do-generator-value-list
(vars
(apply #'multi-gen
(loop for g in gens
collect (list g :fill-value fillspec))))
(apply #'yield (mapcar #'car vars))))))
(defun enumerate (generator)
(izip (icount 0) generator))
(defgenerator compress (data selectors)
"Moves through the data and selectors generators in parallel, only yielding elements of data that are paired with a value from selectors that evaluates to true. Eg:
(compress (list->generator '(a b c d e f))
(list->generator '(t nil t nil t t))) -> a c e f"
(do-generator (d s (izip data selectors))
(when s
(yield d))))
(defgenerator dropwhile (predicate generator)
"Goes through generator, dropping the items until the predicate fails, then emits the rest of the source generator's items. Eg:
(dropwhile (lambda (x) (< x 5))
(list->generator '(1 4 6 4 1))) -> 6 4 1"
(do-generator (item generator)
(unless (funcall predicate item)
(yield item)
(return)))
(yield-all generator))
(defgenerator takewhile (predicate generator)
"Emits items from generator until predicate fails. Eg:
(takewhile (lambda (x) (< x 5))
(list->generator '(1 4 6 4 1))) -> 1 4"
(do-generator (item generator)
(unless (funcall predicate item)
(return))
(yield item)))
;Differs from python version on some points.
(defgenerator groupby (generator &key (key #'identity) (comp #'eq))
"Groups the items from generator by the results of the key func. Each item emitted by groupby will be a pair: the key value, followed by a list of items that produced the key value when passed to the key function. Note: groupby does NOT sort the source generator. In other words it will only group matching items that are adjacent. To get SQL style GROUP BY functionality, sort the source generator before passing it to groupby. Eg:
(groupby (list->generator '(A A A A B B B C C))) ->
(A (A A A A)) (B (B B B)) ..."
(let ((last (gensym))
(res nil))
(do-generator (g generator)
(let ((currkey (funcall key g)))
(if (funcall comp last currkey)
(push g res)
(progn
(when res (yield last (nreverse res)))
(setf last currkey)
(setf res (cons g nil))))))))
(defgenerator ifilter (predicate generator)
"Emits only those items in generator that are true by predicate. The generator equivalent of remove-if-not"
(do-generator (g generator)
(when (funcall predicate g)
(yield g))))
(defgenerator ifilter-false (predicate generator)
"Emits only those items in generator that are false by predicate. The generator equivalent of remove-if"
(do-generator (g generator)
(unless (funcall predicate g)
(yield g))))
(defgenerator islice
(generator stop-or-start
&optional (stop nil has-stop-p) step)
"Emits a subrange of generator. If only stop-or-start is set, islice will emit up to it and stop. If both stop-or-start and stop are set, islice will emit the stretch between stop and start. If step is set, then islice will emit every step-th item between start and stop. If stop is set to nil, then islice will continue through the source generator until it terminates. Eg:
(islice (list->generator '(a b c d e f g)) 2) -> A B
(islice (list->generator '(a b c d e f g)) 2 4) -> C D
(islice (list->generator '(a b c d e f g)) 2 nil) -> C D E F G
(islice (list->generator '(a b c d e f g)) 0 nil 2) -> A C E G"
(let
((start (if has-stop-p stop-or-start nil))
(stop (or stop (unless has-stop-p stop-or-start)))
(step (or step 1)))
(when start
(consume start generator))
(if stop
(dotimes (i (floor (/ (- stop (or start 0)) step)))
(let ((data (take step generator :fail-if-short nil)))
(unless data (return))
(yield (car data))))
(loop do
(let ((data (take step generator :fail-if-short nil)))
(unless data (return))
(yield (car data)))))))
(defgenerator imap (function &rest generators)
"The generator equivalent of the mapcar function. Applies function to the
values of the supplied generators, emitting the result. Eg:
(imap #'* (list->generator '(2 3 4)) (list->generator '(4 5 6))) -> 8 15 24"
(do-generator-value-list (vals (apply #'multi-gen generators))
(yield (apply function (mapcar #'car vals)))))
(defgenerator starmap (function generator)
"Sequentially applies the function to the output of the generator. Like imap, but assumes that the contents of the generator are already merged. Eg:
(starmap #'expt (list->generator '((2 5) (3 2) (10 3)))) -> 32 9 1000"
(do-generator (val generator)
(yield (apply function val))))
(defun append! (lst obj)
(setf (cdr (last lst)) obj))
(defun tee (generator &optional (n 2))
"Creates independent copies of generator, returned as values. If the child generators are consumed at different times, tee will store all of the items from the least consumed child generator through to the most. It can, if used incautiously, require considerable memory.
Note also that this implementation of tee does not create independent copies of the parent items. Modifying items from a child generator and tampering with the parent generator have undefined consequences."
(let ((stor (cons :x nil))) ;First nil is a dummy value
(labels ((get-next ()
(let ((curr (if-generator (g generator)
(cons g nil)
(cons (car g) nil))))
(append! stor curr)
(setf stor curr))))
(apply #'values
(with-collectors (gens<)
(dotimes (i n)
(gens<
(let ((ldata stor))
(gen-lambda-with-sticky-stop ()
(unless (cdr ldata)
(get-next))
(setf ldata (cdr ldata))
(if (eq (car ldata) 'generator-stop)
(sticky-stop)
(apply #'values (ensure-list (car ldata)))))))))))))
(defgenerator product (&rest lists)
"Iterates through each of the supplied lists with nested dolists, yielding an element from each for every iteration of the innermost loop. The yielded values cycle in an odometer-like fashion, with the first value changing seldom and the last changing on every yield. Eg:
(product '(a b c d) '(x y)) -(values)-> (a x) (a y) (b x) (b y) (c x) (c y)...
"
(labels ((listholder (lists)
(if (null lists)
(lambda (stack)
(apply #'yield (reverse stack)))
(let ((nextfunc
(listholder (cdr lists))))
(lambda (stack)
(dolist (itm (car lists))
(funcall nextfunc (cons itm stack))))))))
(funcall (listholder lists) nil)))
(defgenerator permutations (list &optional r)
"Emits every possible permutation of the items in list in a set of size r. If r is not specified, the size of list is used. Eg:
(permutations '(a b c d) 2) -(values)-> (a b) (a c) (a d) (b a) (b c) (b d)
(c a) (c b) (c d) (d a) (d b) (d c)"
(let ((r (or r (length list))))
(labels ((proc (data output)
(if (= r (length output))
(apply #'yield (reverse output))
(dotimes (i (length data))
(proc (xsubseq data i i :type 'list)
(cons (elt data i) output))))))
(proc list nil))))
(defgenerator combinations (list r)
"Emits every possible combination of the items in list in sets of size r. Eg:
(combinations '(a b c d) 2) -(values)-> (a b) (a c) (a d) (b c) (b d) (c d)"
(labels ((proc (data output r)
(if (zerop r)
(apply #'yield (reverse output))
(loop for d on data do
(proc (cdr d)
(cons (car d) output)
(1- r))))))
(if (< (length list) r)
(lambda () 'generator-stop)
(proc list nil r))))
(defgenerator combinations-with-replacement (list r)
"Emits every possible combination, including repetitions, of the items in list in sets of size r. Eg:
(combinations-with-replacement '(a b c) 2) -(values)-> (a a) (a b) (a c)
(b b) (b c) (c c)"
(labels ((proc (data output r)
(if (zerop r)
(apply #'yield (reverse output))
(loop for d on data do
(proc d
(cons (car d) output)
(1- r))))))
(if (< (length list) r)
(lambda () 'generator-stop)
(proc list nil r))))
(defun reduce-generator (function generator
&key (initial-value (funcall generator) initial-value-p))
"Reduce items from generator using a binary operation
FUNCTION. Reduction is left-associative. Optional INITIAL-VALUE is
logically placed before items in the generator. Eg:
(foldl #'+ (list->generator '(1 2 3 4))) -> 10
For consistency with CL:REDUCE do the following if the generator has insufficient items:
* If GENERATOR has no values, return result of FUNCTION called without arguments.
* If GENERATOR has exactly one item, return that item."
(labels ((foldl (x)
(let ((value (funcall generator)))
(if (eq value 'generator-stop) x
(foldl (funcall function x value))))))
(if (and (not initial-value-p)
(eq initial-value 'generator-stop))
(funcall function)
(foldl initial-value))))
(defgenerator take-as-generator (n generator &key (fail-if-short t))
"This function has the same meaning as TAKE but returns a generator
instead of a list. Eg:
(generator->list (take-as-generator 10 (list->generator '(1 2 3)))) -> Signals a condition
(generator->list (take-as-generator 10 (list->generator '(1 2 3)) :fail-if-short nil)) -> (1 2 3)
(generator->list (take-as-generator 3 (repeat 1))) -> (1 1 1)"
(declare (type (integer 1) n))
(let ((count 0))
(do-generator (g generator)
(yield g)
(incf count)
(when (= count n)
(yield 'generator-stop))))
(when fail-if-short
(error 'insufficient-items
:format-control "Insufficient items in generator")))