/**
 * Associative Array utility functions
 *
 * License:
 *   This Source Code Form is subject to the terms of
 *   the Mozilla Public License, v. 2.0. If a copy of
 *   the MPL was not distributed with this file, You
 *   can obtain one at http://mozilla.org/MPL/2.0/.
 *
 * Authors:
 *   Vladimir Panteleev <vladimir@thecybershadow.net>
 */

module ae.utils.aa;

import std.algorithm;
import std.range;
import std.traits;
import std.typecons;

// ***************************************************************************

/// Polyfill for object.require
static if (!__traits(hasMember, object, "require"))
ref V require(K, V)(ref V[K] aa, K key, lazy V value = V.init)
{
	auto p = key in aa;
	if (p)
		return *p;
	return aa[key] = value;
}

unittest
{
	int[int] aa;
	aa.require(1, 2);
	assert(aa[1] == 2);
	aa.require(2, 3) = 4;
	assert(aa[2] == 4);
	aa.require(1, 5);
	assert(aa[1] == 2);
	aa.require(1, 6) = 7;
	assert(aa[1] == 7);
}

static if (!__traits(hasMember, object, "update"))
{
	/// Polyfill for object.update
	void updatePolyfill(K, V, C, U)(ref V[K] aa, K key, scope C create, scope U update)
	if (is(typeof(create()) : V) && is(typeof(update(aa[K.init])) : V))
	{
		auto p = key in aa;
		if (p)
			*p = update(*p);
		else
			aa[key] = create();
	}

	/// Work around https://issues.dlang.org/show_bug.cgi?id=15795
	alias update = updatePolyfill;
}

// https://github.com/dlang/druntime/pull/3012
private enum haveObjectUpdateWithVoidUpdate = is(typeof({
	int[int] aa;
	.object.update(aa, 0, { return 0; }, (ref int v) { });
}));

static if (!haveObjectUpdateWithVoidUpdate)
{
	/// Polyfill for object.update with void update function
	void updateVoid(K, V, C, U)(ref V[K] aa, K key, scope C create, scope U update)
	if (is(typeof(create()) : V) && is(typeof(update(aa[K.init])) == void))
	{
		// We can polyfill this in two ways.
		// What's more expensive, copying the value, or a second key lookup?
		enum haveObjectUpdate = __traits(hasMember, object, "update");
		enum valueIsExpensiveToCopy = V.sizeof > string.sizeof
			|| hasElaborateCopyConstructor!V
			|| hasElaborateDestructor!V;
		static if (haveObjectUpdate && !valueIsExpensiveToCopy)
		{
			.object.update(aa, key, create,
				(ref V v) { update(v); return v; });
		}
		else
		{
			auto p = key in aa;
			if (p)
				update(*p);
			else
				aa[key] = create();
		}
	}

	/// Work around https://issues.dlang.org/show_bug.cgi?id=15795
	alias update = updateVoid;
}
else
	alias updateVoid = object.update;

// Inject overload
static if (__traits(hasMember, object, "update"))
	private alias update = object.update;

// ***************************************************************************

/// Get a value from an AA, and throw an exception (not an error) if not found
ref auto aaGet(AA, K)(auto ref AA aa, auto ref K key)
	if (is(typeof(key in aa)))
{
	import std.conv;

	auto p = key in aa;
	if (p)
		return *p;
	else
		static if (is(typeof(text(key))))
			throw new Exception("Absent value: " ~ text(key));
		else
			throw new Exception("Absent value");
}

/// If key is not in aa, add it with defaultValue.
/// Returns a reference to the value corresponding to key.
ref V getOrAdd(K, V)(ref V[K] aa, auto ref K key, auto ref V defaultValue)
{
	return aa.require(key, defaultValue);
}

/// ditto
ref V getOrAdd(K, V)(ref V[K] aa, auto ref K key)
{
	return getOrAdd(aa, key, V.init);
}

unittest
{
	int[int] aa;
	aa.getOrAdd(1, 2) = 3;
	assert(aa[1] == 3);
	assert(aa.getOrAdd(1, 4) == 3);
}

/// If key is not in aa, add it with the given value, and return true.
/// Otherwise, return false.
bool addNew(K, V)(ref V[K] aa, auto ref K key, auto ref V value)
{
	bool added = void;
	updateVoid(aa, key,
		delegate V   (       ) { added = true ; return value; },
		delegate void(ref V v) { added = false;               },
	);
	return added;
}

