diff --git a/src/containers/treemap.d b/src/containers/treemap.d index 5003000..33e69bf 100644 --- 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)