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