unittest
{
	int[int] aa;
	assert( aa.addNew(1, 2));
	assert(!aa.addNew(1, 3));
	assert(aa[1] == 2);
}

// ***************************************************************************

struct KeyValuePair(K, V) { K key; V value; }

/// Get key/value pairs from AA
deprecated KeyValuePair!(K, V)[] pairs(K, V)(V[K] aa)
{
	KeyValuePair!(K, V)[] result;
	foreach (key, value; aa)
		result ~= KeyValuePair!(K, V)(key, value);
	return result;
}

/// Get key/value pairs from AA, sorted by keys
KeyValuePair!(K, V)[] sortedPairs(K, V)(V[K] aa)
{
	KeyValuePair!(K, V)[] result;
	foreach (key; aa.keys.sort)
		result ~= KeyValuePair!(K, V)(key, aa[key]);
	return result;
}

/// Get values from AA, sorted by keys
V[] sortedValues(K, V)(in V[K] aa)
{
	V[] result;
	foreach (key; aa.keys.sort())
		result ~= aa[key];
	return result;
}

/// Merge source into target. Return target.
V[K] merge(K, V)(auto ref V[K] target, in V[K] source)
{
	foreach (k, v; source)
		target[k] = v;
	return target;
}

unittest
{
	int[int] target;
	int[int] source = [2:4];
	merge(target, source);
	assert(source == target);

	target = [1:1, 2:2, 3:3];
	merge(target, source);
	assert(target == [1:1, 2:4, 3:3]);

	assert(merge([1:1], [2:2]) == [1:1, 2:2]);
}

/// Slurp a range of two elements (or two-element struct/class) into an AA.
auto toAA(R)(R r)
	if (is(typeof(r.front[1])))
{
	alias K = typeof(r.front[0]);
	alias V = typeof(r.front[1]);
	V[K] result;
	foreach (pair; r)
	{
		assert(pair.length == 2);
		result[pair[0]] = pair[1];
	}
	return result;
}

/// ditto
auto toAA(R)(R r)
	if (is(typeof(r.front.tupleof)) && r.front.tupleof.length == 2 && !is(typeof(r.front[1])))
{
	return r.map!(el => tuple(el.tupleof)).toAA();
}

deprecated unittest
{
	assert([[2, 4]].toAA() == [2:4]);
	assert([2:4].pairs.toAA() == [2:4]);
}

/// Ensure that arr is non-null if empty.
V[K] nonNull(K, V)(V[K] aa)
{
	if (aa !is null)
		return aa;
	aa[K.init] = V.init;
	aa.remove(K.init);
	assert(aa !is null);
	return aa;
}

unittest
{
	int[int] aa;
	assert(aa is null);
	aa = aa.nonNull;
	assert(aa !is null);
	assert(aa.length == 0);
}

// ***************************************************************************

// Helpers for HashCollection
private
{
	alias Void = void[0]; // Zero-sized type
	static assert(Void.sizeof == 0);

	// Abstraction layer for single/multi-value type holding one or many T.
	// Optimizer representation for Void.
	struct SingleOrMultiValue(bool multi, T)
	{
		alias ValueType = Select!(multi,
			// multi==true
			Select!(is(T == Void),
				size_t, // "store" the items by keeping track of their count only.
				T[],
			),

			// multi==false
			Select!(is(T == Void),
				Void,
				T[1],
			),
		);

		// Using free functions instead of struct methods,
		// as structs always have non-zero size.
	static:

		size_t length(ref const ValueType v) nothrow
		{
			static if (is(T == Void))
				static if (multi)
					return v; // count
				else
					return 1;
			else
				return v.length; // static or dynamic array
		}
	}
}

