X-Git-Url: http://nitlanguage.org

diff --git a/lib/more_collections.nit b/lib/more_collections.nit
index aad9f4c..1bcc514 100644
--- a/lib/more_collections.nit
+++ b/lib/more_collections.nit
@@ -16,6 +16,19 @@
 module more_collections
 
 # Simple way to store an `HashMap[K, Array[V]]`
+#
+# Unlike standard HashMap, MultiHashMap provides a new
+# empty array on the first access on a unknown key.
+#
+#     var m = new MultiHashMap[String, Char]
+#     assert not m.has_key("four")
+#     m["four"].add('i')
+#     m["four"].add('i')
+#     m["four"].add('i')
+#     m["four"].add('i')
+#     assert m.has_key("four")
+#     assert m["four"] == ['i', 'i', 'i', 'i']
+#     assert m["zzz"] == new Array[Char]
 class MultiHashMap[K: Object, V]
 	super HashMap[K, Array[V]]
 
@@ -30,12 +43,24 @@ class MultiHashMap[K: Object, V]
 		end
 	end
 
-	init do end
+	redef fun provide_default_value(key) do
+		var res = new Array[V]
+		self[key] = res
+		return res
+	end
 end
 
 # Simple way to store an `HashMap[K1, HashMap[K2, V]]`
+#
+# ~~~~
+# var hm2 = new HashMap2[Int, String, Float]
+# hm2[1, "one"] = 1.0
+# hm2[2, "two"] = 2.0
+# assert hm2[1, "one"] == 1.0
+# assert hm2[2, "not-two"] == null
+# ~~~~
 class HashMap2[K1: Object, K2: Object, V]
-	private var level1: HashMap[K1, HashMap[K2, V]] = new HashMap[K1, HashMap[K2, V]]
+	private var level1 = new HashMap[K1, HashMap[K2, V]]
 
 	# Return the value associated to the keys `k1` and `k2`.
 	# Return `null` if no such a value.
@@ -64,8 +89,16 @@ class HashMap2[K1: Object, K2: Object, V]
 end
 
 # Simple way to store an `HashMap[K1, HashMap[K2, HashMap[K3, V]]]`
+#
+# ~~~~
+# var hm3 = new HashMap3[Int, String, Int, Float]
+# hm3[1, "one", 11] = 1.0
+# hm3[2, "two", 22] = 2.0
+# assert hm3[1, "one", 11] == 1.0
+# assert hm3[2, "not-two", 22] == null
+# ~~~~
 class HashMap3[K1: Object, K2: Object, K3: Object, V]
-	private var level1: HashMap[K1, HashMap2[K2, K3, V]] = new HashMap[K1, HashMap2[K2, K3, V]]
+	private var level1 = new HashMap[K1, HashMap2[K2, K3, V]]
 
 	# Return the value associated to the keys `k1`, `k2`, and `k3`.
 	# Return `null` if no such a value.
@@ -91,3 +124,52 @@ class HashMap3[K1: Object, K2: Object, K3: Object, V]
 		level2[k2, k3] = v
 	end
 end
+
+# A map with a default value.
+#
+# ~~~~
+# var dm = new DefaultMap[String, Int](10)
+# assert dm["a"] == 10
+# ~~~~
+#
+# The default value is used when the key is not present.
+# And getting a default value does not register the key.
+#
+# ~~~~
+# assert dm["a"] == 10
+# assert dm.length == 0
+# assert dm.has_key("a") == false
+# ~~~~
+#
+# It also means that removed key retrieve the default value.
+#
+# ~~~~
+# dm["a"] = 2
+# assert dm["a"] == 2
+# dm.keys.remove("a")
+# assert dm["a"] == 10
+# ~~~~
+#
+# Warning: the default value is used as is, so using mutable object might
+# cause side-effects.
+#
+# ~~~~
+# var dma = new DefaultMap[String, Array[Int]](new Array[Int])
+#
+# dma["a"].add(65)
+# assert dma["a"] == [65]
+# assert dma.default == [65]
+# assert dma["c"] == [65]
+#
+# dma["b"] += [66]
+# assert dma["b"] == [65, 66]
+# assert dma.default == [65]
+# ~~~~
+class DefaultMap[K: Object, V]
+	super HashMap[K, V]
+
+	# The default value.
+	var default: V
+
+	redef fun provide_default_value(key) do return default
+end