@@ -0,0 +1,235 @@
1+module Data.Map
2+
3+-- TODO: write merge and split
4+
5+data Tree : Nat -> Type -> Type -> Type where
6+ Leaf : k -> v -> Tree O 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
9+
10+branch4 :
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
16+branch4 a b c d e f g =
17+ Branch2 (Branch2 a b c) d (Branch2 e f g)
18+
19+branch5 :
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
26+branch5 a b c d e f g h i =
27+ Branch2 (Branch2 a b c) d (Branch3 e f g h i)
28+
29+branch6 :
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
37+branch6 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)
39+
40+branch7 :
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
49+branch7 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)
51+
52+merge1 : Tree n k v -> k -> Tree (S n) k v -> k -> Tree (S n) k v -> Tree (S (S n)) k v
53+merge1 a b (Branch2 c d e) f (Branch2 g h i) = branch5 a b c d e f g h i
54+merge1 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
55+merge1 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
56+merge1 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
57+
58+merge2 : Tree (S n) k v -> k -> Tree n k v -> k -> Tree (S n) k v -> Tree (S (S n)) k v
59+merge2 (Branch2 a b c) d e f (Branch2 g h i) = branch5 a b c d e f g h i
60+merge2 (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
61+merge2 (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
62+merge2 (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
63+
64+merge3 : Tree (S n) k v -> k -> Tree (S n) k v -> k -> Tree n k v -> Tree (S (S n)) k v
65+merge3 (Branch2 a b c) d (Branch2 e f g) h i = branch5 a b c d e f g h i
66+merge3 (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
67+merge3 (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
68+merge3 (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
69+
70+treeLookup : Ord k => k -> Tree n k v -> Maybe v
71+treeLookup k (Leaf k' v) =
72+ if k == k' then
73+ Just v
74+ else
75+ Nothing
76+treeLookup k (Branch2 t1 k' t2) =
77+ if k <= k' then
78+ treeLookup k t1
79+ else
80+ treeLookup k t2
81+treeLookup 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
88+
89+treeInsert' : Ord k => k -> v -> Tree n k v -> Either (Tree n k v) (Tree n k v, k, Tree n k v)
90+treeInsert' 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)
95+treeInsert' 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)
104+treeInsert' 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 k2 t2, k2, Branch2 a b c)
118+
119+treeInsert : Ord k => k -> v -> Tree n k v -> Either (Tree n k v) (Tree (S n) k v)
120+treeInsert 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)
124+
125+delType : Nat -> Type -> Type -> Type
126+delType O k v = ()
127+delType (S n) k v = Tree n k v
128+
129+treeDelete : Ord k => k -> Tree n k v -> Either (Tree n k v) (delType n k v)
130+treeDelete k (Leaf k' v) =
131+ if k == k' then
132+ Right ()
133+ else
134+ Left (Leaf k' v)
135+treeDelete {n=S O} 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
144+treeDelete {n=S O} 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)
157+treeDelete {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')
172+treeDelete {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')
185+
186+-- FIXME: this is very inefficient
187+treeToList : Ord k => Tree n k v -> List (k, v)
188+treeToList (Leaf k v) = [(k, v)]
189+treeToList (Branch2 t1 _ t2) = treeToList t1 ++ treeToList t2
190+treeToList (Branch3 t1 _ t2 _ t3) = treeToList t1 ++ treeToList t2 ++ treeToList t3
191+
192+data Map : Type -> Type -> Type where
193+ Empty : Map k v
194+ M : (n:Nat) -> Tree n k v -> Map k v
195+
196+empty : Map k v
197+empty = Empty
198+
199+lookup : Ord k => k -> Map k v -> Maybe v
200+lookup _ Empty = Nothing
201+lookup k (M _ t) = treeLookup k t
202+
203+insert : Ord k => k -> v -> Map k v -> Map k v
204+insert k v Empty = M O (Leaf k v)
205+insert k v (M _ t) =
206+ case treeInsert k v t of
207+ Left t' => (M _ t')
208+ Right t' => (M _ t')
209+
210+delete : Ord k => k -> Map k v -> Map k v
211+delete _ Empty = Empty
212+delete k (M O t) =
213+ case treeDelete k t of
214+ Left t' => (M _ t')
215+ Right () => Empty
216+delete k (M (S _) t) =
217+ case treeDelete k t of
218+ Left t' => (M _ t')
219+ Right t' => (M _ t')
220+
221+fromList : Ord k => List (k, v) -> Map k v
222+fromList l = foldl (flip (uncurry insert)) empty l
223+
224+toList : Ord k => Map k v -> List (k, v)
225+toList Empty = []
226+toList (M _ t) = treeToList t
227+
228+instance Functor (Tree n k) where
229+ fmap f (Leaf k v) = Leaf k (f v)
230+ fmap f (Branch2 t1 k t2) = Branch2 (fmap f t1) k (fmap f t2)
231+ fmap f (Branch3 t1 k1 t2 k2 t3) = Branch3 (fmap f t1) k1 (fmap f t2) k2 (fmap f t3)
232+
233+instance Functor (Map k) where
234+ fmap _ Empty = Empty
235+ fmap f (M _ t) = M _ (fmap f t)