/// Base type for ordered/unordered single-value/multi-value map/set
/*private*/ struct HashCollection(K, V, bool ordered, bool multi)
{
private:
	enum bool haveValues = !is(V == void); // Not a set

	// The type for values used when a value variable is needed
	alias ValueVarType = Select!(haveValues, V, Void);

	// The type of a single element of the values of `this.lookup`.
	// When ordered==true, we use size_t (index into `this.items`).
	alias LookupItem = Select!(ordered, size_t, ValueVarType);

	// The type of the values of `this.lookup`.
	alias SM = SingleOrMultiValue!(multi, LookupItem);
	alias LookupValue = SM.ValueType;

	static if (haveValues)
	{
		alias ReturnType(Fallback) = V;
		alias SingleIterationType = V;
		alias OpIndexKeyType = K;
		alias OpIndexValueType = V;
	}
	else
	{
		alias SingleIterationType = const(K);
		static if (ordered)
		{
			alias OpIndexKeyType = size_t;
			alias OpIndexValueType = K;
			alias ReturnType(Fallback) = K;
		}
		else
		{
			alias OpIndexKeyType = void;
			alias OpIndexValueType = void;
			alias ReturnType(Fallback) = Fallback;
		}
	}
	enum haveReturnType = !is(ReturnType!void == void);
	enum haveIndexing = haveValues || ordered;

	alias IK = OpIndexKeyType;
	alias IV = OpIndexValueType;

	// The contract we try to follow is that adding/removing items in
	// one copy of the object will not affect other copies.
	// Therefore, when we have array fields, make sure they are dup'd
	// on copy, so that we don't trample older copies' data.
	enum bool needDupOnCopy = ordered;

	static if (haveReturnType)
	{
		static if (ordered)
			/*  */ ref inout(ReturnType!void) lookupToReturnValue(in        LookupItem  lookupItem) inout { return items[lookupItem].returnValue; }
		else
			static ref inout(ReturnType!void) lookupToReturnValue(ref inout(LookupItem) lookupItem)       { return       lookupItem             ; }
	}

	// *** Data ***

	// This is used for all key hash lookups.
	LookupValue[K] lookup;

	static if (ordered)
	{
		struct Item
		{
			K key;
			ValueVarType value;

			static if (haveValues)
				private alias returnValue = value;
			else
				private alias returnValue = key;
		}
		Item[] items;

		enum bool canDup = is(typeof(lookup.dup)) && is(typeof(items.dup));
	}
	else
	{
		enum bool canDup = is(typeof(lookup.dup));
	}

public:

	// *** Lifetime ***

	/// Postblit
	static if (needDupOnCopy)
	{
		static if (canDup)
			this(this)
			{
				lookup = lookup.dup;
				items = items.dup;
			}
		else
			@disable this(this);
	}

	/// Create shallow copy
	static if (canDup)
	typeof(this) dup()
	{
		static if (needDupOnCopy)
			return this;
		else
		{
			typeof(this) copy;
			copy.lookup = lookup.dup;
			static if (ordered)
				copy.items = items.dup;
			return copy;
		}
	}
	
	// *** Conversions (from) ***

	/// Construct from something else
	this(Input)(Input input)
	if (is(typeof(opAssign(input))))
	{
		opAssign(input);
	}

	/// Null assignment
	ref typeof(this) opAssign(typeof(null) _)
	{
		clear();
		return this;
	}

	/// Convert from an associative type
	ref typeof(this) opAssign(AA)(AA aa)
	if (haveValues
		&& !is(AA : typeof(this))
		&& is(typeof({ foreach (ref k, ref v; aa) add(k, v); })))
	{
		clear();
		foreach (ref k, ref v; aa)
			add(k, v);
		return this;
	}

	/// Convert from an associative type of multiple items
	ref typeof(this) opAssign(AA)(AA aa)
	if (haveValues
		&& multi
		&& !is(AA : typeof(this))
		&& is(typeof({ foreach (ref k, ref vs; aa) foreach (ref v; vs) add(k, v); })))
	{
		clear();
		foreach (ref k, ref vs; aa)
			foreach (ref v; vs)
				add(k, v);
		return this;
	}

	/// Convert from a range of tuples
	ref typeof(this) opAssign(R)(R input)
	if (haveValues
		&& is(typeof({ foreach (ref pair; input) add(pair[0], pair[1]); }))
		&& !is(typeof({ foreach (ref k, ref v; input) add(k, v); }))
		&& is(typeof(input.front.length))
		&& input.front.length == 2)
	{
		clear();
		foreach (ref pair; input)
			add(pair[0], pair[1]);
		return this;
	}

	/// Convert from a range of key/value pairs
	ref typeof(this) opAssign(R)(R input)
	if (haveValues
		&& is(typeof({ foreach (ref pair; input) add(pair.key, pair.value); }))
		&& !is(typeof({ foreach (ref k, ref v; input) add(k, v); })))
	{
		clear();
		foreach (ref pair; input)
			add(pair.key, pair.value);
		return this;
	}

	/// Convert from a range of values
	ref typeof(this) opAssign(R)(R input)
	if (!haveValues
		&& !is(R : typeof(this))
		&& is(typeof({ foreach (ref v; input) add(v); })))
	{
		clear();
		foreach (ref v; input)
			add(v);
		return this;
	}

	// *** Conversions (to) ***

	/// Convert to bool (true if non-null)
	bool opCast(T)() const
	if (is(T == bool))
	{
		return lookup !is null;
	}

	/// Convert to D associative array
	static if (!ordered)
	{
		const(LookupValue[K]) toAA() const
		{
			return lookup;
		}

		LookupValue[K] toAA()
		{
			return lookup.dup;
		}

		deprecated alias items = toAA;
	}

	// *** Query (basic) ***

	/// True when there are no items.
	bool empty() pure const nothrow @nogc @trusted
	{
		static if (ordered)
			return items.length == 0; // optimization
		else
			return lookup.byKey.empty; // generic version
	}

	/// Total number of items, including with duplicate keys.
	size_t length() pure const nothrow @nogc @trusted
	{
		static if (ordered)
			return items.length; // optimization
		else
		static if (!multi)
			return lookup.length; // optimization
		else // generic version
		{
			size_t result;
			foreach (ref v; lookup.byValue)
				result += SM.length(v);
			return result;
		}
	}

	// *** Query (by key) ***

	/// Check if item with this key has been added.
	/// When applicable, return a pointer to the last value added with this key.
	Select!(haveReturnType, inout(ReturnType!void)*, bool) opBinaryRight(string op : "in", _K)(auto ref _K key) inout
	if (is(typeof(key in lookup)))
	{
		enum missValue = select!haveReturnType(null, false);

		auto p = key in lookup;
		if (!p)
			return missValue;

		static if (haveReturnType)
			return &lookupToReturnValue((*p)[$-1]);
		else
			return true;
	}

	/// Index operator.
	/// The key must exist. Indexing with a key which does not exist
	/// is an error.
	static if (haveIndexing)
	ref inout(IV) opIndex()(auto ref IK k) inout
	{
		static if (haveValues)
			return lookupToReturnValue(lookup[k][$-1]);
		else
			return items[k].returnValue;
	}

	/// Retrieve last value associated with key, or `defaultValue` if none.
	static if (haveIndexing)
	auto ref inout(IV) get()(auto ref IK k, auto ref inout(IV) defaultValue) inout
	{
		static if (haveValues)
		{
			auto p = k in lookup;
			return p ? lookupToReturnValue((*p)[$-1]) : defaultValue;
		}
		else
			return k < items.length ? items[k].returnValue : defaultValue;
	}

	// *** Query (ranges) ***

	/// Return a range which iterates over key/value pairs.
	static if (haveValues)
	auto byKeyValue(this This)()
	{
		static if (ordered)
			return items;
		else
		{
			return lookup
				.byKeyValue
				.map!(pair =>
					pair
					.value
					.map!(value => KeyValuePair!(K, V)(pair.key, value))
				)
				.joiner;
		}
	}

	/// ditto
	static if (haveValues)
	auto byPair(this This)()
	{
		return byKeyValue
			.map!(pair => tuple!("key", "value")(pair.key, pair.value));
	}

	/// Return a range which iterates over all keys.
	/// Duplicate keys will occur several times in the range.
	auto byKey(this This)()
	{
		static if (ordered)
		{
			static ref getKey(MItem)(ref MItem item) { return item.key; }
			return items.map!getKey;
		}
		else
		{
			return lookup
				.byKeyValue
				.map!(pair =>
					pair.key.repeat(SM.length(pair.value))
				)
				.joiner;
		}
	}

	/// Return a range which iterates over all values.
	static if (haveValues)
	auto byValue(this This)()
	{
		static if (ordered)
		{
			static ref getValue(MItem)(ref MItem item) { return item.value; }
			return items.map!getValue;
		}
		else
		{
			return lookup
				.byKeyValue
				.map!(pair =>
					pair
					.value
				)
				.joiner;
		}
	}

	@property auto keys(this This)() { return byKey.array; }
	@property auto values(this This)() { return byValue.array; }

	// *** Query (search by key) ***

	static if (ordered)
	{
		size_t indexOf()(auto ref const K k)
		{
			auto p = k in lookup;
			return p ? (*p)[0] : -1;
		}

		size_t[] indicesOf()(auto ref const K k)
		{
			auto p = k in lookup;
			return p ? (*p)[] : null;
		}
	}

	/// Return the number of items with the given key.
	/// When multi==false, always returns 0 or 1.
	size_t count()(auto ref K k)
	{
		static if (ordered)
			return indicesOf(k).length;
		else
		{
			auto p = k in lookup;
			return p ? SM.length(*p) : 0;
		}
	}

	/// Return a range with all values with the given key.
	/// If the key is not present, returns an empty range.
	static if (haveValues)
	auto byValueOf(this This)(auto ref K k)
	{
		static if (ordered)
			return indicesOf(k).map!(index => items[index].value);
		else
			return valuesOf(k);
	}

	/// Return an array with all values with the given key.
	/// If the key is not present, returns an empty array.
	static if (haveValues)
	V[] valuesOf()(auto ref K k)
	{
		static if (ordered)
			return byValueOf(k).array;
		else
		{
			static if (multi)
				return lookup.get(k, null);
			else
			{
				auto p = k in lookup;
				return p ? (*p)[] : null;
			}
		}
	}

	static if (haveValues)
	deprecated alias getAll = valuesOf;

	// *** Iteration ***

	private int opApplyImpl(this This, Dg)(Dg dg)
	{
		enum single = arity!dg == 1;

		int result = 0;

		static if (ordered)
		{
			foreach (ref item; items)
			{
				static if (single)
					result = dg(item.returnValue);
				else
					result = dg(item.key, item.value);
				if (result)
					break;
			}
		}
		else
		{
		outer:
			foreach (ref key, ref values; lookup)
				static if (haveValues)
				{
					foreach (ref value; values)
					{
						static if (single)
							result = dg(value);
						else
							result = dg(key, value);
						if (result)
							break outer;
					}
				}
				else
				{
					foreach (iteration; 0 .. SM.length(values))
					{
						static assert(single);
						result = dg(key);
						if (result)
							break outer;
					}
				}
		}
		return result;
	}

	/// Iterate over keys (sets) / values (maps).
	int opApply(int delegate(ref SingleIterationType x) dg)
	{
		return opApplyImpl(dg);
	}

	/// ditto
	int opApply(int delegate(const ref SingleIterationType x) dg) const
	{
		return opApplyImpl(dg);
	}

	static if (haveValues)
	{
		/// Iterate over keys and values.
		int opApply(int delegate(K k, ref V v) dg)
		{
			return opApplyImpl(dg);
		}

		/// ditto
		int opApply(int delegate(K k, const ref V v) dg) const
		{
			return opApplyImpl(dg);
		}
	}

	// *** Mutation (addition) ***

	private enum AddMode
	{
		add,     /// Always add value
		replace, /// Replace all previous values
		require, /// Only add value if it did not exist before
	}

	private ref ReturnType!void addImpl(AddMode mode, AK, GV)(ref AK key, scope GV getValue)
	if (is(AK : K))
	{
		static if (ordered)
		{
			size_t addedIndex;

			static if (multi && mode == AddMode.add)
			{
				addedIndex = items.length;
				lookup[key] ~= addedIndex;
				items ~= Item(key, getValue());
			}
			else
			{
				lookup.updateVoid(key,
					delegate LookupValue()
					{
						addedIndex = items.length;
						items ~= Item(key, getValue());
						return [addedIndex];
					},
					delegate void(ref LookupValue existingIndex)
					{
						addedIndex = existingIndex[0];
						static if (mode != AddMode.require)
						{
							static if (multi)
							{
								static assert(mode == AddMode.replace);
								existingIndex = existingIndex[0 .. 1];
							}
							items[addedIndex].value = getValue();
						}
					});
			}

			return items[addedIndex].returnValue;
		}
		else // ordered
		{
			static if (haveValues)
			{
				static if (mode == AddMode.require)
					return (lookup.require(key, [getValue()]))[0];
				else
				static if (multi && mode == AddMode.add)
					return (lookup[key] ~= getValue())[$-1];
				else
					return (lookup[key] = [getValue()])[0];
			}
			else
			{
				static if (multi)
				{
					static if (mode == AddMode.require)
						lookup.require(key, 1);
					else
					static if (mode == AddMode.add)
						lookup[key]++;
					else
						lookup[key] = 1;
				}
				else
					lookup[key] = LookupValue.init;
				// This branch returns void, as there is no reasonable
				// ref to an AA key that we can return here.
			}
		}
	}

	/*private*/ template addSetFunc(AddMode mode)
	{
		static if (haveValues)
		{
			ref ReturnType!void addSetFunc(AK, AV)(auto ref AK key, auto ref AV value)
			if (is(AK : K) && is(AV : V))
			{
				return addImpl!mode(key, () => value);
			}
		}
		else
		{
			ref ReturnType!void addSetFunc(AK)(auto ref AK key)
			if (is(AK : K))
			{
				ValueVarType value; // void[0]
				return addImpl!mode(key, () => value);
			}
		}
	}

	/// Add an item.
	alias add = addSetFunc!(AddMode.add);

	/// Ensure a key exists (with the given value).
	/// When `multi==true`, replaces all previous entries with this key.
	/// Otherwise, behaves identically to `add`.
	alias set = addSetFunc!(AddMode.replace);

	/// Add `value` only if `key` is not present.
	static if (haveValues)
	ref V require()(auto ref K key, lazy V value = V.init)
	{
		return addImpl!(AddMode.require)(key, () => value);
	}

	deprecated alias getOrAdd = require;

	private alias UpdateFuncRT(U) = typeof({ U u = void; V v = void; return u(v); }());

	/// If `key` is present, call `update` for every value;
	/// otherwise, add new value with `create`.
	static if (haveValues)
	void update(C, U)(auto ref K key, scope C create, scope U update)
	if (is(typeof(create()) : V) && (is(UpdateFuncRT!U : V) || is(UpdateFuncRT!U == void)))
	{
		static if (ordered)
		{
			lookup.updateVoid(key,
				delegate LookupValue()
				{
					auto addedIndex = items.length;
					items ~= Item(key, create());
					return [addedIndex];
				},
				delegate void(ref LookupValue existingIndex)
				{
					foreach (i; existingIndex)
						static if (is(UpdateFuncRT!U == void))
							update(items[i].value);
						else
							items[i].value = update(items[i].value);
				});
		}
		else // ordered
		{
			lookup.updateVoid(key,
				delegate LookupValue ()
				{
					return [create()];
				},
				delegate void (ref LookupValue values)
				{
					foreach (ref value; values)
						static if (is(UpdateFuncRT!U == void))
							update(value);
						else
							value = update(value);
				});
		}
	}

	// *** Mutation (editing) ***

	static if (haveIndexing)
	{
		static if (haveValues)
		{
			/// Same as `set(k, v)`.
			ref IV opIndexAssign(AK, AV)(auto ref AV v, auto ref AK k)
			if (is(AK : K) && is(AV : V))
			{
				return set(k, v);
			}

			/// Perform cumulative operation with value
			/// (initialized with `.init` if the key does not exist).
			ref IV opIndexOpAssign(string op, AK, AV)(auto ref AV v, auto ref AK k)
			if (is(AK : K) && is(AV : V))
			{
				auto pv = &require(k);
				return mixin("(*pv) " ~ op ~ "= v");
			}

			/// Perform unary operation with value
			/// (initialized with `.init` if the key does not exist).
			ref IV opIndexUnary(string op, AK)(auto ref AK k)
			if (is(AK : K))
			{
				auto pv = &require(k);
				mixin("(*pv) " ~ op ~ ";");
				return *pv;
			}
		}
		else
		{
			private ref K editIndex(size_t index, scope void delegate(ref K) edit)
			{
				auto item = &items[index];
				K oldKey = item.key;
				auto pOldIndices = oldKey in lookup;
				assert(pOldIndices);

				edit(item.key);

				// Add new value

				lookup.updateVoid(item.key,
					delegate LookupValue()
					{
						// New value did not exist.
						if ((*pOldIndices).length == 1)
						{
							// Optimization - migrate the Indexes value
							assert((*pOldIndices)[0] == index);
							return *pOldIndices;
						}
						else
							return [index];
					},
					delegate void(ref LookupValue existingIndex)
					{
						// Value(s) with the new key already existed
						static if (multi)
							existingIndex ~= index;
						else
							assert(false, "Collision after in-place edit of a non-multi ordered set element");
					});

				// Remove old value

				if ((*pOldIndices).length == 1)
					lookup.remove(oldKey);
				else
				static if (multi)
					*pOldIndices = (*pOldIndices).remove!(i => i == index);
				else
					assert(false); // Should be unreachable (`if` above will always be true)

				return item.key;
			}

			/// Allows writing to ordered sets by index.
			/// The total number of elements never changes as a result
			/// of such an operation - a consequence of which is that
			/// if multi==false, changing the value to one that's
			/// already in the set is an error.
			ref IV opIndexAssign()(auto ref IV v, auto ref IK k)
			{
				static if (haveValues)
					return set(k, v);
				else
					return editIndex(k, (ref IV e) { e = v; });
			}

			/// Perform cumulative operation with value at index.
			ref IV opIndexOpAssign(string op)(auto ref VV v, auto ref IK k)
			{
				return editIndex(k, (ref IV e) { mixin("e " ~ op ~ "= v;"); });
			}

			/// Perform unary operation with value at index.
			ref IV opIndexUnary(string op)(auto ref IK k)
			{
				return editIndex(k, (ref IV e) { mixin("e " ~ op ~ ";"); });
			}
		}
	}

	// *** Mutation (removal) ***

	/// Removes all elements with the given key.
	bool remove(AK)(auto ref AK key)
	if (is(typeof(lookup.remove(key))))
	{
		static if (ordered)
		{
			auto p = key in lookup;
			if (!p)
				return false;

			auto targets = *p;
			foreach (target; targets)
			{
				items = items.remove!(SwapStrategy.stable)(target);
				foreach (ref k, ref vs; lookup)
					foreach (ref v; vs)
						if (v > target)
							v--;
			}
			auto success = lookup.remove(key);
			assert(success);
			return true;
		}
		else
			return lookup.remove(key);
	}

	/// Removes all elements.
	void clear()
	{
		lookup.clear();
		static if (ordered)
			items.length = 0;
	}
}

