Add back support for Cosmos nested dictionaries (#34312)

* Add back support for Cosmos nested dictionaries

Fixes #34105

This code will be consolidated with the relational UTF8 JSON code when #29825 is implemented. For now, just adding back what was already in Cosmos.

* Review updates

Author: ajcvickers

Diff for: src/EFCore.Cosmos/ChangeTracking/Internal/NullableStringDictionaryComparer.cs

@@ -19,11 +19,11 @@ public sealed class NullableStringDictionaryComparer<TElement, TCollection> : Va
     ///     any release. You should only use it directly in your code with extreme caution and knowing that
     ///     doing so can result in application failures when updating to a new Entity Framework Core release.
     /// </summary>
-    public NullableStringDictionaryComparer(ValueComparer elementComparer, bool readOnly)
+    public NullableStringDictionaryComparer(ValueComparer elementComparer)
         : base(
             (a, b) => Compare(a, b, (ValueComparer<TElement>)elementComparer),
             o => GetHashCode(o, (ValueComparer<TElement>)elementComparer),
-            source => Snapshot(source, (ValueComparer<TElement>)elementComparer, readOnly))
+            source => Snapshot(source, (ValueComparer<TElement>)elementComparer))
     {
     }
 
@@ -92,13 +92,8 @@ private static int GetHashCode(TCollection source, ValueComparer<TElement> eleme
         return hash.ToHashCode();
     }
 
