git.iko.soy/Idris-dev
on f6edbf0 help sign in
Idris-dev/lib/Data/SortedMap.idr
236 lines 7.6 kB history raw
1module Data.SortedMap
3import Prelude
4-- TODO: write merge and split
6data Tree : Nat -> Type -> Type -> Type where
7 Leaf : k -> v -> Tree Z k v
8 Branch2 : Tree n k v -> k -> Tree n k v -> Tree (S n) k v
9 Branch3 : Tree n k v -> k -> Tree n k v -> k -> Tree n k v -> Tree (S n) k v
11branch4 :
12 Tree n k v -> k ->
13 Tree n k v -> k ->
14 Tree n k v -> k ->
15 Tree n k v ->
16 Tree (S (S n)) k v
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 -> k ->
22 Tree n k v -> k ->
23 Tree n k v -> k ->
24 Tree n k v -> k ->
25 Tree n k v ->
26 Tree (S (S n)) k v
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 -> k ->
32 Tree n k v -> k ->
33 Tree n k v -> k ->
34 Tree n k v -> k ->
35 Tree n k v -> k ->
36 Tree n k v ->
37 Tree (S (S n)) k v
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 -> 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 -> k ->
48 Tree n k v ->
49 Tree (S (S n)) k v
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 -> k -> Tree (S n) k v -> k -> Tree (S n) k v -> Tree (S (S n)) k v
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 -> k -> Tree n k v -> k -> Tree (S n) k v -> Tree (S (S n)) k v
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 -> k -> Tree (S n) k v -> k -> Tree n k v -> Tree (S (S n)) k v
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 : Ord k => k -> Tree n k v -> 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' : Ord k => k -> v -> Tree n k v -> Either (Tree n k v) (Tree n k v, k, Tree n k v)
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 k2 t2, k2, Branch2 a b c)
120treeInsert : Ord k => k -> v -> Tree n k v -> Either (Tree n k v) (Tree (S n) k v)
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 -> Type -> Type -> Type
127delType Z k v = ()
128delType (S n) k v = Tree n k v
130treeDelete : Ord k => k -> Tree n k v -> Either (Tree n k v) (delType n k v)
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')
187-- FIXME: this is very inefficient
188treeToList : Ord k => Tree n k v -> List (k, v)
189treeToList (Leaf k v) = [(k, v)]
190treeToList (Branch2 t1 _ t2) = treeToList t1 ++ treeToList t2
191treeToList (Branch3 t1 _ t2 _ t3) = treeToList t1 ++ treeToList t2 ++ treeToList t3
193data SortedMap : Type -> Type -> Type where
194 Empty : SortedMap k v
195 M : (n:Nat) -> Tree n k v -> SortedMap k v
197empty : SortedMap k v
198empty = Empty
200lookup : Ord k => k -> SortedMap k v -> Maybe v
201lookup _ Empty = Nothing
202lookup k (M _ t) = treeLookup k t
204insert : Ord k => k -> v -> SortedMap k v -> SortedMap k v
205insert k v Empty = M Z (Leaf k v)
206insert k v (M _ t) =
207 case treeInsert k v t of
208 Left t' => (M _ t')
209 Right t' => (M _ t')
211delete : Ord k => k -> SortedMap k v -> SortedMap k v
212delete _ Empty = Empty
213delete k (M Z t) =
214 case treeDelete k t of
215 Left t' => (M _ t')
216 Right () => Empty
217delete k (M (S _) t) =
218 case treeDelete k t of
219 Left t' => (M _ t')
220 Right t' => (M _ t')
222fromList : Ord k => List (k, v) -> SortedMap k v
223fromList l = foldl (flip (uncurry insert)) empty l
225toList : Ord k => SortedMap k v -> List (k, v)
226toList Empty = []
227toList (M _ t) = treeToList t
229instance Functor (Tree n k) where
230 map f (Leaf k v) = Leaf k (f v)
231 map f (Branch2 t1 k t2) = Branch2 (map f t1) k (map f t2)
232 map f (Branch3 t1 k1 t2 k2 t3) = Branch3 (map f t1) k1 (map f t2) k2 (map f t3)
234instance Functor (SortedMap k) where
235 map _ Empty = Empty
236 map f (M n t) = M _ (map f t)