/// An associative array which retains the order in which elements were added.
alias OrderedMap(K, V) = HashCollection!(K, V, true, false);

unittest
{
	alias M = OrderedMap!(string, int);
	M m;
	m["a"] = 1;
	m["b"] = 2;
	m["c"] = 3;
	assert(m.length == 3);
	assert("a" in m);
	assert("d" !in m);

	{
		auto r = m.byKeyValue;
		assert(!r.empty);
		assert(r.front.key == "a");
		r.popFront();
		assert(!r.empty);
		assert(r.front.key == "b");
		r.popFront();
		assert(!r.empty);
		assert(r.front.key == "c");
		r.popFront();
		assert(r.empty);
	}

	assert(m.byKey.equal(["a", "b", "c"]));
	assert(m.byValue.equal([1, 2, 3]));
	assert(m.byKeyValue.map!(p => p.key).equal(m.byKey));
	assert(m.byKeyValue.map!(p => p.value).equal(m.byValue));
	assert(m.keys == ["a", "b", "c"]);
	assert(m.values == [1, 2, 3]);

	{
		const(M)* c = &m;
		assert(c.byKey.equal(["a", "b", "c"]));
		assert(c.byValue.equal([1, 2, 3]));
		assert(c.keys == ["a", "b", "c"]);
		assert(c.values == [1, 2, 3]);
	}

	m.byValue.front = 5;
	assert(m.byValue.equal([5, 2, 3]));

	m.remove("a");
	assert(m.length == 2);
	m["x"] -= 1;
	assert(m["x"] == -1);
	++m["y"];
	assert(m["y"] == 1);
	auto cm = cast(const)m.dup;
	foreach (k, v; cm)
		if (k == "x")
			assert(v == -1);
}