-    private static TCollection Snapshot(TCollection source, ValueComparer<TElement> elementComparer, bool readOnly)
+    private static TCollection Snapshot(TCollection source, ValueComparer<TElement> elementComparer)
     {
-        if (readOnly)
-        {
-            return source;
-        }
-
         var snapshot = new Dictionary<string, TElement?>(((IReadOnlyDictionary<string, TElement?>)source).Count);
         foreach (var (key, element) in source)
         {

Diff for: src/EFCore.Cosmos/ChangeTracking/Internal/StringDictionaryComparer.cs

@@ -1,6 +1,9 @@
 // Licensed to the .NET Foundation under one or more agreements.
 // The .NET Foundation licenses this file to you under the MIT license.
 
+using System.Collections;
+using Microsoft.EntityFrameworkCore.Cosmos.Internal;
+
 namespace Microsoft.EntityFrameworkCore.Cosmos.ChangeTracking.Internal;
 
 /// <summary>
@@ -9,21 +12,30 @@ namespace Microsoft.EntityFrameworkCore.Cosmos.ChangeTracking.Internal;
 ///     any release. You should only use it directly in your code with extreme caution and knowing that
 ///     doing so can result in application failures when updating to a new Entity Framework Core release.
 /// </summary>
-public sealed class StringDictionaryComparer<TElement, TCollection> : ValueComparer<TCollection>
-    where TCollection : class, IEnumerable<KeyValuePair<string, TElement>>
+public sealed class StringDictionaryComparer<TDictionary, TElement> : ValueComparer<object>, IInfrastructure<ValueComparer>
 {
+    private static readonly MethodInfo CompareMethod = typeof(StringDictionaryComparer<TDictionary, TElement>).GetMethod(
+        nameof(Compare), BindingFlags.Static | BindingFlags.NonPublic, [typeof(object), typeof(object), typeof(ValueComparer)])!;
+
+    private static readonly MethodInfo GetHashCodeMethod = typeof(StringDictionaryComparer<TDictionary, TElement>).GetMethod(
+        nameof(GetHashCode), BindingFlags.Static | BindingFlags.NonPublic, [typeof(IEnumerable), typeof(ValueComparer)])!;
+
+    private static readonly MethodInfo SnapshotMethod = typeof(StringDictionaryComparer<TDictionary, TElement>).GetMethod(
+        nameof(Snapshot), BindingFlags.Static | BindingFlags.NonPublic, [typeof(object), typeof(ValueComparer)])!;
+
     /// <summary>
     ///     This is an internal API that supports the Entity Framework Core infrastructure and not subject to
     ///     the same compatibility standards as public APIs. It may be changed or removed without notice in
     ///     any release. You should only use it directly in your code with extreme caution and knowing that
     ///     doing so can result in application failures when updating to a new Entity Framework Core release.
     /// </summary>
-    public StringDictionaryComparer(ValueComparer elementComparer, bool readOnly)
+    public StringDictionaryComparer(ValueComparer elementComparer)
         : base(
-            (a, b) => Compare(a, b, (ValueComparer<TElement>)elementComparer),
-            o => GetHashCode(o, (ValueComparer<TElement>)elementComparer),
-            source => Snapshot(source, (ValueComparer<TElement>)elementComparer, readOnly))
+            CompareLambda(elementComparer),
+            GetHashCodeLambda(elementComparer),
+            SnapshotLambda(elementComparer))
     {
+        ElementComparer = elementComparer;
     }
 
     /// <summary>
@@ -32,63 +44,136 @@ public StringDictionaryComparer(ValueComparer elementComparer, bool readOnly)
     ///     any release. You should only use it directly in your code with extreme caution and knowing that
     ///     doing so can result in application failures when updating to a new Entity Framework Core release.
     /// </summary>
-    public override Type Type
-        => typeof(TCollection);
+    public ValueComparer ElementComparer { get; }
+
+    ValueComparer IInfrastructure<ValueComparer>.Instance => ElementComparer;
+
+    private static Expression<Func<object?, object?, bool>> CompareLambda(ValueComparer elementComparer)
+    {
+        var prm1 = Expression.Parameter(typeof(object), "a");
+        var prm2 = Expression.Parameter(typeof(object), "b");
+
+        return Expression.Lambda<Func<object?, object?, bool>>(
+            Expression.Call(
+                CompareMethod,
+                prm1,
+                prm2,
+#pragma warning disable EF9100
+                elementComparer.ConstructorExpression),
+#pragma warning restore EF9100
+            prm1,
+            prm2);
+    }
+
+    private static Expression<Func<object, int>> GetHashCodeLambda(ValueComparer elementComparer)
+    {
+        var prm = Expression.Parameter(typeof(object), "o");
+
+        return Expression.Lambda<Func<object, int>>(
+            Expression.Call(
+                GetHashCodeMethod,
+                Expression.Convert(
+                    prm,
+                    typeof(IEnumerable)),
+#pragma warning disable EF9100
+                elementComparer.ConstructorExpression),
+#pragma warning restore EF9100
+            prm);
+    }
+
+    private static Expression<Func<object, object>> SnapshotLambda(ValueComparer elementComparer)
+    {
+        var prm = Expression.Parameter(typeof(object), "source");
 
-    private static bool Compare(TCollection? a, TCollection? b, ValueComparer<TElement> elementComparer)
+        return Expression.Lambda<Func<object, object>>(
+            Expression.Call(
+                SnapshotMethod,
+                prm,
+#pragma warning disable EF9100
+                elementComparer.ConstructorExpression),
+#pragma warning restore EF9100
+            prm);
+    }
+
+    private static bool Compare(object? a, object? b, ValueComparer elementComparer)
     {
-        if (a is not IReadOnlyDictionary<string, TElement> aDict)
+        if (ReferenceEquals(a, b))
         {
-            return b is not IReadOnlyDictionary<string, TElement>;
+            return true;
         }
 
-        if (b is not IReadOnlyDictionary<string, TElement> bDict || aDict.Count != bDict.Count)
+        if (a is null)
         {
-            return false;
+            return b is null;
         }
 
-        if (ReferenceEquals(aDict, bDict))
+        if (b is null)
         {
-            return true;
+            return false;
         }
 
-        foreach (var (key, element) in aDict)
+        if (a is IReadOnlyDictionary<string, TElement?> aDictionary && b is IReadOnlyDictionary<string, TElement?> bDictionary)
         {
-            if (!bDict.TryGetValue(key, out var bValue)
-                || !elementComparer.Equals(element, bValue))
+            if (aDictionary.Count != bDictionary.Count)
             {
                 return false;
             }
+
+            foreach (var pair in aDictionary)
+            {
+                if (!bDictionary.TryGetValue(pair.Key, out var bValue)
+                    || !elementComparer.Equals(pair.Value, bValue))
+                {
+                    return false;
+                }
+            }
+
+            return true;
         }
 
-        return true;
+        throw new InvalidOperationException(
+            CosmosStrings.BadDictionaryType(
+                (a is IDictionary<string, TElement?> ? b : a).GetType().ShortDisplayName(),
+                typeof(IDictionary<,>).MakeGenericType(typeof(string), elementComparer.Type).ShortDisplayName()));
     }
 
-    private static int GetHashCode(TCollection source, ValueComparer<TElement> elementComparer)
+    private static int GetHashCode(IEnumerable source, ValueComparer elementComparer)
     {
+        if (source is not IReadOnlyDictionary<string, TElement?> sourceDictionary)
+        {
+            throw new InvalidOperationException(
+                CosmosStrings.BadDictionaryType(
+                    source.GetType().ShortDisplayName(),
+                    typeof(IList<>).MakeGenericType(elementComparer.Type).ShortDisplayName()));
+        }
+
         var hash = new HashCode();
-        foreach (var (key, element) in source)
+
+        foreach (var pair in sourceDictionary)
         {
-            hash.Add(key);
-            hash.Add(element, elementComparer);
+            hash.Add(pair.Key);
+            hash.Add(pair.Value == null ? 0 : elementComparer.GetHashCode(pair.Value));
         }
 
         return hash.ToHashCode();
     }
 
-    private static TCollection Snapshot(TCollection source, ValueComparer<TElement> elementComparer, bool readOnly)
+    private static IReadOnlyDictionary<string, TElement?> Snapshot(object source, ValueComparer elementComparer)
     {
-        if (readOnly)
+        if (source is not IReadOnlyDictionary<string, TElement?> sourceDictionary)
         {
-            return source;
+            throw new InvalidOperationException(
+                CosmosStrings.BadDictionaryType(
+                    source.GetType().ShortDisplayName(),
+                    typeof(IDictionary<,>).MakeGenericType(typeof(string), elementComparer.Type).ShortDisplayName()));
         }
 
-        var snapshot = new Dictionary<string, TElement>(((IReadOnlyDictionary<string, TElement>)source).Count);
-        foreach (var (key, element) in source)
+        var snapshot = new Dictionary<string, TElement?>();
+        foreach (var pair in sourceDictionary)
         {
-            snapshot.Add(key, element is null ? default! : elementComparer.Snapshot(element));
+            snapshot[pair.Key] = pair.Value == null ? default : (TElement?)elementComparer.Snapshot(pair.Value);
         }
 
-        return (TCollection)(object)snapshot;
+        return snapshot;
     }
 }

Diff for: src/EFCore.Cosmos/Properties/CosmosStrings.Designer.cs

@@ -120,6 +120,9 @@
   <data name="AnalyticalTTLMismatch" xml:space="preserve">
     <value>The time to live for analytical store was configured to '{ttl1}' on '{entityType1}', but on '{entityType2}' it was configured to '{ttl2}'. All entity types mapped to the same container '{container}' must be configured with the same time to live for analytical store.</value>
   </data>
+  <data name="BadDictionaryType" xml:space="preserve">
+    <value>The type '{givenType}' cannot be mapped as a dictionary because it does not implement '{dictionaryType}'.</value>
+  </data>
   <data name="CanConnectNotSupported" xml:space="preserve">
     <value>The Cosmos database does not support 'CanConnect' or 'CanConnectAsync'.</value>
   </data>

Diff for: src/EFCore.Cosmos/Properties/CosmosStrings.resx

@@ -1,7 +1,11 @@
 // Licensed to the .NET Foundation under one or more agreements.
 // The .NET Foundation licenses this file to you under the MIT license.
 
+using System.Diagnostics.CodeAnalysis;
+using System.Text.Json;
 using Microsoft.EntityFrameworkCore.Cosmos.ChangeTracking.Internal;
+using Microsoft.EntityFrameworkCore.Storage.Internal;
+using Microsoft.EntityFrameworkCore.Storage.Json;
 using Newtonsoft.Json.Linq;
 
 namespace Microsoft.EntityFrameworkCore.Cosmos.Storage.Internal;
@@ -103,11 +107,12 @@ public CosmosTypeMappingSource(TypeMappingSourceDependencies dependencies)
             return null;
         }
 
