004cc24763e5ef9d18afda73c180d002b8f6f29f
[nit.git] / lib / more_collections.nit
1 # This file is part of NIT ( http://www.nitlanguage.org ).
2 #
3 # Licensed under the Apache License, Version 2.0 (the "License");
4 # you may not use this file except in compliance with the License.
5 # You may obtain a copy of the License at
6 #
7 # http://www.apache.org/licenses/LICENSE-2.0
8 #
9 # Unless required by applicable law or agreed to in writing, software
10 # distributed under the License is distributed on an "AS IS" BASIS,
11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 # See the License for the specific language governing permissions and
13 # limitations under the License.
14
15 # Highly specific, but useful, collections-related classes.
16 module more_collections
17
18 # Simple way to store an `HashMap[K, Array[V]]`
19 #
20 # Unlike standard HashMap, MultiHashMap provides a new
21 # empty array on the first access on a unknown key.
22 #
23 # var m = new MultiHashMap[String, Char]
24 # assert not m.has_key("four")
25 # m["four"].add('i')
26 # m["four"].add('i')
27 # m["four"].add('i')
28 # m["four"].add('i')
29 # assert m.has_key("four")
30 # assert m["four"] == ['i', 'i', 'i', 'i']
31 # assert m["zzz"] == new Array[Char]
32 class MultiHashMap[K, V]
33 super HashMap[K, Array[V]]
34
35 # Add `v` to the array associated with `k`.
36 # If there is no array associated, then create it.
37 fun add_one(k: K, v: V)
38 do
39 var x = self.get_or_null(k)
40 if x != null then
41 x.add(v)
42 else
43 self[k] = [v]
44 end
45 end
46
47 redef fun provide_default_value(key) do
48 var res = new Array[V]
49 self[key] = res
50 return res
51 end
52 end
53
54 # Simple way to store an `HashMap[K1, HashMap[K2, V]]`
55 #
56 # ~~~~
57 # var hm2 = new HashMap2[Int, String, Float]
58 # hm2[1, "one"] = 1.0
59 # hm2[2, "two"] = 2.0
60 # assert hm2[1, "one"] == 1.0
61 # assert hm2[2, "not-two"] == null
62 # ~~~~
63 class HashMap2[K1, K2, V]
64 private var level1 = new HashMap[K1, HashMap[K2, V]]
65
66 # Return the value associated to the keys `k1` and `k2`.
67 # Return `null` if no such a value.
68 fun [](k1: K1, k2: K2): nullable V
69 do
70 var level1 = self.level1
71 var level2 = level1.get_or_null(k1)
72 if level2 == null then return null
73 return level2.get_or_null(k2)
74 end
75
76 # Set `v` the value associated to the keys `k1` and `k2`.
77 fun []=(k1: K1, k2: K2, v: V)
78 do
79 var level1 = self.level1
80 var level2 = level1.get_or_null(k1)
81 if level2 == null then
82 level2 = new HashMap[K2, V]
83 level1[k1] = level2
84 end
85 level2[k2] = v
86 end
87
88 # Remove the item at `k1` and `k2`
89 fun remove_at(k1: K1, k2: K2)
90 do
91 var level1 = self.level1
92 var level2 = level1.get_or_null(k1)
93 if level2 == null then return
94 level2.keys.remove(k2)
95 end
96 end
97
98 # Simple way to store an `HashMap[K1, HashMap[K2, HashMap[K3, V]]]`
99 #
100 # ~~~~
101 # var hm3 = new HashMap3[Int, String, Int, Float]
102 # hm3[1, "one", 11] = 1.0
103 # hm3[2, "two", 22] = 2.0
104 # assert hm3[1, "one", 11] == 1.0
105 # assert hm3[2, "not-two", 22] == null
106 # ~~~~
107 class HashMap3[K1, K2, K3, V]
108 private var level1 = new HashMap[K1, HashMap2[K2, K3, V]]
109
110 # Return the value associated to the keys `k1`, `k2`, and `k3`.
111 # Return `null` if no such a value.
112 fun [](k1: K1, k2: K2, k3: K3): nullable V
113 do
114 var level1 = self.level1
115 var level2 = level1.get_or_null(k1)
116 if level2 == null then return null
117 return level2[k2, k3]
118 end
119
120 # Set `v` the value associated to the keys `k1`, `k2`, and `k3`.
121 fun []=(k1: K1, k2: K2, k3: K3, v: V)
122 do
123 var level1 = self.level1
124 var level2 = level1.get_or_null(k1)
125 if level2 == null then
126 level2 = new HashMap2[K2, K3, V]
127 level1[k1] = level2
128 end
129 level2[k2, k3] = v
130 end
131
132 # Remove the item at `k1`, `k2` and `k3`
133 fun remove_at(k1: K1, k2: K2, k3: K3)
134 do
135 var level1 = self.level1
136 var level2 = level1.get_or_null(k1)
137 if level2 == null then return
138 level2.remove_at(k2, k3)
139 end
140 end
141
142 # A map with a default value.
143 #
144 # ~~~~
145 # var dm = new DefaultMap[String, Int](10)
146 # assert dm["a"] == 10
147 # ~~~~
148 #
149 # The default value is used when the key is not present.
150 # And getting a default value does not register the key.
151 #
152 # ~~~~
153 # assert dm["a"] == 10
154 # assert dm.length == 0
155 # assert dm.has_key("a") == false
156 # ~~~~
157 #
158 # It also means that removed key retrieve the default value.
159 #
160 # ~~~~
161 # dm["a"] = 2
162 # assert dm["a"] == 2
163 # dm.keys.remove("a")
164 # assert dm["a"] == 10
165 # ~~~~
166 #
167 # Warning: the default value is used as is, so using mutable object might
168 # cause side-effects.
169 #
170 # ~~~~
171 # var dma = new DefaultMap[String, Array[Int]](new Array[Int])
172 #
173 # dma["a"].add(65)
174 # assert dma["a"] == [65]
175 # assert dma.default == [65]
176 # assert dma["c"] == [65]
177 #
178 # dma["b"] += [66]
179 # assert dma["b"] == [65, 66]
180 # assert dma.default == [65]
181 # ~~~~
182 class DefaultMap[K, V]
183 super HashMap[K, V]
184
185 # The default value.
186 var default: V
187
188 redef fun provide_default_value(key) do return default
189 end