unittest
{
	OrderedMap!(string, int) m;
	m["a"] = 1;
	m["b"] = 2;
	m.remove("a");
	assert(m["b"] == 2);
}

unittest
{
	OrderedMap!(string, int) m;
	m["a"] = 1;
	auto m2 = m;
	m2.remove("a");
	m2["b"] = 2;
	assert(m["a"] == 1);
}

unittest
{
	OrderedMap!(string, int) m;
	m["a"] = 1;
	m["b"] = 2;
	auto m2 = m;
	m.remove("a");
	assert(m2["a"] == 1);
}

unittest
{
	class C {}
	const OrderedMap!(string, C) m;
	cast(void)m.byKeyValue;
}

unittest
{
	OrderedMap!(int, int) m;
	m.update(10,
		{ return 20; },
		(ref int k) { k++; return 30; },
	);
	assert(m.length == 1 && m[10] == 20);
	m.update(10,
		{ return 40; },
		(ref int k) { k++; return 50; },
	);
	assert(m.length == 1 && m[10] == 50);
}

// https://issues.dlang.org/show_bug.cgi?id=18606
unittest
{
	struct S
	{
		struct T
		{
			int foo;
			int[] bar;
		}

		OrderedMap!(int, T) m;
	}
}

unittest
{
	OrderedMap!(string, int) m;
	static assert(is(typeof(m.keys)));
	static assert(is(typeof(m.values)));
}

