1(*
2 * PIntMap: Maps over integers implemented as Patricia trees.
3 * Copyright (C) 2000 Jean-Christophe FILLIATRE
4 *
5 * This software is free software; you can redistribute it and/or
6 * modify it under the terms of the GNU Library General Public
7 * License version 2, as published by the Free Software Foundation.
8 *
9 * This software is distributed in the hope that it will be useful,
10 * but WITHOUT ANY WARRANTY; without even the implied warranty of
11 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
12 *
13 * See the GNU Library General Public License version 2 for more details
14 * (enclosed in the file LGPL).
15 *
16 * Translated to SML by Michael Norrish, 2001
17 *)
18
19(*s Maps of integers implemented as Patricia trees, following Chris
20    Okasaki and Andrew Gill's paper {\em Fast Mergeable Integer Maps}
21    ({\tt\small http://www.cs.columbia.edu/\~{}cdo/papers.html\#ml98maps}).
22    See the documentation of module [Ptset] which is also based on the
23    same data-structure. *)
24
25structure PIntMap :> PIntMap =
26struct
27
28type key = int
29
30datatype 'a t =
31         Empty
32       | Leaf of int * 'a
33       | Branch of int * int * 'a t * 'a t * int
34exception NotFound
35
36val empty = Empty
37
38fun land (p, q) = Word.toIntX(Word.andb(Word.fromInt p, Word.fromInt q))
39infix land
40
41fun zero_bit k m = (k land m) = 0
42
43fun mem k t =
44    case t of
45      Empty => false
46    | Leaf (j,_) => k = j
47    | Branch (_, m, l, r, _) => mem k (if zero_bit k m then l else r)
48
49fun find k t =
50    case t of
51      Empty => raise NotFound
52    | Leaf (j,x) => if k = j then x else raise NotFound
53    | Branch (_, m, l, r, _) => find k (if zero_bit k m then l else r)
54
55fun lowest_bit x = x land ~x
56
57fun branching_bit p0 p1 =
58    lowest_bit (Word.toIntX(Word.xorb(Word.fromInt p0, Word.fromInt p1)))
59
60fun mask p m = p land (m-1)
61
62fun size Empty = 0
63  | size (Leaf _) = 1
64  | size (Branch(_, _, _, _, sz)) = sz
65
66fun join (p0,t0,p1,t1) = let
67  val m = branching_bit p0 p1
68  val sz = size t0 + size t1
69in
70  if zero_bit p0 m then
71    Branch (mask p0 m, m, t0, t1, sz)
72  else
73    Branch (mask p0 m, m, t1, t0, sz)
74end
75
76fun match_prefix k p m = (mask k m) = p
77
78fun addf f k x t = let
79  fun ins t =
80      case t of
81        Empty => (Leaf (k,x), x)
82      | Leaf (j,old) => if j = k then let val new = f old
83                                      in (Leaf (k,new), new) end
84                        else (join (k, Leaf (k,x), j, t), x)
85      | Branch (p,m,t0,t1,sz) => if match_prefix k p m then
86                                   if zero_bit k m then let
87                                       val (t0', new) = ins t0
88                                     in
89                                       (Branch (p, m, t0', t1, sz+1), new)
90                                     end
91                                else let
92                                    val (t1', new) = ins t1
93                                  in
94                                    (Branch (p, m, t0, t1',sz+1), new)
95                                  end
96                              else
97                                (join (k, Leaf (k,x), p, t), x)
98in
99  ins t
100end
101
102fun add k x t = #1 (addf (fn _ => x) k x t)
103
104val branch = fn (_,_,Empty,t) => t
105              | (_,_,t,Empty) => t
106              | (p,m,t0,t1)   => Branch (p,m,t0,t1,size t0 + size t1)
107
108fun remove k t = let
109  fun rmv t =
110      case t of
111        Empty => Empty
112      | Leaf (j,_) => if k = j then Empty else t
113      | Branch (p,m,t0,t1,_) => if match_prefix k p m then
114                                if zero_bit k m then
115                                  branch (p, m, rmv t0, t1)
116                                else
117                                  branch (p, m, t0, rmv t1)
118                              else
119                                t
120in
121  rmv t
122end
123
124fun choose t =
125    case t of
126      Empty => raise NotFound
127    | Leaf(k, x) => (Empty, (k, x))
128    | Branch(p, m, t0, t1, _) => let
129        val (t0', x) = choose t0
130      in
131        (branch(p,m,t0',t1), x)
132      end
133
134fun app f = fn Empty => ()
135              | Leaf (k,x) => f (k, x)
136              | Branch (_,_,t0,t1,_) => (app f t0; app f t1)
137
138fun map f = fn Empty => Empty
139             | Leaf (k,x) => Leaf (k, f x)
140             | Branch (p,m,t0,t1,s) => Branch (p, m, map f t0, map f t1, s)
141
142fun mapi f = fn Empty => Empty
143              | Leaf (k,x) => Leaf (k, f k x)
144              | Branch (p,m,t0,t1,s) => Branch (p, m, mapi f t0, mapi f t1, s)
145
146fun fold f accu s = case s of
147                      Empty => accu
148                    | Leaf (k,x) => f (k, x, accu)
149                    | Branch (_,_,t0,t1,_) => fold f (fold f accu t1) t0
150
151end
152