1(* ========================================================================= *)
2(* mlibPatricia: Maps over integers implemented as mlibPatricia trees.               *)
3(* Copyright (c) 2000 Jean-Christophe Filliatre, 2001 Michael Norrish        *)
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
14(* Maps of integers implemented as mlibPatricia trees, following Chris
15   Okasaki and Andrew Gill's paper {\em Fast Mergeable Integer Maps}
16   ({\tt\small http://www.cs.columbia.edu/\~{}cdo/papers.html\#ml98maps}).
17   See the documentation of module [Ptset] which is also based on the
18   same data-structure. *)
19
20structure mlibPatricia :> mlibPatricia =
21struct
22type key = int
23
24datatype 'a t =
25         Empty
26       | Leaf of int * 'a
27       | Branch of int * int * 'a t * 'a t * int
28exception NotFound
29
30val empty = Empty
31
32fun land (p, q) = Word.toIntX(Word.andb(Word.fromInt p, Word.fromInt q))
33infix land
34
35fun zero_bit k m = (k land m) = 0
36
37fun mem k t =
38    case t of
39      Empty => false
40    | Leaf (j,_) => k = j
41    | Branch (_, m, l, r, _) => mem k (if zero_bit k m then l else r)
42
43fun find k t =
44    case t of
45      Empty => raise NotFound
46    | Leaf (j,x) => if k = j then x else raise NotFound
47    | Branch (_, m, l, r, _) => find k (if zero_bit k m then l else r)
48
49fun lowest_bit x = x land ~x
50
51fun branching_bit p0 p1 =
52    lowest_bit (Word.toIntX(Word.xorb(Word.fromInt p0, Word.fromInt p1)))
53
54fun mask p m = p land (m-1)
55
56fun size Empty = 0
57  | size (Leaf _) = 1
58  | size (Branch(_, _, _, _, sz)) = sz
59
60fun join (p0,t0,p1,t1) = let
61  val m = branching_bit p0 p1
62  val sz = size t0 + size t1
63in
64  if zero_bit p0 m then
65    Branch (mask p0 m, m, t0, t1, sz)
66  else
67    Branch (mask p0 m, m, t1, t0, sz)
68end
69
70fun match_prefix k p m = (mask k m) = p
71
72fun addf f k x t = let
73  fun ins t =
74      case t of
75        Empty => Leaf (k,x)
76      | Leaf (j,old) => if j = k then Leaf (k, f old)
77                        else join (k, Leaf (k, x), j, t)
78      | Branch (p,m,t0,t1,sz) =>
79        if match_prefix k p m then
80            if zero_bit k m then Branch (p, m, ins t0, t1, sz+1)
81            else Branch (p, m, t0, ins t1,sz+1)
82        else join (k, Leaf (k,x), p, t)
83in
84    ins t
85end
86
87fun add k x t = addf (fn _ => x) k x t
88
89val branch = fn (_,_,Empty,t) => t
90              | (_,_,t,Empty) => t
91              | (p,m,t0,t1)   => Branch (p,m,t0,t1,size t0 + size t1)
92
93fun remove k t = let
94  fun rmv t =
95      case t of
96        Empty => Empty
97      | Leaf (j,_) => if k = j then Empty else t
98      | Branch (p,m,t0,t1,_) => if match_prefix k p m then
99                                if zero_bit k m then
100                                  branch (p, m, rmv t0, t1)
101                                else
102                                  branch (p, m, t0, rmv t1)
103                              else
104                                t
105in
106  rmv t
107end
108
109fun choose t =
110    case t of
111      Empty      => raise NotFound
112    | Leaf(k, x) => ((k, x), Empty)
113    | Branch(p, m, t0, t1, _) => let
114          val (x, t0) = choose t0
115      in
116          (x, branch(p,m,t0,t1))
117      end
118
119fun getItem t = SOME(choose t) handle NotFound => NONE
120
121
122
123
124fun app f = fn Empty => ()
125              | Leaf (k,x) => f (k, x)
126              | Branch (_,_,t0,t1,_) => (app f t0; app f t1)
127
128fun map f = fn Empty => Empty
129             | Leaf (k,x) => Leaf (k, f x)
130             | Branch (p,m,t0,t1,s) => Branch (p, m, map f t0, map f t1, s)
131
132fun mapi f = fn Empty => Empty
133              | Leaf (k,x) => Leaf (k, f k x)
134              | Branch (p,m,t0,t1,s) => Branch (p, m, mapi f t0, mapi f t1, s)
135
136fun fold f accu s = case s of
137                      Empty => accu
138                    | Leaf (k,x) => f (k, x, accu)
139                    | Branch (_,_,t0,t1,_) => fold f (fold f accu t1) t0
140
141end
142