unittest
{
	OrderedMap!(string, int) m;
	foreach (k, v; m)
		k = k ~ k;
}

/// Like assocArray
auto orderedMap(R)(R input)
if (is(typeof(input.front.length) : size_t) && input.front.length == 2)
{
	alias K = typeof(input.front[0]);
	alias V = typeof(input.front[1]);
	return OrderedMap!(K, V)(input);
}

auto orderedMap(R)(R input) /// ditto
if (is(typeof(input.front.key)) && is(typeof(input.front.value)) && !is(typeof(input.front.length)))
{
	alias K = typeof(input.front.key);
	alias V = typeof(input.front.value);
	return OrderedMap!(K, V)(input);
}

unittest
{
	auto map = 3.iota.map!(n => tuple(n, n + 1)).orderedMap;
	assert(map.length == 3 && map[1] == 2);
}

unittest
{
	OrderedMap!(string, int) m;
	m = m.byKeyValue.orderedMap;
	m = m.byPair.orderedMap;
}

// ***************************************************************************

/// Helper/wrapper for void[0][T]
alias HashSet(T) = HashCollection!(T, void, false, false);

unittest
{
	HashSet!int s;
	assert(!s);
	assert(s.length == 0);
	assert(!(1 in s));
	assert(1 !in s);
	s.add(1);
	assert(1 in s);
	assert(s.length == 1);
	foreach (k; s)
		assert(k == 1);
	s.remove(1);
	assert(s.length == 0);

	s.add(1);
	auto t = s.dup;
	s.add(2);
	assert(t.length==1);
	t.remove(1);
	assert(t.length==0);
}

