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