From 9006bdffa238876fa5e21acac4c414b7d1d074a0 Mon Sep 17 00:00:00 2001
From: Oliver Booth <me@olivr.me>
Date: Fri, 29 Apr 2022 18:10:41 +0100
Subject: [PATCH] Add IDictionary.AddOrUpdate

---
 X10D.Tests/src/Collections/DictionaryTests.cs |  90 ++++++++++
 X10D/src/Collections/DictionaryExtensions.cs  | 164 ++++++++++++++++++
 2 files changed, 254 insertions(+)
 create mode 100644 X10D.Tests/src/Collections/DictionaryTests.cs

diff --git a/X10D.Tests/src/Collections/DictionaryTests.cs b/X10D.Tests/src/Collections/DictionaryTests.cs
new file mode 100644
index 0000000..b6b33fc
--- /dev/null
+++ b/X10D.Tests/src/Collections/DictionaryTests.cs
@@ -0,0 +1,90 @@
+using Microsoft.VisualStudio.TestTools.UnitTesting;
+using X10D.Collections;
+
+namespace X10D.Tests.Collections;
+
+[TestClass]
+public class DictionaryTests
+{
+    [TestMethod]
+    public void AddOrUpdate_ShouldAddNewKey_IfNotExists()
+    {
+        var dictionary = new Dictionary<int, string>();
+        Assert.IsFalse(dictionary.ContainsKey(1));
+
+        dictionary.AddOrUpdate(1, "one", (_, _) => string.Empty);
+        Assert.IsTrue(dictionary.ContainsKey(1));
+        Assert.AreEqual("one", dictionary[1]);
+
+        dictionary.Clear();
+        Assert.IsFalse(dictionary.ContainsKey(1));
+
+        dictionary.AddOrUpdate(1, _ => "one", (_, _) => string.Empty);
+        Assert.IsTrue(dictionary.ContainsKey(1));
+        Assert.AreEqual("one", dictionary[1]);
+
+        dictionary.Clear();
+        Assert.IsFalse(dictionary.ContainsKey(1));
+
+        dictionary.AddOrUpdate(1, (_, _) => "one", (_, _, _) => string.Empty, 0);
+        Assert.IsTrue(dictionary.ContainsKey(1));
+        Assert.AreEqual("one", dictionary[1]);
+    }
+
+    [TestMethod]
+    public void AddOrUpdate_ShouldUpdateKey_IfExists()
+    {
+        var dictionary = new Dictionary<int, string> {[1] = "one"};
+        Assert.IsTrue(dictionary.ContainsKey(1));
+        Assert.AreEqual("one", dictionary[1]);
+
+        dictionary.AddOrUpdate(1, string.Empty, (_, _) => "two");
+        Assert.IsTrue(dictionary.ContainsKey(1));
+        Assert.AreEqual("two", dictionary[1]);
+
+        dictionary.Clear();
+        Assert.IsFalse(dictionary.ContainsKey(1));
+        dictionary[1] = "one";
+
+        dictionary.AddOrUpdate(1, _ => string.Empty, (_, _) => "two");
+        Assert.IsTrue(dictionary.ContainsKey(1));
+        Assert.AreEqual("two", dictionary[1]);
+
+        dictionary.Clear();
+        Assert.IsFalse(dictionary.ContainsKey(1));
+        dictionary[1] = "one";
+
+        dictionary.AddOrUpdate(1, (_, _) => string.Empty, (_, _, _) => "two", 0);
+        Assert.IsTrue(dictionary.ContainsKey(1));
+        Assert.AreEqual("two", dictionary[1]);
+    }
+
+    [TestMethod]
+    public void AddOrUpdate_ShouldThrow_GivenNullDictionary()
+    {
+        Dictionary<int, string>? dictionary = null;
+        Assert.ThrowsException<ArgumentNullException>(() => dictionary!.AddOrUpdate(1, string.Empty, (_, _) => string.Empty));
+        Assert.ThrowsException<ArgumentNullException>(() =>
+            dictionary!.AddOrUpdate(1, _ => string.Empty, (_, _) => string.Empty));
+        Assert.ThrowsException<ArgumentNullException>(() =>
+            dictionary!.AddOrUpdate(1, (_, _) => string.Empty, (_, _, _) => string.Empty, 0));
+    }
+
+    [TestMethod]
+    public void AddOrUpdate_ShouldThrow_GivenNullUpdateValueFactory()
+    {
+        var dictionary = new Dictionary<int, string>();
+        Assert.ThrowsException<ArgumentNullException>(() => dictionary.AddOrUpdate(1, string.Empty, null!));
+        Assert.ThrowsException<ArgumentNullException>(() => dictionary.AddOrUpdate(1, _ => string.Empty, null!));
+        Assert.ThrowsException<ArgumentNullException>(() => dictionary.AddOrUpdate(1, (_, _) => string.Empty, null!, 0));
+    }
+
+    [TestMethod]
+    public void AddOrUpdate_ShouldThrow_GivenNullAddValueFactory()
+    {
+        var dictionary = new Dictionary<int, string>();
+        Func<int, string>? addValueFactory = null;
+        Assert.ThrowsException<ArgumentNullException>(() => dictionary.AddOrUpdate(1, addValueFactory!, (_, _) => "one"));
+        Assert.ThrowsException<ArgumentNullException>(() => dictionary.AddOrUpdate(1, null!, (_, _, _) => "one", 0));
+    }
+}
diff --git a/X10D/src/Collections/DictionaryExtensions.cs b/X10D/src/Collections/DictionaryExtensions.cs
index 80d2811..498b277 100644
--- a/X10D/src/Collections/DictionaryExtensions.cs
+++ b/X10D/src/Collections/DictionaryExtensions.cs
@@ -7,6 +7,170 @@ namespace X10D.Collections;
 /// </summary>
 public static class DictionaryExtensions
 {
+    /// <summary>
+    ///     Adds a key/value pair to the <see cref="IDictionary{TKey,TValue}" /> if the key does not already exist, or updates a
+    ///     key/value pair in the <see cref="IDictionary{TKey,TValue}" /> by using the specified function if the key already
+    ///     exists.
+    /// </summary>
+    /// <param name="dictionary">The dictionary to update.</param>
+    /// <param name="key">The key to be added or whose value should be updated.</param>
+    /// <param name="addValue">The value to be added for an absent key.</param>
+    /// <param name="updateValueFactory">
+    ///     The function used to generate a new value for an existing key based on the key's existing value.
+    /// </param>
+    /// <typeparam name="TKey">The type of the keys in the dictionary.</typeparam>
+    /// <typeparam name="TValue">The type of the values in the dictionary.</typeparam>
+    /// <returns>
+    ///     The new value for the key. This will be either be <paramref name="addValue" /> (if the key was absent) or the result
+    ///     of <paramref name="updateValueFactory" /> (if the key was present).
+    /// </returns>
+    /// <exception cref="ArgumentNullException">
+    ///     <para><paramref name="dictionary" /> is <see langword="null" />.</para>
+    ///     -or-
+    ///     <para><paramref name="updateValueFactory" /> is <see langword="null" />.</para>
+    /// </exception>
+    public static TValue AddOrUpdate<TKey, TValue>(this IDictionary<TKey, TValue> dictionary, TKey key, TValue addValue,
+        Func<TKey, TValue, TValue> updateValueFactory)
+        where TKey : notnull
+    {
+        if (dictionary is null)
+        {
+            throw new ArgumentNullException(nameof(dictionary));
+        }
+
+        if (updateValueFactory is null)
+        {
+            throw new ArgumentNullException(nameof(updateValueFactory));
+        }
+
+        if (dictionary.ContainsKey(key))
+        {
+            dictionary[key] = updateValueFactory(key, dictionary[key]);
+        }
+        else
+        {
+            dictionary.Add(key, addValue);
+        }
+
+        return dictionary[key];
+    }
+
+    /// <summary>
+    ///     Uses the specified functions to add a key/value pair to the <see cref="IDictionary{TKey,TValue}" /> if the key does
+    ///     not already exist, or to update a key/value pair in the <see cref="IDictionary{TKey,TValue}" /> if the key already
+    ///     exists.
+    /// </summary>
+    /// <param name="dictionary">The dictionary to update.</param>
+    /// <param name="key">The key to be added or whose value should be updated.</param>
+    /// <param name="addValueFactory">The function used to generate a value for an absent key.</param>
+    /// <param name="updateValueFactory">
+    ///     The function used to generate a new value for an existing key based on the key's existing value.
+    /// </param>
+    /// <typeparam name="TKey">The type of the keys in the dictionary.</typeparam>
+    /// <typeparam name="TValue">The type of the values in the dictionary.</typeparam>
+    /// <returns>
+    ///     The new value for the key. This will be either be the result of <paramref name="addValueFactory "/> (if the key was
+    ///     absent) or the result of <paramref name="updateValueFactory" /> (if the key was present).
+    /// </returns>
+    /// <exception cref="ArgumentNullException">
+    ///     <para><paramref name="dictionary" /> is <see langword="null" />.</para>
+    ///     -or-
+    ///     <para><paramref name="addValueFactory" /> is <see langword="null" />.</para>
+    ///     -or-
+    ///     <para><paramref name="updateValueFactory" /> is <see langword="null" />.</para>
+    /// </exception>
+    public static TValue AddOrUpdate<TKey, TValue>(this IDictionary<TKey, TValue> dictionary, TKey key,
+        Func<TKey, TValue> addValueFactory, Func<TKey, TValue, TValue> updateValueFactory)
+        where TKey : notnull
+    {
+        if (dictionary is null)
+        {
+            throw new ArgumentNullException(nameof(dictionary));
+        }
+
+        if (addValueFactory is null)
+        {
+            throw new ArgumentNullException(nameof(addValueFactory));
+        }
+
+        if (updateValueFactory is null)
+        {
+            throw new ArgumentNullException(nameof(updateValueFactory));
+        }
+
+        if (dictionary.ContainsKey(key))
+        {
+            dictionary[key] = updateValueFactory(key, dictionary[key]);
+        }
+        else
+        {
+            dictionary.Add(key, addValueFactory(key));
+        }
+
+        return dictionary[key];
+    }
+
+    /// <summary>
+    ///     Uses the specified functions and argument to add a key/value pair to the <see cref="IDictionary{TKey,TValue}" /> if
+    ///     the key does not already exist, or to update a key/value pair in the <see cref="IDictionary{TKey,TValue}" /> if th
+    ///     key already exists.
+    /// </summary>
+    /// <param name="dictionary">The dictionary to update.</param>
+    /// <param name="key">The key to be added or whose value should be updated.</param>
+    /// <param name="addValueFactory">The function used to generate a value for an absent key.</param>
+    /// <param name="updateValueFactory">
+    ///     The function used to generate a new value for an existing key based on the key's existing value.
+    /// </param>
+    /// <param name="factoryArgument">
+    ///     An argument to pass into <paramref name="addValueFactory" /> and <paramref name="updateValueFactory" />.
+    /// </param>
+    /// <typeparam name="TKey">The type of the keys in the dictionary.</typeparam>
+    /// <typeparam name="TValue">The type of the values in the dictionary.</typeparam>
+    /// <typeparam name="TArg">
+    ///     The type of an argument to pass into <paramref name="addValueFactory" /> and <paramref name="updateValueFactory" />.
+    /// </typeparam>
+    /// <returns>
+    ///     The new value for the key. This will be either be the result of <paramref name="addValueFactory "/> (if the key was
+    ///     absent) or the result of <paramref name="updateValueFactory" /> (if the key was present).
+    /// </returns>
+    /// <exception cref="ArgumentNullException">
+    ///     <para><paramref name="dictionary" /> is <see langword="null" />.</para>
+    ///     -or-
+    ///     <para><paramref name="addValueFactory" /> is <see langword="null" />.</para>
+    ///     -or-
+    ///     <para><paramref name="updateValueFactory" /> is <see langword="null" />.</para>
+    /// </exception>
+    public static TValue AddOrUpdate<TKey, TValue, TArg>(this IDictionary<TKey, TValue> dictionary, TKey key,
+        Func<TKey, TArg, TValue> addValueFactory, Func<TKey, TValue, TArg, TValue> updateValueFactory, TArg factoryArgument)
+        where TKey : notnull
+    {
+        if (dictionary is null)
+        {
+            throw new ArgumentNullException(nameof(dictionary));
+        }
+
+        if (addValueFactory is null)
+        {
+            throw new ArgumentNullException(nameof(addValueFactory));
+        }
+
+        if (updateValueFactory is null)
+        {
+            throw new ArgumentNullException(nameof(updateValueFactory));
+        }
+
+        if (dictionary.ContainsKey(key))
+        {
+            dictionary[key] = updateValueFactory(key, dictionary[key], factoryArgument);
+        }
+        else
+        {
+            dictionary.Add(key, addValueFactory(key, factoryArgument));
+        }
+
+        return dictionary[key];
+    }
+
     /// <summary>
     ///     Converts an <see cref="IEnumerable{T}" /> of <see cref="KeyValuePair{TKey, TValue}" /> to a data connection
     ///     string.