1 /**
2  * Associative Array utility functions
3  *
4  * License:
5  *   This Source Code Form is subject to the terms of
6  *   the Mozilla Public License, v. 2.0. If a copy of
7  *   the MPL was not distributed with this file, You
8  *   can obtain one at http://mozilla.org/MPL/2.0/.
9  *
10  * Authors:
11  *   Vladimir Panteleev <vladimir@thecybershadow.net>
12  */
13 
14 module ae.utils.aa;
15 
16 import std.algorithm;
17 import std.range;
18 import std.typecons;
19 
20 // ***************************************************************************
21 
22 /// Get a value from an AA, and throw an exception (not an error) if not found
23 ref auto aaGet(AA, K)(AA aa, K key)
24 	if (is(typeof(key in aa)))
25 {
26 	import std.conv;
27 
28 	auto p = key in aa;
29 	if (p)
30 		return *p;
31 	else
32 		static if (is(typeof(text(key))))
33 			throw new Exception("Absent value: " ~ text(key));
34 		else
35 			throw new Exception("Absent value");
36 }
37 
38 /// If key is not in aa, add it with defaultValue.
39 /// Returns a reference to the value corresponding to key.
40 ref V getOrAdd(K, V)(ref V[K] aa, K key, V defaultValue = V.init)
41 {
42 	auto p = key in aa;
43 	if (!p)
44 	{
45 		aa[key] = defaultValue;
46 		p = key in aa;
47 	}
48 	return *p;
49 }
50 
51 unittest
52 {
53 	int[int] aa;
54 	aa.getOrAdd(1, 2) = 3;
55 	assert(aa[1] == 3);
56 	assert(aa.getOrAdd(1, 4) == 3);
57 }
58 
59 struct KeyValuePair(K, V) { K key; V value; }
60 
61 /// Get key/value pairs from AA
62 KeyValuePair!(K, V)[] pairs(K, V)(V[K] aa)
63 {
64 	KeyValuePair!(K, V)[] result;
65 	foreach (key, value; aa)
66 		result ~= KeyValuePair!(K, V)(key, value);
67 	return result;
68 }
69 
70 /// Get key/value pairs from AA, sorted by keys
71 KeyValuePair!(K, V)[] sortedPairs(K, V)(V[K] aa)
72 {
73 	KeyValuePair!(K, V)[] result;
74 	foreach (key; aa.keys.sort)
75 		result ~= KeyValuePair!(K, V)(key, aa[key]);
76 	return result;
77 }
78 
79 /// Get values from AA, sorted by keys
80 V[] sortedValues(K, V)(in V[K] aa)
81 {
82 	V[] result;
83 	foreach (key; aa.keys.sort())
84 		result ~= aa[key];
85 	return result;
86 }
87 
88 /// Merge source into target. Return target.
89 V[K] merge(K, V)(auto ref V[K] target, in V[K] source)
90 {
91 	foreach (k, v; source)
92 		target[k] = v;
93 	return target;
94 }
95 
96 unittest
97 {
98 	int[int] target;
99 	int[int] source = [2:4];
100 	merge(target, source);
101 	assert(source == target);
102 
103 	target = [1:1, 2:2, 3:3];
104 	merge(target, source);
105 	assert(target == [1:1, 2:4, 3:3]);
106 
107 	assert(merge([1:1], [2:2]) == [1:1, 2:2]);
108 }
109 
110 /// Slurp a range of two elements (or two-element struct/class) into an AA.
111 auto toAA(R)(R r)
112 	if (is(typeof(r.front[1])))
113 {
114 	alias K = typeof(r.front[0]);
115 	alias V = typeof(r.front[1]);
116 	V[K] result;
117 	foreach (pair; r)
118 	{
119 		assert(pair.length == 2);
120 		result[pair[0]] = pair[1];
121 	}
122 	return result;
123 }
124 
125 /// ditto
126 auto toAA(R)(R r)
127 	if (is(typeof(r.front.tupleof)) && r.front.tupleof.length == 2 && !is(typeof(r.front[1])))
128 {
129 	return r.map!(el => tuple(el.tupleof)).toAA();
130 }
131 
132 unittest
133 {
134 	assert([[2, 4]].toAA() == [2:4]);
135 	assert([2:4].pairs.toAA() == [2:4]);
136 }
137 
138 // ***************************************************************************
139 
140 /// An associative array which retains the order in which elements were added.
141 struct OrderedMap(K, V)
142 {
143 	K[] keys;
144 	V[] values;
145 	size_t[K] index;
146 
147 	/// Convert from regular AA
148 	this(V[K] aa)
149 	{
150 		foreach (ref k, ref v; aa)
151 		{
152 			index[k] = values.length;
153 			keys ~= k;
154 			values ~= v;
155 		}
156 	}
157 
158 	ref inout(V) opIndex()(auto ref K k) inout
159 	{
160 		return values[index[k]];
161 	}
162 
163 	ref V opIndexAssign()(auto ref V v, auto ref K k)
164 	{
165 		auto pi = k in index;
166 		if (pi)
167 		{
168 			auto pv = &values[*pi];
169 			*pv = v;
170 			return *pv;
171 		}
172 
173 		index[k] = values.length;
174 		keys ~= k;
175 		values ~= v;
176 		return values[$-1];
177 	}
178 
179 	ref V getOrAdd()(auto ref K k)
180 	{
181 		auto pi = k in index;
182 		V* pv;
183 		if (pi)
184 			pv = &values[*pi];
185 		else
186 		{
187 			index[k] = values.length;
188 			keys ~= k;
189 			values ~= V.init;
190 			pv = &values[$-1];
191 		}
192 		return *pv;
193 	}
194 
195 	ref V opIndexUnary(string op)(auto ref K k)
196 	{
197 		auto pv = &getOrAdd(k);
198 		mixin("(*pv) " ~ op ~ ";");
199 		return *pv;
200 	}
201 
202 	ref V opIndexOpAssign(string op)(auto ref V v, auto ref K k)
203 	{
204 		auto pv = &getOrAdd(k);
205 		mixin("(*pv) " ~ op ~ "= v;");
206 		return *pv;
207 	}
208 
209 	inout(V) get()(auto ref K k, inout(V) defaultValue) inout
210 	{
211 		auto p = k in index;
212 		return p ? values[*p] : defaultValue;
213 	}
214 
215 	inout(V)* opIn_r()(auto ref K k) inout
216 	{
217 		auto p = k in index;
218 		return p ? &values[*p] : null;
219 	}
220 
221 	void remove()(auto ref K k)
222 	{
223 		auto i = index[k];
224 		index.remove(k);
225 		keys = keys.remove(i);
226 		values = values.remove(i);
227 	}
228 
229 	@property size_t length() const { return values.length; }
230 
231 	private int opApplyImpl(this This, Dg)(Dg dg)
232 	{
233 		int result = 0;
234 
235 		foreach (i, ref v; values)
236 		{
237 			result = dg(keys[i], v);
238 			if (result)
239 				break;
240 		}
241 		return result;
242 	}
243 
244 	int opApply(int delegate(ref K k, ref V v) dg)
245 	{
246 		return opApplyImpl(dg);
247 	}
248 
249 	int opApply(int delegate(const ref K k, const ref V v) dg) const
250 	{
251 		return opApplyImpl(dg);
252 	}
253 
254 	@property typeof(this) dup()
255 	{
256 		typeof(this) result;
257 		result.keys = keys.dup;
258 		result.values = values.dup;
259 		result.index = index.dup;
260 		return result;
261 	}
262 }
263 
264 unittest
265 {
266 	OrderedMap!(string, int) m;
267 	m["a"] = 1;
268 	m["b"] = 2;
269 	m["c"] = 3;
270 	assert(m.length == 3);
271 	assert("a" in m);
272 	assert("d" !in m);
273 	m.remove("a");
274 	assert(m.length == 2);
275 	m["x"] -= 1;
276 	assert(m["x"] == -1);
277 	++m["y"];
278 	assert(m["y"] == 1);
279 	auto cm = cast(const)m.dup;
280 	foreach (k, v; cm)
281 		if (k == "x")
282 			assert(v == -1);
283 }
284 
285 // ***************************************************************************
286 
287 /// Helper/wrapper for void[0][T]
288 struct HashSet(T)
289 {
290 	void[0][T] data;
291 
292 	alias data this;
293 
294 	this(R)(R r)
295 	{
296 		foreach (k; r)
297 			add(k);
298 	}
299 
300 	void add(T k)
301 	{
302 		void[0] v;
303 		data[k] = v;
304 	}
305 
306 	void remove(T k)
307 	{
308 		data.remove(k);
309 	}
310 
311 	@property HashSet!T dup() const
312 	{
313 		// Can't use .dup with void[0] value
314 		HashSet!T result;
315 		foreach (k, v; data)
316 			result.add(k);
317 		return result;
318 	}
319 
320 	int opApply(scope int delegate(ref T) dg)
321 	{
322 		int result;
323 		foreach (k, v; data)
324 			if ((result = dg(k)) != 0)
325 				break;
326 		return result;
327 	}
328 }
329 
330 unittest
331 {
332 	HashSet!int s;
333 	assert(s.length == 0);
334 	assert(!(1 in s));
335 	assert(1 !in s);
336 	s.add(1);
337 	assert(1 in s);
338 	assert(s.length == 1);
339 	foreach (k; s)
340 		assert(k == 1);
341 	s.remove(1);
342 	assert(s.length == 0);
343 
344 	s.add(1);
345 	auto t = s.dup;
346 	s.add(2);
347 	assert(t.length==1);
348 	t.remove(1);
349 	assert(t.length==0);
350 }
351 
352 auto toSet(R)(R r)
353 {
354 	alias E = ElementType!R;
355 	return HashSet!E(r);
356 }
357 
358 unittest
359 {
360 	auto set = [1, 2, 3].toSet();
361 	assert(2 in set);
362 	assert(4 !in set);
363 }
364 
365 // ***************************************************************************
366 
367 /// An object which acts mostly as an associative array,
368 /// with the added property of being able to hold keys with
369 /// multiple values. These are only exposed explicitly and
370 /// through iteration
371 struct MultiAA(K, V)
372 {
373 	V[][K] items;
374 
375 	/// If multiple items with this name are present,
376 	/// only the first one is returned.
377 	ref inout(V) opIndex(K key) inout
378 	{
379 		return items[key][0];
380 	}
381 
382 	V opIndexAssign(V value, K key)
383 	{
384 		items[key] = [value];
385 		return value;
386 	}
387 
388 	inout(V)* opIn_r(K key) inout @nogc
389 	{
390 		auto pvalues = key in items;
391 		if (pvalues && (*pvalues).length)
392 			return &(*pvalues)[0];
393 		return null;
394 	}
395 
396 	void remove(K key)
397 	{
398 		items.remove(key);
399 	}
400 
401 	// D forces these to be "ref"
402 	int opApply(int delegate(ref K key, ref V value) dg)
403 	{
404 		int ret;
405 		outer:
406 		foreach (key, values; items)
407 			foreach (ref value; values)
408 			{
409 				ret = dg(key, value);
410 				if (ret)
411 					break outer;
412 			}
413 		return ret;
414 	}
415 
416 	// Copy-paste because of https://issues.dlang.org/show_bug.cgi?id=7543
417 	int opApply(int delegate(ref const(K) key, ref const(V) value) dg) const
418 	{
419 		int ret;
420 		outer:
421 		foreach (key, values; items)
422 			foreach (ref value; values)
423 			{
424 				ret = dg(key, value);
425 				if (ret)
426 					break outer;
427 			}
428 		return ret;
429 	}
430 
431 	void add(K key, V value)
432 	{
433 		if (key !in items)
434 			items[key] = [value];
435 		else
436 			items[key] ~= value;
437 	}
438 
439 	V get(K key, lazy V def) const
440 	{
441 		auto pvalue = key in this;
442 		return pvalue ? *pvalue : def;
443 	}
444 
445 	inout(V)[] getAll(K key) inout
446 	{
447 		inout(V)[] result;
448 		foreach (ref value; items.get(key, null))
449 			result ~= value;
450 		return result;
451 	}
452 
453 	this(typeof(null) Null)
454 	{
455 	}
456 
457 	this(V[K] aa)
458 	{
459 		foreach (ref key, ref value; aa)
460 			add(key, value);
461 	}
462 
463 	this(V[][K] aa)
464 	{
465 		foreach (ref key, values; aa)
466 			foreach (ref value; values)
467 				add(key, value);
468 	}
469 
470 	@property auto keys() inout { return items.keys; }
471 
472 	// https://issues.dlang.org/show_bug.cgi?id=14626
473 
474 	@property V[] values()
475 	{
476 		return items.byValue.join;
477 	}
478 
479 	@property const(V)[] values() const
480 	{
481 		return items.byValue.join;
482 	}
483 
484 	@property typeof(V[K].init.pairs) pairs()
485 	{
486 		alias Pair = typeof(V[K].init.pairs[0]);
487 		Pair[] result;
488 		result.reserve(length);
489 		foreach (ref k, ref v; this)
490 			result ~= Pair(k, v);
491 		return result;
492 	}
493 
494 	@property size_t length() const { return items.byValue.map!(item => item.length).sum(); }
495 
496 	auto byKey() { return items.byKey(); }
497 	auto byValue() { return items.byValue().joiner(); }
498 
499 	bool opCast(T)() inout
500 		if (is(T == bool))
501 	{
502 		return !!items;
503 	}
504 
505 	/// Warning: discards repeating items
506 	V[K] opCast(T)() const
507 		if (is(T == V[K]))
508 	{
509 		V[K] result;
510 		foreach (key, value; this)
511 			result[key] = value;
512 		return result;
513 	}
514 
515 	V[][K] opCast(T)() inout
516 		if (is(T == V[][K]))
517 	{
518 		V[][K] result;
519 		foreach (k, v; this)
520 			result[k] ~= v;
521 		return result;
522 	}
523 }
524 
525 unittest
526 {
527 	MultiAA!(string, string) aa;
528 }