1;;; fxmap.ss
2;;; Copyright 1984-2017 Cisco Systems, Inc.
3;;;
4;;; Licensed under the Apache License, Version 2.0 (the "License");
5;;; you may not use this file except in compliance with the License.
6;;; You may obtain a copy of the License at
7;;;
8;;; http://www.apache.org/licenses/LICENSE-2.0
9;;;
10;;; Unless required by applicable law or agreed to in writing, software
11;;; distributed under the License is distributed on an "AS IS" BASIS,
12;;; WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13;;; See the License for the specific language governing permissions and
14;;; limitations under the License.
15
16;; Based on Okasaki and Gill's "Fast Mergeable Integer Maps" (1998).
17
18(module fxmap
19 (fxmap?
20  empty-fxmap
21  fxmap-empty?
22  fxmap-count
23  fxmap-ref
24  fxmap-set
25  fxmap-remove
26  fxmap-remove/base
27  fxmap-reset/base
28  fxmap-advance/base
29  fxmap-for-each
30  fxmap-for-each/diff
31  fxmap-changes
32
33  ;; internals
34  ; $branch? make-$branch $branch-prefix $branch-mask $branch-left $branch-right
35  ; $leaf? make-$leaf $leaf-key $leaf-val
36
37  ;; We treat $empty as a singleton, so don't use these functions.
38  ; $empty? make-$empty
39  )
40
41 ;; record types
42
43 (define-record-type $branch
44   (fields prefix mask left right count changes)
45   (nongenerative #{$branch pfv8jpsat5jrk6vq7vclc3ntg-1})
46   (sealed #t))
47
48 (define-record-type $leaf
49   (fields key val changes)
50   (nongenerative #{$leaf pfv8jq2dzw50ox4f6vqm1ff5v-1})
51   (sealed #t))
52
53 (define-record-type $empty
54   (nongenerative #{$empty pfwk1nal7cs5dornqtzvda91m-0})
55   (sealed #t))
56
57 (define-syntax let-branch
58   (syntax-rules ()
59     [(_ ([(p m l r) d] ...) exp ...)
60      (let ([p ($branch-prefix d)] ...
61            [m ($branch-mask d)] ...
62            [l ($branch-left d)] ...
63            [r ($branch-right d)] ...)
64        exp ...)]))
65
66 ;; constants & empty
67
68 (define empty-fxmap (make-$empty))
69
70 (define (fxmap-empty? x) (eq? empty-fxmap x))
71
72 ;; predicate
73
74 (define (fxmap? x)
75   (or ($branch? x)
76       ($leaf? x)
77       (eq? empty-fxmap x)))
78
79 ;; count & changes
80
81 (define (fxmap-count d)
82   (cond
83     [($branch? d)
84      ($branch-count d)]
85     [($leaf? d) 1]
86     [else 0]))
87
88 (define (fxmap-changes d)
89   (cond
90     [($branch? d)
91      ($branch-changes d)]
92     [($leaf? d)
93      ($leaf-changes d)]
94     [else 0]))
95
96 ;; ref
97
98 (define (fxmap-ref/leaf d key)
99   (cond
100     [($branch? d)
101      (let-branch ([(p m l r) d])
102        (cond
103          [(fx<= key p)
104           (fxmap-ref/leaf l key)]
105          [else
106           (fxmap-ref/leaf r key)]))]
107
108     [($leaf? d)
109      (if (fx= key ($leaf-key d))
110          d
111          #f)]
112
113     [else
114      #f]))
115
116 (define (fxmap-ref d key default)
117   (let ([d (fxmap-ref/leaf d key)])
118     (if d
119         ($leaf-val d)
120         default)))
121
122 (define (fxmap-ref/changes d key)
123   (let ([d (fxmap-ref/leaf d key)])
124     (if d
125         ($leaf-changes d)
126         0)))
127
128 ;; set
129
130 (define (fxmap-set/changes d key val changes)
131   (cond
132    [($branch? d)
133     (let-branch ([(p m l r) d])
134       (cond
135        [(nomatch? key p m)
136         (join key (make-$leaf key val (or changes 1)) p d)]
137        [(fx<= key p)
138         (br p m (fxmap-set/changes l key val changes) r)]
139        [else
140         (br p m l (fxmap-set/changes r key val changes))]))]
141
142    [($leaf? d)
143     (let ([k ($leaf-key d)])
144       (if (fx= key k)
145           (make-$leaf key val (or changes (fx+ ($leaf-changes d) 1)))
146           (join key (make-$leaf key val (or changes 1)) k d)))]
147
148    [else
149     (make-$leaf key val (or changes 1))]))
150
151 (define (fxmap-set d key val)
152   (fxmap-set/changes d key val #f))
153
154 ;; remove
155
156 (define (fxmap-remove d key)
157   (cond
158    [($branch? d)
159     (let-branch ([(p m l r) d])
160       (cond
161        [(nomatch? key p m) d]
162        [(fx<= key p)       (br* p m (fxmap-remove l key) r)]
163        [else               (br* p m l (fxmap-remove r key))]))]
164
165    [($leaf? d)
166     (if (fx= key ($leaf-key d))
167         empty-fxmap
168         d)]
169
170    [else
171     empty-fxmap]))
172
173 (define (fxmap-remove/base d key base)
174   ; Remove key from d, but try to reuse the branches from base when possible
175   ; instead of creating new ones.
176   ; TODO: This assumes that all the keys in base are in d too.
177   ; Perhaps this restriction can be removed.
178   (cond
179     [($branch? base)
180      (cond
181        [($branch? d)
182         (let-branch ([(p0 m0 l0 r0) base]
183                      [(p m l r) d])
184           (let ([sub-base (cond
185                             [(fx< m0 m) base]
186                             [(fx<= key p0) l0]
187                             [else r0])])
188             (cond
189               [(nomatch? key p m)
190                d]
191               [(fx<= key p)
192                (br*/base p m (fxmap-remove/base l key sub-base) r base)]
193               [else
194                (br*/base p m l (fxmap-remove/base r key sub-base) base)])))]
195
196        [($leaf? d)
197         (if (fx= key ($leaf-key d))
198             empty-fxmap
199             d)]
200
201        [else
202         empty-fxmap])]
203    [else
204     (fxmap-remove d key)]))
205
206 ;; reset and advance
207
208 (define (fxmap-reset/base d key base)
209   ; Reset key in d to the value it has in base, but try to reuse the branches
210   ; from base when possible instead of creating new ones.
211   ; TODO: This assumes that all the keys in base are in d too.
212   ; Perhaps this restriction can be removed.
213   (cond
214     [($branch? d)
215      (let-branch ([(p m l r) d])
216        (let ([sub-base (cond
217                          [($branch? base)
218                           (let-branch ([(p0 m0 l0 r0) base])
219                             (cond
220                               [(fx< m0 m) base]
221                               [(fx<= key p0) l0]
222                               [else r0]))]
223                          [else base])])
224           (cond
225             [(nomatch? key p m)
226              d]
227             [(fx<= key p)
228              (br*/base p m (fxmap-reset/base l key sub-base) r base)]
229             [else
230              (br*/base p m l (fxmap-reset/base r key sub-base) base)])))]
231
232      [(and ($leaf? d)
233            (fx= key ($leaf-key d))
234            ($leaf? base)
235            (fx= key ($leaf-key base)))
236       base]
237
238      [else
239       (error 'fxmap-reset/base "")]))
240
241 (define (fxmap-advance/base d key base)
242   (let ([changes (fx+ (fxmap-ref/changes base key) 1)]
243         [l (fxmap-ref/leaf d key)])
244     (if l
245         (if (fx= changes ($leaf-changes l))
246             d
247            (fxmap-set/changes d key ($leaf-val l) changes))
248         (error 'fxmap-advance/base ""))))
249
250 ;; set and remove utilities
251
252 (define-syntax define-syntax-rule
253  (syntax-rules ()
254    [(_ (name arg ...) e ...)
255     (define-syntax name
256       (syntax-rules ()
257         [(_ arg ...) e ...]))]))
258
259 (define (br p m l r)
260   (make-$branch p m l r
261                 (fx+ (fxmap-count l) (fxmap-count r))
262                 (fx+ (fxmap-changes l) (fxmap-changes r))))
263
264 (define (br* p m l r)
265   (cond [(eq? empty-fxmap r) l]
266         [(eq? empty-fxmap l) r]
267         [else (br p m l r)]))
268
269 (define (br*/base p m l r base)
270   (cond [(eq? empty-fxmap r) l]
271         [(eq? empty-fxmap l) r]
272         [(and ($branch? base)
273               (eq? l ($branch-left base))
274               (eq? r ($branch-right base)))
275          base]
276         [else (br p m l r)]))
277
278 (define (join p0 d0 p1 d1)
279   (let ([m (branching-bit p0 p1)])
280     (if (fx<= p0 p1)
281         (br (mask p0 m) m d0 d1)
282         (br (mask p0 m) m d1 d0))))
283
284 (define (join* p1 d1 p2 d2)
285   (cond
286     [(eq? empty-fxmap d1) d2]
287     [(eq? empty-fxmap d2) d1]
288     [else (join p1 d1 p2 d2)]))
289
290 (define (branching-bit p m)
291   (highest-set-bit (fxxor p m)))
292
293 (define-syntax-rule (mask h m)
294   (fxand (fxior h (fx1- m)) (fxnot m)))
295
296 (define highest-set-bit
297   (if (fx= (fixnum-width) 61)
298       (lambda (x1)
299         (let* ([x2 (fxior x1 (fxsrl x1 1))]
300                [x3 (fxior x2 (fxsrl x2 2))]
301                [x4 (fxior x3 (fxsrl x3 4))]
302                [x5 (fxior x4 (fxsrl x4 8))]
303                [x6 (fxior x5 (fxsrl x5 16))]
304                [x7 (fxior x6 (fxsrl x6 32))])
305           (fxxor x7 (fxsrl x7 1))))
306       (lambda (x1)
307         (let* ([x2 (fxior x1 (fxsrl x1 1))]
308                [x3 (fxior x2 (fxsrl x2 2))]
309                [x4 (fxior x3 (fxsrl x3 4))]
310                [x5 (fxior x4 (fxsrl x4 8))]
311                [x6 (fxior x5 (fxsrl x5 16))])
312           (fxxor x6 (fxsrl x6 1))))))
313
314
315 (define-syntax-rule (nomatch? h p m)
316   (not (fx= (mask h m) p)))
317
318 ;; merge
319
320 (define (fxmap-merge bin f id g1 g2 d1 d2)
321   (define-syntax go
322     (syntax-rules ()
323       [(_ d1 d2) (fxmap-merge bin f id g1 g2 d1 d2)]))
324
325   (cond
326    [(eq? d1 d2) (id d1)]
327
328    [($branch? d1)
329     (cond
330      [($branch? d2)
331       (let-branch ([(p1 m1 l1 r1) d1]
332                    [(p2 m2 l2 r2) d2])
333        (cond
334         [(fx> m1 m2) (cond
335                       [(nomatch? p2 p1 m1) (join* p1 (g1 d1) p2 (g2 d2))]
336                       [(fx<= p2 p1)        (bin p1 m1 (go l1 d2) (g1 r1))]
337                       [else                (bin p1 m1 (g1 l1) (go r1 d2))])]
338         [(fx> m2 m1) (cond
339                       [(nomatch? p1 p2 m2) (join* p1 (g1 d1) p2 (g2 d2))]
340                       [(fx<= p1 p2)        (bin p2 m2 (go d1 l2) (g2 r2))]
341                       [else                (bin p2 m2 (g2 l2) (go d1 r2))])]
342         [(fx= p1 p2) (bin p1 m1 (go l1 l2) (go r1 r2))]
343         [else        (join* p1 (g1 d1) p2 (g2 d2))]))]
344
345      [($leaf? d2)
346       (let ([k2 ($leaf-key d2)])
347         (let merge0 ([d1 d1])
348           (cond
349            [(eq? d1 d2)
350             (id d1)]
351
352            [($branch? d1)
353             (let-branch ([(p1 m1 l1 r1) d1])
354              (cond [(nomatch? k2 p1 m1) (join* p1 (g1 d1) k2 (g2 d2))]
355                    [(fx<= k2 p1)        (bin p1 m1 (merge0 l1) (g1 r1))]
356                    [else                (bin p1 m1 (g1 l1) (merge0 r1))]))]
357
358            [($leaf? d1)
359             (let ([k1 ($leaf-key d1)])
360               (cond [(fx= k1 k2) (f d1 d2)]
361                     [else        (join* k1 (g1 d1) k2 (g2 d2))]))]
362
363            [else ; (eq? empty-fxmap d1)
364             (g2 d2)])))]
365
366      [else ; (eq? empty-fxmap d2)
367       (g1 d1)])]
368
369    [($leaf? d1)
370     (let ([k1 ($leaf-key d1)])
371       (let merge0 ([d2 d2])
372         (cond
373          [(eq? d1 d2)
374           (id d1)]
375
376          [($branch? d2)
377           (let-branch ([(p2 m2 l2 r2) d2])
378            (cond [(nomatch? k1 p2 m2) (join* k1 (g1 d1) p2 (g2 d2))]
379                  [(fx<= k1 p2)        (bin p2 m2 (merge0 l2) (g2 r2))]
380                  [else                (bin p2 m2 (g2 l2) (merge0 r2))]))]
381
382          [($leaf? d2)
383           (let ([k2 ($leaf-key d2)])
384             (cond [(fx= k1 k2) (f d1 d2)]
385                   [else        (join* k1 (g1 d1) k2 (g2 d2))]))]
386
387          [else ; (eq? empty-fxmap d2)
388           (g1 d1)])))]
389
390    [else ; (eq? empty-fxmap d1)
391     (g2 d2)]))
392
393 ;; merge*
394 ; like merge, but the result is (void)
395
396 (define (fxmap-merge* f id g1 g2 d1 d2)
397   (define (merge* f id g1 g2 d1 d2)
398     (define-syntax go
399       (syntax-rules ()
400         [(_ d1 d2) (merge* f id g1 g2 d1 d2)]))
401
402     (cond
403      [(eq? d1 d2) (id d1)]
404
405      [($branch? d1)
406       (cond
407        [($branch? d2)
408         (let-branch ([(p1 m1 l1 r1) d1]
409                      [(p2 m2 l2 r2) d2])
410          (cond
411           [(fx> m1 m2) (cond
412                         [(nomatch? p2 p1 m1) (g1 d1) (g2 d2)]
413                         [(fx<= p2 p1)        (go l1 d2) (g1 r1)]
414                         [else                (g1 l1) (go r1 d2)])]
415           [(fx> m2 m1) (cond
416                         [(nomatch? p1 p2 m2) (g1 d1) (g2 d2)]
417                         [(fx<= p1 p2)        (go d1 l2) (g2 r2)]
418                         [else                (g2 l2) (go d1 r2)])]
419           [(fx= p1 p2) (go l1 l2) (go r1 r2)]
420           [else        (g1 d1) (g2 d2)]))]
421
422        [else ; ($leaf? d2)
423         (let ([k2 ($leaf-key d2)])
424           (let merge*0 ([d1 d1])
425             (cond
426              [(eq? d1 d2)
427               (id d1)]
428
429              [($branch? d1)
430               (let-branch ([(p1 m1 l1 r1) d1])
431                (cond [(nomatch? k2 p1 m1) (g1 d1) (g2 d2)]
432                      [(fx<= k2 p1)        (merge*0 l1) (g1 r1)]
433                      [else                (g1 l1) (merge*0 r1)]))]
434
435              [else ; ($leaf? d1)
436               (let ([k1 ($leaf-key d1)])
437                 (cond [(fx= k1 k2) (f d1 d2)]
438                       [else        (g1 d1) (g2 d2)]))])))])]
439
440      [($leaf? d1)
441       (let ([k1 ($leaf-key d1)])
442         (let merge*0 ([d2 d2])
443           (cond
444            [(eq? d1 d2)
445             (id d1)]
446
447            [($branch? d2)
448             (let-branch ([(p2 m2 l2 r2) d2])
449              (cond [(nomatch? k1 p2 m2) (g1 d1) (g2 d2)]
450                    [(fx<= k1 p2)        (merge*0 l2) (g2 r2)]
451                    [else                (g2 l2) (merge*0 r2)]))]
452
453            [else ; ($leaf? d2)
454             (let ([k2 ($leaf-key d2)])
455               (cond [(fx= k1 k2) (f d1 d2)]
456                     [else        (g1 d1) (g2 d2)]))])))]))
457
458   (cond
459    [(eq? d1 d2)
460     (id d1)]
461    [(eq? empty-fxmap d1)
462     (g2 d2)]
463    [(eq? empty-fxmap d2)
464     (g1 d1)]
465    [else
466     (merge* f id g1 g2 d1 d2)])
467   (void))
468
469 ;; for-each
470
471 (define (fxmap-for-each g1 d1)
472   (cond
473     [($branch? d1)
474      (fxmap-for-each g1 ($branch-left d1))
475      (fxmap-for-each g1 ($branch-right d1))]
476     [($leaf? d1)
477      (g1 ($leaf-key d1) ($leaf-val d1))]
478     [else ; (eq? empty-fxmap d1)
479      (void)])
480   (void))
481
482 (define (fxmap-for-each/diff f g1 g2 d1 d2)
483   (fxmap-merge* (lambda (x y) (f ($leaf-key x) ($leaf-val x) ($leaf-val y)))
484                 (lambda (x) (void))
485                 (lambda (x) (fxmap-for-each g1 x))
486                 (lambda (x) (fxmap-for-each g2 x))
487                 d1
488                 d2)
489   (void))
490)
491