on bf4e535 help sign in
Idris-dev/libs/contrib/Data/SortedMap.idr
248 lines 7.9 kB history raw
1module Data.SortedMap
3-- TODO: write merge and split
5private
6data Tree : Nat -> (k : Type) -> Type -> Ord k -> Type where
7 Leaf : k -> v -> Tree Z k v o
8 Branch2 : Tree n k v o -> k -> Tree n k v o -> Tree (S n) k v o
9 Branch3 : Tree n k v o -> k -> Tree n k v o -> k -> Tree n k v o -> Tree (S n) k v o
11branch4 :
12 Tree n k v o -> k ->
13 Tree n k v o -> k ->
14 Tree n k v o -> k ->
15 Tree n k v o ->
16 Tree (S (S n)) k v o
17branch4 a b c d e f g =
18 Branch2 (Branch2 a b c) d (Branch2 e f g)
20branch5 :
21 Tree n k v o -> k ->
22 Tree n k v o -> k ->
23 Tree n k v o -> k ->
24 Tree n k v o -> k ->
25 Tree n k v o ->
26 Tree (S (S n)) k v o
27branch5 a b c d e f g h i =
28 Branch2 (Branch2 a b c) d (Branch3 e f g h i)
30branch6 :
31 Tree n k v o -> k ->
32 Tree n k v o -> k ->
33 Tree n k v o -> k ->
34 Tree n k v o -> k ->
35 Tree n k v o -> k ->
36 Tree n k v o ->
37 Tree (S (S n)) k v o
38branch6 a b c d e f g h i j k =
39 Branch3 (Branch2 a b c) d (Branch2 e f g) h (Branch2 i j k)
41branch7 :
42 Tree n k v o -> k ->
43 Tree n k v o -> k ->
44 Tree n k v o -> k ->
45 Tree n k v o -> k ->
46 Tree n k v o -> k ->
47 Tree n k v o -> k ->
48 Tree n k v o ->
49 Tree (S (S n)) k v o
50branch7 a b c d e f g h i j k l m =
51 Branch3 (Branch3 a b c d e) f (Branch2 g h i) j (Branch2 k l m)
53merge1 : Tree n k v o -> k -> Tree (S n) k v o -> k -> Tree (S n) k v o -> Tree (S (S n)) k v o
54merge1 a b (Branch2 c d e) f (Branch2 g h i) = branch5 a b c d e f g h i
55merge1 a b (Branch2 c d e) f (Branch3 g h i j k) = branch6 a b c d e f g h i j k
56merge1 a b (Branch3 c d e f g) h (Branch2 i j k) = branch6 a b c d e f g h i j k
57merge1 a b (Branch3 c d e f g) h (Branch3 i j k l m) = branch7 a b c d e f g h i j k l m
59merge2 : Tree (S n) k v o -> k -> Tree n k v o -> k -> Tree (S n) k v o -> Tree (S (S n)) k v o
60merge2 (Branch2 a b c) d e f (Branch2 g h i) = branch5 a b c d e f g h i
61merge2 (Branch2 a b c) d e f (Branch3 g h i j k) = branch6 a b c d e f g h i j k
62merge2 (Branch3 a b c d e) f g h (Branch2 i j k) = branch6 a b c d e f g h i j k
63merge2 (Branch3 a b c d e) f g h (Branch3 i j k l m) = branch7 a b c d e f g h i j k l m
65merge3 : Tree (S n) k v o -> k -> Tree (S n) k v o -> k -> Tree n k v o -> Tree (S (S n)) k v o
66merge3 (Branch2 a b c) d (Branch2 e f g) h i = branch5 a b c d e f g h i
67merge3 (Branch2 a b c) d (Branch3 e f g h i) j k = branch6 a b c d e f g h i j k
68merge3 (Branch3 a b c d e) f (Branch2 g h i) j k = branch6 a b c d e f g h i j k
69merge3 (Branch3 a b c d e) f (Branch3 g h i j k) l m = branch7 a b c d e f g h i j k l m
71treeLookup : k -> Tree n k v o -> Maybe v
72treeLookup k (Leaf k' v) =
73 if k == k' then
74 Just v
75 else
76 Nothing
77treeLookup k (Branch2 t1 k' t2) =
78 if k <= k' then
79 treeLookup k t1
80 else
81 treeLookup k t2
82treeLookup k (Branch3 t1 k1 t2 k2 t3) =
83 if k <= k1 then
84 treeLookup k t1
85 else if k <= k2 then
86 treeLookup k t2
87 else
88 treeLookup k t3
90treeInsert' : k -> v -> Tree n k v o -> Either (Tree n k v o) (Tree n k v o, k, Tree n k v o)
91treeInsert' k v (Leaf k' v') =
92 case compare k k' of
93 LT => Right (Leaf k v, k, Leaf k' v')
94 EQ => Left (Leaf k v)
95 GT => Right (Leaf k' v', k', Leaf k v)
96treeInsert' k v (Branch2 t1 k' t2) =
97 if k <= k' then
98 case treeInsert' k v t1 of
99 Left t1' => Left (Branch2 t1' k' t2)
100 Right (a, b, c) => Left (Branch3 a b c k' t2)
101 else
102 case treeInsert' k v t2 of
103 Left t2' => Left (Branch2 t1 k' t2')
104 Right (a, b, c) => Left (Branch3 t1 k' a b c)
105treeInsert' k v (Branch3 t1 k1 t2 k2 t3) =
106 if k <= k1 then
107 case treeInsert' k v t1 of
108 Left t1' => Left (Branch3 t1' k1 t2 k2 t3)
109 Right (a, b, c) => Right (Branch2 a b c, k1, Branch2 t2 k2 t3)
110 else
111 if k <= k2 then
112 case treeInsert' k v t2 of
113 Left t2' => Left (Branch3 t1 k1 t2' k2 t3)
114 Right (a, b, c) => Right (Branch2 t1 k1 a, b, Branch2 c k2 t3)
115 else
116 case treeInsert' k v t3 of
117 Left t3' => Left (Branch3 t1 k1 t2 k2 t3')
118 Right (a, b, c) => Right (Branch2 t1 k1 t2, k2, Branch2 a b c)
120treeInsert : k -> v -> Tree n k v o -> Either (Tree n k v o) (Tree (S n) k v o)
121treeInsert k v t =
122 case treeInsert' k v t of
123 Left t' => Left t'
124 Right (a, b, c) => Right (Branch2 a b c)
126delType : Nat -> (k : Type) -> Type -> Ord k -> Type
127delType Z k v o = ()
128delType (S n) k v o = Tree n k v o
130treeDelete : k -> Tree n k v o -> Either (Tree n k v o) (delType n k v o)
131treeDelete k (Leaf k' v) =
132 if k == k' then
133 Right ()
134 else
135 Left (Leaf k' v)
136treeDelete {n=S Z} k (Branch2 t1 k' t2) =
137 if k <= k' then
138 case treeDelete k t1 of
139 Left t1' => Left (Branch2 t1' k' t2)
140 Right () => Right t2
141 else
142 case treeDelete k t2 of
143 Left t2' => Left (Branch2 t1 k' t2')
144 Right () => Right t1
145treeDelete {n=S Z} k (Branch3 t1 k1 t2 k2 t3) =
146 if k <= k1 then
147 case treeDelete k t1 of
148 Left t1' => Left (Branch3 t1' k1 t2 k2 t3)
149 Right () => Left (Branch2 t2 k2 t3)
150 else if k <= k2 then
151 case treeDelete k t2 of
152 Left t2' => Left (Branch3 t1 k1 t2' k2 t3)
153 Right () => Left (Branch2 t1 k1 t3)
154 else
155 case treeDelete k t3 of
156 Left t3' => Left (Branch3 t1 k1 t2 k2 t3')
157 Right () => Left (Branch2 t1 k1 t2)
158treeDelete {n=S (S _)} k (Branch2 t1 k' t2) =
159 if k <= k' then
160 case treeDelete k t1 of
161 Left t1' => Left (Branch2 t1' k' t2)
162 Right t1' =>
163 case t2 of
164 Branch2 a b c => Right (Branch3 t1' k' a b c)
165 Branch3 a b c d e => Left (branch4 t1' k' a b c d e)
166 else
167 case treeDelete k t2 of
168 Left t2' => Left (Branch2 t1 k' t2')
169 Right t2' =>
170 case t1 of
171 Branch2 a b c => Right (Branch3 a b c k' t2')
172 Branch3 a b c d e => Left (branch4 a b c d e k' t2')
173treeDelete {n=(S (S _))} k (Branch3 t1 k1 t2 k2 t3) =
174 if k <= k1 then
175 case treeDelete k t1 of
176 Left t1' => Left (Branch3 t1' k1 t2 k2 t3)
177 Right t1' => Left (merge1 t1' k1 t2 k2 t3)
178 else if k <= k2 then
179 case treeDelete k t2 of
180 Left t2' => Left (Branch3 t1 k1 t2' k2 t3)
181 Right t2' => Left (merge2 t1 k1 t2' k2 t3)
182 else
183 case treeDelete k t3 of
184 Left t3' => Left (Branch3 t1 k1 t2 k2 t3')
185 Right t3' => Left (merge3 t1 k1 t2 k2 t3')
187treeToList : Tree n k v o -> List (k, v)
188treeToList = treeToList' (:: [])
189 where
190 treeToList' : ((k, v) -> List (k, v)) -> Tree n k v o -> List (k, v)
191 treeToList' cont (Leaf k v) = cont (k, v)
192 treeToList' cont (Branch2 t1 _ t2) = treeToList' (:: treeToList' cont t2) t1
193 treeToList' cont (Branch3 t1 _ t2 _ t3) = treeToList' (:: treeToList' (:: treeToList' cont t3) t2) t1
195abstract
196data SortedMap : Type -> Type -> Type where
197 Empty : Ord k => SortedMap k v
198 M : (o : Ord k) => (n:Nat) -> Tree n k v o -> SortedMap k v
200public
201empty : Ord k => SortedMap k v
202empty = Empty
204public
205lookup : k -> SortedMap k v -> Maybe v
206lookup _ Empty = Nothing
207lookup k (M _ t) = treeLookup k t
209public
210insert : k -> v -> SortedMap k v -> SortedMap k v
211insert k v Empty = M Z (Leaf k v)
212insert k v (M _ t) =
213 case treeInsert k v t of
214 Left t' => (M _ t')
215 Right t' => (M _ t')
217public
218delete : k -> SortedMap k v -> SortedMap k v
219delete _ Empty = Empty
220delete k (M Z t) =
221 case treeDelete k t of
222 Left t' => (M _ t')
223 Right () => Empty
224delete k (M (S _) t) =
225 case treeDelete k t of
226 Left t' => (M _ t')
227 Right t' => (M _ t')
229public
230fromList : Ord k => List (k, v) -> SortedMap k v
231fromList l = foldl (flip (uncurry insert)) empty l
233public
234toList : SortedMap k v -> List (k, v)
235toList Empty = []
236toList (M _ t) = treeToList t
238public
239treeMap : (a -> b) -> Tree n k a o -> Tree n k b o
240treeMap f (Leaf k v) = Leaf k (f v)
241treeMap f (Branch2 t1 k t2) = Branch2 (treeMap f t1) k (treeMap f t2)
242treeMap f (Branch3 t1 k1 t2 k2 t3)
243 = Branch3 (treeMap f t1) k1 (treeMap f t2) k2 (treeMap f t3)
245instance Functor (SortedMap k) where
246 map _ Empty = Empty
247 map f (M n t) = M _ (treeMap f t)