basis/persistent/vectors/vectors.factor

   1 ! Based on Clojure's PersistentVector by Rich Hickey.
   2
   3 USING: math accessors kernel sequences.private sequences arrays
   4 combinators combinators.short-circuit parser prettyprint.custom
   5 persistent.sequences ;
   6 IN: persistent.vectors
   7
   8 <PRIVATE
   9
  10 TUPLE: node { children array } { level fixnum } ;
  11
  12 PRIVATE>
  13
  14 ERROR: empty-error pvec ;
  15
  16 TUPLE: persistent-vector
  17 { count fixnum }
  18 { root node initial: T{ node f { } 1 } }
  19 { tail node initial: T{ node f { } 0 } } ;
  20
  21 M: persistent-vector length count>> ;
  22
  23 CONSTANT: node-size 32
  24
  25 : node-mask ( m -- n ) node-size mod ; inline
  26
  27 : node-shift ( m n -- x ) -5 * shift ; inline
  28
  29 : node-nth ( i node -- obj )
  30     [ node-mask ] [ children>> ] bi* nth ;
  31
  32 : body-nth ( i node -- i node' )
  33     dup level>> [
  34         dupd [ level>> node-shift ] keep node-nth
  35     ] times ;
  36
  37 : tail-offset ( pvec -- n )
  38     [ count>> ] [ tail>> children>> length ] bi - ;
  39
  40 M: persistent-vector nth-unsafe
  41     2dup tail-offset >=
  42     [ tail>> ] [ root>> body-nth ] if
  43     node-nth ;
  44
  45 : node-add ( val node -- node' )
  46     clone [ ppush ] change-children ;
  47
  48 : ppush-tail ( val pvec -- pvec' )
  49     [ node-add ] change-tail ;
  50
  51 : full? ( node -- ? )
  52     children>> length node-size = ;
  53
  54 : 1node ( val level -- node )
  55     [ 1array ] dip node boa ;
  56
  57 : 2node ( first second -- node )
  58     [ 2array ] [ drop level>> 1 + ] 2bi node boa ;
  59
  60 : new-child ( new-child node -- node' expansion/f )
  61     dup full? [ tuck level>> 1node ] [ node-add f ] if ;
  62
  63 : new-last ( val seq -- seq' )
  64     [ length 1 - ] keep new-nth ;
  65
  66 : node-set-last ( child node -- node' )
  67     clone [ new-last ] change-children ;
  68
  69 : (ppush-new-tail) ( tail node -- node' expansion/f )
  70     dup level>> 1 = [
  71         new-child
  72     ] [
  73         tuck children>> last (ppush-new-tail)
  74         [ swap new-child ] [ swap node-set-last f ] ?if
  75     ] if ;
  76
  77 : do-expansion ( pvec root expansion/f -- pvec )
  78     [ 2node ] when* >>root ;
  79
  80 : ppush-new-tail ( val pvec -- pvec' )
  81     [ ] [ tail>> ] [ root>> ] tri
  82     (ppush-new-tail) do-expansion
  83     swap 0 1node >>tail ;
  84
  85 M: persistent-vector ppush ( val pvec -- pvec' )
  86     clone
  87     dup tail>> full?
  88     [ ppush-new-tail ] [ ppush-tail ] if
  89     [ 1 + ] change-count ;
  90
  91 : node-set-nth ( val i node -- node' )
  92     clone [ new-nth ] change-children ;
  93
  94 : node-change-nth ( i node quot -- node' )
  95     [ clone ] dip [
  96         [ clone ] dip [ change-nth ] 2keep drop
  97     ] curry change-children ; inline
  98
  99 : (new-nth) ( val i node -- node' )
 100     dup level>> 0 = [
 101         [ node-mask ] dip node-set-nth
 102     ] [
 103         [ dupd level>> node-shift node-mask ] keep
 104         [ (new-nth) ] node-change-nth
 105     ] if ;
 106
 107 M: persistent-vector new-nth ( obj i pvec -- pvec' )
 108     2dup count>> = [ nip ppush ] [
 109         clone
 110         2dup tail-offset >= [
 111             [ node-mask ] dip
 112             [ node-set-nth ] change-tail
 113         ] [
 114             [ (new-nth) ] change-root
 115         ] if
 116     ] if ;
 117
 118 ! The pop code is really convoluted. I don't understand Rich Hickey's
 119 ! original code. It uses a 'Box' out parameter which is passed around
 120 ! inside a recursive function, and gets mutated along the way to boot.
 121 ! Super-confusing.
 122 : ppop-tail ( pvec -- pvec' )
 123     [ clone [ ppop ] change-children ] change-tail ;
 124
 125 : (ppop-contraction) ( node -- node' tail' )
 126     clone [ unclip-last swap ] change-children swap ;
 127
 128 : ppop-contraction ( node -- node' tail' )
 129     dup children>> length 1 =
 130     [ children>> last f swap ]
 131     [ (ppop-contraction) ]
 132     if ;
 133
 134 : (ppop-new-tail) ( root -- root' tail' )
 135     dup level>> 1 > [
 136         dup children>> last (ppop-new-tail) [
 137             dup
 138             [ swap node-set-last ]
 139             [ drop ppop-contraction drop ]
 140             if
 141         ] dip
 142     ] [
 143         ppop-contraction
 144     ] if ;
 145
 146 : trivial? ( node -- ? )
 147     { [ level>> 1 > ] [ children>> length 1 = ] } 1&& ;
 148
 149 : ppop-new-tail ( pvec -- pvec' )
 150     dup root>> (ppop-new-tail) [
 151         {
 152             { [ dup not ] [ drop T{ node f { } 1 } ] }
 153             { [ dup trivial? ] [ children>> first ] }
 154             [ ]
 155         } cond
 156     ] dip [ >>root ] [ >>tail ] bi* ;
 157
 158 PRIVATE>
 159
 160 M: persistent-vector ppop ( pvec -- pvec' )
 161     dup count>> {
 162         { 0 [ empty-error ] }
 163         { 1 [ drop T{ persistent-vector } ] }
 164         [
 165             [
 166                 clone
 167                 dup tail>> children>> length 1 >
 168                 [ ppop-tail ] [ ppop-new-tail ] if
 169             ] dip 1 - >>count
 170         ]
 171     } case ;
 172
 173 M: persistent-vector like
 174     drop T{ persistent-vector } [ swap ppush ] reduce ;
 175
 176 M: persistent-vector equal?
 177     over persistent-vector? [ sequence= ] [ 2drop f ] if ;
 178
 179 : >persistent-vector ( seq -- pvec )
 180     T{ persistent-vector } like ;
 181
 182 SYNTAX: PV{ \ } [ >persistent-vector ] parse-literal ;
 183
 184 M: persistent-vector pprint-delims drop \ PV{ \ } ;
 185 M: persistent-vector >pprint-sequence ;
 186 M: persistent-vector pprint* pprint-object ;
 187
 188 INSTANCE: persistent-vector immutable-sequence