-        var jsonValueReaderWriter = Dependencies.JsonValueReaderWriterSource.FindReaderWriter(clrType);
-
         if (clrType is { IsGenericType: true, IsGenericTypeDefinition: false })
         {
             var genericTypeDefinition = clrType.GetGenericTypeDefinition();
+
+            // This is legacy type mapping support for dictionaries in Cosmos. This needs to be consolidated with the relational
+            // support, but for now this is being added back in to avoid a regression in EF9.
             if (genericTypeDefinition == typeof(Dictionary<,>)
                 || genericTypeDefinition == typeof(IDictionary<,>)
                 || genericTypeDefinition == typeof(IReadOnlyDictionary<,>))
@@ -122,11 +127,24 @@ public CosmosTypeMappingSource(TypeMappingSourceDependencies dependencies)
                 var elementMappingInfo = new TypeMappingInfo(elementType);
                 elementMapping = FindPrimitiveMapping(elementMappingInfo)
                     ?? FindCollectionMapping(elementMappingInfo);
-                return elementMapping == null
-                    ? null
-                    : new CosmosTypeMapping(
-                        clrType, CreateStringDictionaryComparer(elementMapping, elementType, clrType),
+
+                if (elementMapping != null)
+                {
+                    var jsonValueReaderWriter = Dependencies.JsonValueReaderWriterSource.FindReaderWriter(clrType);
+                    if (jsonValueReaderWriter == null
+                        && elementMapping.JsonValueReaderWriter != null)
+                    {
+                        jsonValueReaderWriter = (JsonValueReaderWriter?)Activator.CreateInstance(
+                            typeof(PlaceholderJsonStringKeyedDictionaryReaderWriter<>)
+                                .MakeGenericType(elementMapping.JsonValueReaderWriter.ValueType),
+                            elementMapping.JsonValueReaderWriter);
+                    }
+
+                    return new CosmosTypeMapping(
+                        clrType,
+                        CreateStringDictionaryComparer(elementMapping, elementType, clrType),
                         jsonValueReaderWriter: jsonValueReaderWriter);
+                }
             }
         }
 
