1 ! Based on Clojure's PersistentVector by Rich Hickey.
3 USING: math accessors kernel sequences.private sequences arrays
4 combinators combinators.short-circuit parser prettyprint.custom
10 TUPLE: node { children array } { level fixnum } ;
14 ERROR: empty-error pvec ;
16 TUPLE: persistent-vector
18 { root node initial: T{ node f { } 1 } }
19 { tail node initial: T{ node f { } 0 } } ;
21 M: persistent-vector length count>> ;
25 CONSTANT: node-size 32
27 : node-mask ( m -- n ) node-size mod ; inline
29 : node-shift ( m n -- x ) -5 * shift ; inline
31 : node-nth ( i node -- obj )
32 [ node-mask ] [ children>> ] bi* nth ;
34 : body-nth ( i node -- i node' )
36 dupd [ level>> node-shift ] keep node-nth
39 : tail-offset ( pvec -- n )
40 [ count>> ] [ tail>> children>> length ] bi - ;
42 M: persistent-vector nth-unsafe
44 [ tail>> ] [ root>> body-nth ] if
47 : node-add ( val node -- node' )
48 clone [ ppush ] change-children ;
50 : ppush-tail ( val pvec -- pvec' )
51 [ node-add ] change-tail ;
54 children>> length node-size = ;
56 : 1node ( val level -- node )
57 [ 1array ] dip node boa ;
59 : 2node ( first second -- node )
60 [ 2array ] [ drop level>> 1 + ] 2bi node boa ;
62 : new-child ( new-child node -- node' expansion/f )
63 dup full? [ [ level>> 1node ] keep swap ] [ node-add f ] if ;
65 : new-last ( val seq -- seq' )
66 [ length 1 - ] keep new-nth ;
68 : node-set-last ( child node -- node' )
69 clone [ new-last ] change-children ;
71 : (ppush-new-tail) ( tail node -- node' expansion/f )
75 [ nip ] 2keep children>> last (ppush-new-tail)
76 [ swap new-child ] [ swap node-set-last f ] ?if
79 : do-expansion ( pvec root expansion/f -- pvec )
80 [ 2node ] when* >>root ;
82 : ppush-new-tail ( val pvec -- pvec' )
83 [ ] [ tail>> ] [ root>> ] tri
84 (ppush-new-tail) do-expansion
87 M: persistent-vector ppush ( val pvec -- pvec' )
90 [ ppush-new-tail ] [ ppush-tail ] if
91 [ 1 + ] change-count ;
93 : node-set-nth ( val i node -- node' )
94 clone [ new-nth ] change-children ;
96 : node-change-nth ( i node quot -- node' )
98 [ clone ] dip [ change-nth ] 2keep drop
99 ] curry change-children ; inline
101 : (new-nth) ( val i node -- node' )
103 [ node-mask ] dip node-set-nth
105 [ dupd level>> node-shift node-mask ] keep
106 [ (new-nth) ] node-change-nth
109 M: persistent-vector new-nth ( obj i pvec -- pvec' )
110 2dup count>> = [ nip ppush ] [
112 2dup tail-offset >= [
114 [ node-set-nth ] change-tail
116 [ (new-nth) ] change-root
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.
124 : ppop-tail ( pvec -- pvec' )
125 [ clone [ ppop ] change-children ] change-tail ;
127 : (ppop-contraction) ( node -- node' tail' )
128 clone [ unclip-last swap ] change-children swap ;
130 : ppop-contraction ( node -- node' tail' )
131 dup children>> length 1 =
132 [ children>> last f swap ]
133 [ (ppop-contraction) ]
136 : (ppop-new-tail) ( root -- root' tail' )
138 dup children>> last (ppop-new-tail) [
140 [ swap node-set-last ]
141 [ drop ppop-contraction drop ]
148 : trivial? ( node -- ? )
149 { [ level>> 1 > ] [ children>> length 1 = ] } 1&& ;
151 : ppop-new-tail ( pvec -- pvec' )
152 dup root>> (ppop-new-tail) [
154 { [ dup not ] [ drop T{ node f { } 1 } ] }
155 { [ dup trivial? ] [ children>> first ] }
158 ] dip [ >>root ] [ >>tail ] bi* ;
162 M: persistent-vector ppop ( pvec -- pvec' )
164 { 0 [ throw-empty-error ] }
165 { 1 [ drop T{ persistent-vector } ] }
169 dup tail>> children>> length 1 >
170 [ ppop-tail ] [ ppop-new-tail ] if
175 M: persistent-vector like
176 drop T{ persistent-vector } [ swap ppush ] reduce ;
178 M: persistent-vector equal?
179 over persistent-vector? [ sequence= ] [ 2drop f ] if ;
181 : >persistent-vector ( seq -- pvec )
182 T{ persistent-vector } like ;
184 SYNTAX: PV{ \ } [ >persistent-vector ] parse-literal ;
186 M: persistent-vector pprint-delims drop \ PV{ \ } ;
187 M: persistent-vector >pprint-sequence ;
188 M: persistent-vector pprint* pprint-object ;
190 INSTANCE: persistent-vector immutable-sequence