unittest
{
	struct S { int[int] aa; }
	HashSet!S set;
	S s;
	set.add(s);
	assert(s in set);
}

auto toSet(R)(R r)
{
	alias E = ElementType!R;
	return HashSet!E(r);
}

unittest
{
	auto set = [1, 2, 3].toSet();
	assert(2 in set);
	assert(4 !in set);
}

unittest
{
	HashSet!int m;
	const int i;
	m.remove(i);
}

unittest
{
	HashSet!Object m;
	Object o;
	m.remove(o);
}

// ***************************************************************************

alias OrderedSet(T) = HashCollection!(T, void, true, false);

unittest
{
	OrderedSet!int set;

	assert(1 !in set);
	set.add(1);
	assert(1 in set);
	set.remove(1);
	assert(1 !in set);

	set.add(1);
	set.clear();
	assert(1 !in set);

	set = set.init;
	assert(!set);
	set.add(1);
	assert(!!set);

	assert(set[0] == 1);
	set[0] = 2;
	assert(set[0] == 2);
	assert(1 !in set);
	assert(2 in set);

	assert(set.length == 1);
	set.remove(2);
	assert(set.length == 0);

	set.add(1);
	auto set2 = set;
	set.remove(1);
	set.add(2);
	assert(1 !in set && 2 in set);
	assert(1 in set2 && 2 !in set2);

	foreach (v; set)
		assert(v == 2);
}

auto orderedSet(R)(R r)
{
	alias E = ElementType!R;
	return OrderedSet!E(r);
}

// ***************************************************************************

/// An object which acts mostly as an associative array,
/// with the added property of being able to hold keys with
/// multiple values. These are only exposed explicitly and
/// through iteration
alias MultiAA(K, V) = HashCollection!(K, V, false, true);

unittest
{
	alias MASS = MultiAA!(string, int);
	MASS aa;
	aa.add("foo", 42);
	assert(aa["foo"] == 42);
	assert(aa.valuesOf("foo") == [42]);
	assert(aa.byPair.front.key == "foo");

	auto aa2 = MASS([tuple("foo", 42)]);
	aa2 = ["a":1,"b":2];

	const int i;
	aa["a"] = i;
}

unittest
{
	MultiAA!(int, int) m;
	int[][int] a;
	m = a;
}