core/assocs/assocs.factor

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