]> gitweb.factorcode.org Git - factor.git/blob - core/assocs/assocs.factor
assocs: refactor collect-by to use collect-by!
[factor.git] / core / assocs / assocs.factor
1 ! Copyright (C) 2007, 2010 Daniel Ehrenberg, Slava Pestov
2 ! See http://factorcode.org/license.txt for BSD license.
3 USING: accessors arrays kernel math sequences sequences.private
4 vectors ;
5 IN: assocs
6
7 MIXIN: assoc
8
9 GENERIC: at* ( key assoc -- value/f ? )
10 GENERIC: value-at* ( value assoc -- key/f ? )
11 GENERIC: set-at ( value key assoc -- )
12 GENERIC: new-assoc ( capacity exemplar -- newassoc )
13 GENERIC: delete-at ( key assoc -- )
14 GENERIC: clear-assoc ( assoc -- )
15 GENERIC: assoc-size ( assoc -- n )
16 GENERIC: assoc-like ( assoc exemplar -- newassoc )
17 GENERIC: assoc-clone-like ( assoc exemplar -- newassoc )
18 GENERIC: >alist ( assoc -- newassoc )
19 GENERIC: keys ( assoc -- keys )
20 GENERIC: values ( assoc -- values )
21 GENERIC: unzip ( assoc -- keys values )
22
23 M: assoc assoc-like drop ; inline
24
25 : ?at ( key assoc -- value/key ? )
26     2dup at* [ 2nip t ] [ 2drop f ] if ; inline
27
28 : maybe-set-at ( value key assoc -- changed? )
29     3dup at* [ = [ 3drop f ] [ set-at t ] if ] [ 2drop set-at t ] if ;
30
31 <PRIVATE
32
33 : assoc-operator ( assoc quot -- alist quot' )
34     [ >alist ] dip [ first2 ] prepose ; inline
35
36 : assoc-stack-from ( key i seq -- value/f )
37     over 0 < [
38         3drop f
39     ] [
40         3dup nth-unsafe at*
41         [ 3nip ] [ drop [ 1 - ] dip assoc-stack-from ] if
42     ] if ; inline recursive
43
44 : search-alist ( key alist -- pair/f i/f )
45     [ first = ] with find swap ; inline
46
47 : substituter ( assoc -- quot )
48     [ ?at drop ] curry ; inline
49
50 : with-assoc ( assoc quot: ( ..a value key assoc -- ..b ) -- quot: ( ..a key value -- ..b ) )
51     curry [ swap ] prepose ; inline
52
53 PRIVATE>
54
55 : assoc-find ( ... assoc quot: ( ... key value -- ... ? ) -- ... key value ? )
56     assoc-operator find swap [ first2-unsafe t ] [ drop f f f ] if ; inline
57
58 : key? ( key assoc -- ? ) at* nip ; inline
59
60 : assoc-each ( ... assoc quot: ( ... key value -- ... ) -- ... )
61     assoc-operator each ; inline
62
63 : assoc>map ( ... assoc quot: ( ... key value -- ... elt ) exemplar -- ... seq )
64     [ assoc-operator ] dip map-as ; inline
65
66 : assoc-map-as ( ... assoc quot: ( ... key value -- ... newkey newvalue ) exemplar -- ... newassoc )
67     [ [ 2array ] compose { } assoc>map ] dip assoc-like ; inline
68
69 : assoc-map ( ... assoc quot: ( ... key value -- ... newkey newvalue ) -- ... newassoc )
70     over assoc-map-as ; inline
71
72 : assoc-filter-as ( ... assoc quot: ( ... key value -- ... ? ) exemplar -- ... subassoc )
73     [ assoc-operator filter ] dip assoc-like ; inline
74
75 : assoc-filter ( ... assoc quot: ( ... key value -- ... ? ) -- ... subassoc )
76     over assoc-filter-as ; inline
77
78 : assoc-reject-as ( ... assoc quot: ( ... key value -- ... ? ) exemplar -- ... subassoc )
79     [ [ not ] compose ] [ assoc-filter-as ] bi* ; inline
80
81 : assoc-reject ( ... assoc quot: ( ... key value -- ... ? ) -- ... subassoc )
82     over assoc-reject-as ; inline
83
84 : assoc-filter! ( ... assoc quot: ( ... key value -- ... ? ) -- ... assoc )
85     [
86         over [ [ [ drop ] 2bi ] dip [ delete-at ] 2curry unless ] 2curry
87         assoc-each
88     ] [ drop ] 2bi ; inline
89
90 : assoc-reject! ( ... assoc quot: ( ... key value -- ... ? ) -- ... assoc )
91     [ not ] compose assoc-filter! ; inline
92
93 : sift-keys ( assoc -- assoc' )
94     [ drop ] assoc-filter ; inline
95
96 : sift-values ( assoc -- assoc' )
97     [ nip ] assoc-filter ; inline
98
99 : harvest-keys ( assoc -- assoc' )
100     [ drop empty? ] assoc-reject ; inline
101
102 : harvest-values ( assoc -- assoc' )
103     [ nip empty? ] assoc-reject ; inline
104
105 : assoc-partition ( ... assoc quot: ( ... key value -- ... ? ) -- ... true-assoc false-assoc )
106     [ assoc-operator partition ] [ drop ] 2bi
107     [ assoc-like ] curry bi@ ; inline
108
109 : assoc-any? ( ... assoc quot: ( ... key value -- ... ? ) -- ... ? )
110     assoc-find 2nip ; inline
111
112 : assoc-all? ( ... assoc quot: ( ... key value -- ... ? ) -- ... ? )
113     [ not ] compose assoc-any? not ; inline
114
115 : at ( key assoc -- value/f )
116     at* drop ; inline
117
118 : ?of ( assoc key -- value/key ? )
119     swap ?at ; inline
120
121 : of ( assoc key -- value/f )
122     swap at ; inline
123
124 M: assoc assoc-clone-like
125     [ dup assoc-size ] dip new-assoc
126     [ [ set-at ] with-assoc assoc-each ] keep ; inline
127
128 M: assoc keys [ drop ] { } assoc>map ;
129
130 M: assoc values [ nip ] { } assoc>map ;
131
132 : delete-at* ( key assoc -- value/f ? )
133     [ at* ] 2keep delete-at ;
134
135 : ?delete-at ( key assoc -- value/key ? )
136     [ ?at ] 2keep delete-at ;
137
138 : rename-at ( newkey key assoc -- )
139     [ delete-at* ] keep [ set-at ] with-assoc [ 2drop ] if ;
140
141 : assoc-empty? ( assoc -- ? )
142     assoc-size 0 = ; inline
143
144 : assoc-stack ( key seq -- value )
145     index-of-last assoc-stack-from ; flushable
146
147 : assoc-subset? ( assoc1 assoc2 -- ? )
148     [ at* [ = ] [ 2drop f ] if ] with-assoc assoc-all? ;
149
150 : assoc= ( assoc1 assoc2 -- ? )
151     2dup [ assoc-size ] bi@ = [ assoc-subset? ] [ 2drop f ] if ;
152
153 : assoc-hashcode ( n assoc -- code )
154     >alist hashcode* ;
155
156 : assoc-intersect ( assoc1 assoc2 -- intersection )
157     swap [ nip key? ] curry assoc-filter ;
158
159 : assoc-union! ( assoc1 assoc2 -- assoc1 )
160     over [ set-at ] with-assoc assoc-each ;
161
162 : assoc-union-as ( assoc1 assoc2 exemplar -- union )
163     [ [ [ assoc-size ] bi@ + ] dip new-assoc ] 2keepd
164     [ assoc-union! ] bi@ ;
165
166 : assoc-union ( assoc1 assoc2 -- union )
167     over assoc-union-as ;
168
169 : assoc-union-all ( seq -- union )
170     H{ } clone [ assoc-union! ] reduce ;
171
172 : assoc-intersect-all ( seq -- assoc )
173     [ f ] [ [ ] [ assoc-intersect ] map-reduce ] if-empty ;
174
175 : assoc-differ ( key -- quot )
176     [ nip key? not ] curry ; inline
177
178 : assoc-diff ( assoc1 assoc2 -- diff )
179     assoc-differ assoc-filter ;
180
181 : assoc-diff! ( assoc1 assoc2 -- assoc1 )
182     assoc-differ assoc-filter! ;
183
184 : substitute ( seq assoc -- newseq )
185     substituter map ;
186
187 : cache ( ... key assoc quot: ( ... key -- ... value ) -- ... value )
188     [ [ at* ] 2keep ] dip
189     [ [ nip call dup ] [ drop ] 3bi set-at ] 3curry
190     [ drop ] prepose
191     unless ; inline
192
193 : 2cache ( ... key1 key2 assoc quot: ( ... key1 key2 -- ... value ) -- ... value )
194     [ 2array ] 2dip [ first2-unsafe ] prepose cache ; inline
195
196 : change-at ( ..a key assoc quot: ( ..a value -- ..b newvalue ) -- ..b )
197     [ [ at ] dip call ] [ drop ] 3bi set-at ; inline
198
199 : ?change-at ( ..a key assoc quot: ( ..a value -- ..b newvalue ) -- ..b )
200     2over [ set-at ] 2curry compose [ at* ] dip [ drop ] if ; inline
201
202 : at+ ( n key assoc -- ) [ 0 or + ] change-at ; inline
203
204 : inc-at ( key assoc -- ) [ 1 ] 2dip at+ ; inline
205
206 : map>assoc ( ... seq quot: ( ... elt -- ... key value ) exemplar -- ... assoc )
207     dup sequence? [
208         [ [ 2array ] compose ] dip map-as
209     ] [
210         [ over assoc-size ] dip new-assoc
211         [ [ swapd set-at ] curry compose each ] keep
212     ] if ; inline
213
214 : map>alist ( ... seq quot: ( ... elt -- ... key value ) -- ... alist )
215     { } map>assoc ; inline
216
217 : extract-keys ( seq assoc -- subassoc )
218     [ [ dupd at ] curry ] keep map>assoc ;
219
220 M: assoc value-at* swap [ = nip ] curry assoc-find nip ;
221
222 : value-at ( value assoc -- key/f ) value-at* drop ;
223
224 : ?value-at ( value assoc -- key/value ? )
225     2dup value-at* [ 2nip t ] [ 2drop f ] if ; inline
226
227 : value? ( value assoc -- ? ) value-at* nip ;
228
229 : push-at ( value key assoc -- )
230     [ ?push ] change-at ;
231
232 : zip-as ( keys values exemplar -- assoc )
233     dup sequence? [
234         [ 2array ] swap 2map-as
235     ] [
236         [ 2dup min-length ] dip new-assoc
237         [ [ set-at ] with-assoc 2each ] keep
238     ] if ; inline
239
240 : zip ( keys values -- alist )
241      { } zip-as ; inline
242
243 : zip-index-as ( values exemplar -- assoc )
244     [ dup length <iota> ] dip zip-as ; inline
245
246 : zip-index ( values -- alist )
247     { } zip-index-as ; inline
248
249 M: assoc unzip
250     dup assoc-empty? [ drop { } { } ] [ >alist flip first2 ] if ;
251
252 : zip-with-as ( ... seq quot: ( ... key -- ... value ) exemplar -- ... assoc )
253     [ [ keep swap ] curry ] dip map>assoc ; inline
254
255 : zip-with ( ... seq quot: ( ... key -- ... value ) -- ... alist )
256     { } zip-with-as ; inline
257
258 : collect-by! ( ... assoc seq quot: ( ... obj -- ... key ) -- ... assoc )
259     [ keep swap ] curry rot [
260         [ push-at ] curry compose each
261     ] keep ; inline
262
263 : collect-by ( ... seq quot: ( ... obj -- ... key ) -- ... assoc )
264     [ H{ } clone ] 2dip collect-by! ; inline
265
266 M: sequence at*
267     search-alist [ second t ] [ f ] if ;
268
269 M: sequence set-at
270     2dup search-alist
271     [ 2nip set-second ]
272     [ drop [ swap 2array ] dip push ] if ;
273
274 M: sequence new-assoc drop <vector> ; inline
275
276 M: sequence clear-assoc delete-all ; inline
277
278 M: sequence delete-at
279     [ nip ] [ search-alist nip ] 2bi
280     [ swap remove-nth! drop ] [ drop ] if* ;
281
282 M: sequence assoc-size length ; inline
283
284 M: sequence assoc-clone-like
285     [ >alist ] dip clone-like ; inline
286
287 M: sequence assoc-like
288     [ >alist ] dip like ; inline
289
290 M: sequence >alist ; inline
291
292 ! Override sequence => assoc instance for f
293 M: f at* 2drop f f ; inline
294
295 M: f assoc-size drop 0 ; inline
296
297 M: f clear-assoc drop ; inline
298
299 M: f assoc-like drop dup assoc-empty? [ drop f ] when ; inline
300
301 INSTANCE: sequence assoc
302
303 TUPLE: enumerated { seq read-only } ;
304
305 C: <enumerated> enumerated
306
307 M: enumerated at*
308     seq>> 2dup bounds-check?
309     [ nth-unsafe t ] [ 2drop f f ] if ; inline
310
311 M: enumerated set-at seq>> set-nth ; inline
312
313 M: enumerated delete-at seq>> remove-nth! drop ; inline
314
315 M: enumerated >alist ; inline
316
317 M: enumerated keys seq>> length <iota> >array ; inline
318
319 M: enumerated values seq>> >array ; inline
320
321 M: enumerated unzip seq>> [ length <iota> ] keep [ >array ] bi@ ;
322
323 M: enumerated assoc-size seq>> length ; inline
324
325 M: enumerated clear-assoc seq>> delete-all ; inline
326
327 INSTANCE: enumerated assoc
328
329 M: enumerated length seq>> length ; inline
330
331 M: enumerated nth-unsafe dupd seq>> nth-unsafe 2array ; inline
332
333 INSTANCE: enumerated immutable-sequence