Merge pull request #154 from dlang-community/more-treemap-fixes

dlang-bot · web-flow · commit f4e793755395 · 2019-10-10T22:27:47.000+02:00
Fix a few problems with TreeMap.getOrAdd
merged-on-behalf-of: Brian Schott &lt;Hackerpilot@users.noreply.github.com&gt;
diff --git a/src/containers/treemap.d b/src/containers/treemap.d
@@ -103,19 +103,25 @@ struct TreeMap(K, V, Allocator = Mallocator, alias less = "a < b",
 	 * Params:
 	 *     key = the key to look up
 	 *     value = the default value
+	 *
+	 * Returns: A pointer to the existing value, or a pointer to the inserted
+	 *     value.
 	 */
-	auto getOrAdd(this This)(const K key, lazy V value) @safe
+	auto getOrAdd(this This)(const K key, lazy V defaultValue)
 	{
 		alias CET = ContainerElementType!(This, V);
 		auto tme = TreeMapElement(key);
 		auto er = tree.equalRange(tme);
 		if (er.empty)
 		{
-			insert(value, key);
-			return value;
+			// TODO: This does two lookups and should be made faster.
+			tree.insert(TreeMapElement(key, defaultValue));
+			return cast(CET*) &tree.equalRange(tme)._containersFront().value;
 		}
 		else
-			return er.front.value;
+		{
+			return cast(CET*) &er._containersFront().value;
+		}
 	}
 
 	/**
@@ -412,14 +418,14 @@ version(emsi_containers_unittest) unittest
 {
 	TreeMap!(int, int) map;
 	auto p = map.getOrAdd(1, 1);
-	assert(p == 1);
+	assert(*p == 1);
 }
 
 version(emsi_containers_unittest) unittest
 {
 	import std.uuid : randomUUID;
 	import std.range.primitives : walkLength;
-	import std.stdio;
+	//import std.stdio;
 
 	auto hm = TreeMap!(string, int)();
 	foreach (i; 0 .. 1_000_000)
