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