@@ -143,9 +161,59 @@ private static ValueComparer CreateStringDictionaryComparer(
 
         return (ValueComparer)Activator.CreateInstance(
             elementType == unwrappedType
-                ? typeof(StringDictionaryComparer<,>).MakeGenericType(elementType, dictType)
+                ? typeof(StringDictionaryComparer<,>).MakeGenericType(dictType, elementType)
                 : typeof(NullableStringDictionaryComparer<,>).MakeGenericType(unwrappedType, dictType),
-            elementMapping.Comparer,
-            readOnly)!;
+            elementMapping.Comparer)!;
+    }
+
+    // This ensures that the element reader/writers are not null when using Cosmos dictionary type mappings, but
+    // is never actually used because Cosmos does not (yet) read and write JSON using this mechanism.
+    /// <summary>
+    ///     This is an internal API that supports the Entity Framework Core infrastructure and not subject to
+    ///     the same compatibility standards as public APIs. It may be changed or removed without notice in
+    ///     any release. You should only use it directly in your code with extreme caution and knowing that
+    ///     doing so can result in application failures when updating to a new Entity Framework Core release.
+    /// </summary>
+#pragma warning disable EF1001
+    public sealed class PlaceholderJsonStringKeyedDictionaryReaderWriter<TElement>(JsonValueReaderWriter elementReaderWriter)
+        : JsonValueReaderWriter<IEnumerable<KeyValuePair<string, TElement>>>, ICompositeJsonValueReaderWriter
+#pragma warning restore EF1001
+    {
+        private readonly JsonValueReaderWriter<TElement> _elementReaderWriter = (JsonValueReaderWriter<TElement>)elementReaderWriter;
+
+        /// <summary>
+        ///     This is an internal API that supports the Entity Framework Core infrastructure and not subject to
+        ///     the same compatibility standards as public APIs. It may be changed or removed without notice in
+        ///     any release. You should only use it directly in your code with extreme caution and knowing that
+        ///     doing so can result in application failures when updating to a new Entity Framework Core release.
+        /// </summary>
+        public override IEnumerable<KeyValuePair<string, TElement>> FromJsonTyped(
+            ref Utf8JsonReaderManager manager,
+            object? existingObject = null)
+            => throw new NotImplementedException("JsonValueReaderWriter infrastructure is not supported on Cosmos.");
+
+        /// <summary>
+        ///     This is an internal API that supports the Entity Framework Core infrastructure and not subject to
+        ///     the same compatibility standards as public APIs. It may be changed or removed without notice in
+        ///     any release. You should only use it directly in your code with extreme caution and knowing that
+        ///     doing so can result in application failures when updating to a new Entity Framework Core release.
+        /// </summary>
+        public override void ToJsonTyped(Utf8JsonWriter writer, IEnumerable<KeyValuePair<string, TElement>> value)
+            => throw new NotImplementedException("JsonValueReaderWriter infrastructure is not supported on Cosmos.");
+
+        JsonValueReaderWriter ICompositeJsonValueReaderWriter.InnerReaderWriter
+            => _elementReaderWriter;
+
+        private readonly ConstructorInfo _constructorInfo
+            = typeof(PlaceholderJsonStringKeyedDictionaryReaderWriter<TElement>)
+                .GetConstructor([typeof(JsonValueReaderWriter<TElement>)])!;
+
+        /// <inheritdoc />
+        public override Expression ConstructorExpression
+#pragma warning disable EF9100
+#pragma warning disable EF1001
+            => Expression.New(_constructorInfo, ((ICompositeJsonValueReaderWriter)this).InnerReaderWriter.ConstructorExpression);
+#pragma warning restore EF1001
+#pragma warning restore EF9100
     }
 }