CSHARP-4443: Add comprehensive dictionary LINQ support for all 3 representations

src/MongoDB.Bson/Serialization/Serializers/KeyValuePairSerializer.cs

@@ -29,6 +29,22 @@ public interface IKeyValuePairSerializer
         BsonType Representation { get;  }
     }
 
+    /// <summary>
+    /// An extended interface for KeyValuePairSerializer that provides access to key and value serializers.
+    /// </summary>
+    public interface IKeyValuePairSerializerV2 : IKeyValuePairSerializer
+    {
+        /// <summary>
+        /// Gets the key serializer.
+        /// </summary>
+        IBsonSerializer KeySerializer { get; }
+
+        /// <summary>
+        /// Gets the value serializer.
+        /// </summary>
+        IBsonSerializer ValueSerializer { get; }
+    }
+
     /// <summary>
     /// Static factory class for KeyValuePairSerializers.
     /// </summary>
@@ -61,7 +77,7 @@ public static IBsonSerializer Create(
     public sealed class KeyValuePairSerializer<TKey, TValue> :
         StructSerializerBase<KeyValuePair<TKey, TValue>>,
         IBsonDocumentSerializer,
-        IKeyValuePairSerializer
+        IKeyValuePairSerializerV2
     {
         // private constants
         private static class Flags
@@ -191,6 +207,16 @@ public IBsonSerializer<TValue> ValueSerializer
             get { return _lazyValueSerializer.Value; }
         }
 
+        /// <summary>
+        /// Gets the key serializer.
+        /// </summary>
+        IBsonSerializer IKeyValuePairSerializerV2.KeySerializer => KeySerializer;
+
+        /// <summary>
+        /// Gets the value serializer.
+        /// </summary>
+        IBsonSerializer IKeyValuePairSerializerV2.ValueSerializer => ValueSerializer;
+
         // public methods
         /// <summary>
         /// Deserializes a value.

src/MongoDB.Driver/Linq/Linq3Implementation/Ast/Expressions/AstExpression.cs

@@ -134,6 +134,11 @@ public static AstExpression ArrayElemAt(AstExpression array, AstExpression index
             return new AstBinaryExpression(AstBinaryOperator.ArrayElemAt, array, index);
         }
 
+        public static AstExpression ArrayToObject(AstExpression arg)
+        {
+            return new AstUnaryExpression(AstUnaryOperator.ArrayToObject, arg);
+        }
+
         public static AstExpression Avg(AstExpression array)
         {
             return new AstUnaryExpression(AstUnaryOperator.Avg, array);

src/MongoDB.Driver/Linq/Linq3Implementation/Ast/Optimizers/AstSimplifier.cs

@@ -452,34 +452,36 @@ public override AstNode VisitMapExpression(AstMapExpression node)
 
             // { $map : { input : { $map : { input : <innerInput>, as : "inner", in : { A : <exprA>, B : <exprB>, ... } } }, as: "outer", in : { F : '$$outer.A', G : "$$outer.B", ... } } }
             // => { $map : { input : <innerInput>, as: "inner", in : { F : <exprA>, G : <exprB>, ... } } }
-            if (node.Input is AstMapExpression innerMapExpression &&
-                node.As is var outerVar &&
-                node.In is AstComputedDocumentExpression outerComputedDocumentExpression &&
-                innerMapExpression.Input is var innerInput &&
-                innerMapExpression.As is var innerVar &&
-                innerMapExpression.In is AstComputedDocumentExpression innerComputedDocumentExpression &&
-                outerComputedDocumentExpression.Fields.All(outerField =>
-                    outerField.Value is AstGetFieldExpression outerGetFieldExpression &&
-                    outerGetFieldExpression.Input == outerVar &&
-                    outerGetFieldExpression.FieldName is AstConstantExpression { Value : BsonString { Value : var matchingFieldName } } &&
-                    innerComputedDocumentExpression.Fields.Any(innerField => innerField.Path == matchingFieldName)))
             {
-                var rewrittenOuterFields = new List<AstComputedField>();
-                foreach (var outerField in outerComputedDocumentExpression.Fields)
-                {
-                    var outerGetFieldExpression = (AstGetFieldExpression)outerField.Value;
-                    var matchingFieldName = ((AstConstantExpression)outerGetFieldExpression.FieldName).Value.AsString;
-                    var matchingInnerField = innerComputedDocumentExpression.Fields.Single(innerField => innerField.Path == matchingFieldName);
-                    var rewrittenOuterField = AstExpression.ComputedField(outerField.Path, matchingInnerField.Value);
-                    rewrittenOuterFields.Add(rewrittenOuterField);
-                }
+                if (node.Input is AstMapExpression innerMapExpression &&
+                    node.As is var outerVar &&
+                    node.In is AstComputedDocumentExpression outerComputedDocumentExpression &&
+                    innerMapExpression.Input is var innerInput &&
+                    innerMapExpression.As is var innerVar &&
+                    innerMapExpression.In is AstComputedDocumentExpression innerComputedDocumentExpression &&
+                    outerComputedDocumentExpression.Fields.All(outerField =>
+                        outerField.Value is AstGetFieldExpression outerGetFieldExpression &&
+                        outerGetFieldExpression.Input == outerVar &&
+                        outerGetFieldExpression.FieldName is AstConstantExpression { Value : BsonString { Value : var matchingFieldName } } &&
+                        innerComputedDocumentExpression.Fields.Any(innerField => innerField.Path == matchingFieldName)))
+                {
+                    var rewrittenOuterFields = new List<AstComputedField>();
+                    foreach (var outerField in outerComputedDocumentExpression.Fields)
+                    {
+                        var outerGetFieldExpression = (AstGetFieldExpression)outerField.Value;
+                        var matchingFieldName = ((AstConstantExpression)outerGetFieldExpression.FieldName).Value.AsString;
+                        var matchingInnerField = innerComputedDocumentExpression.Fields.Single(innerField => innerField.Path == matchingFieldName);
+                        var rewrittenOuterField = AstExpression.ComputedField(outerField.Path, matchingInnerField.Value);
+                        rewrittenOuterFields.Add(rewrittenOuterField);
+                    }
 
-                var simplified = AstExpression.Map(
-                    input: innerInput,
-                    @as: innerVar,
-                    @in: AstExpression.ComputedDocument(rewrittenOuterFields));
+                    var simplified = AstExpression.Map(
+                        input: innerInput,
+                        @as: innerVar,
+                        @in: AstExpression.ComputedDocument(rewrittenOuterFields));
 
-                return Visit(simplified);
+                    return Visit(simplified);
+                }
             }
 
             // { $map : { input : [{ A : <exprA1>, B : <exprB1>, ... }, { A : <exprA2>, B : <exprB2>, ... }, ...], as : "item", in: { F : "$$item.A", G : "$$item.B", ... } } }
@@ -523,6 +525,32 @@ mappedField.Value is AstGetFieldExpression mappedGetField &&
                 return Visit(simplified);
             }
 
+            // { $map : { input : { $map : { input : <array>Expr, as : <innerVar>, in : { ..., fieldName : <fieldExpr>, ... } } }, as : <outerVar>, in : "$$<outerVar>.fieldName" } }
+            // => { $map : { input : <arrayExpr>, as : <innerVar>, in : <fieldExpr> } }
+            {
+                if (node.Input is AstMapExpression innerMap &&
+                    node.As is var outerVar &&
+                    node.In is AstGetFieldExpression outerGetFieldExpression &&
+                    outerGetFieldExpression.Input == outerVar &&
+                    outerGetFieldExpression.FieldName is AstConstantExpression { Value : { BsonType : BsonType.String } } fieldNameExpression &&
+                    fieldNameExpression.Value.AsString is var fieldName &&
+                    innerMap.Input is var arrayExpression &&
+                    innerMap.As is var innerVar &&
+                    innerMap.In is AstComputedDocumentExpression innerComputedDocument &&
+                    innerComputedDocument.Fields.SingleOrDefault(f => f.Path == fieldName) is var computedField &&
+                    computedField != null &&
+                    computedField.Value is var fieldExpr)
+                {
+                    var simplified =
+                        AstExpression.Map(
+                            input: arrayExpression,
+                            @as: innerVar,
+                            @in: fieldExpr);
+
+                    return Visit(simplified);
+                }
+            }
+
             return base.VisitMapExpression(node);
 
             static AstExpression UltimateGetFieldInput(AstGetFieldExpression getField)
@@ -656,6 +684,14 @@ item is AstComputedDocumentExpression computedDocumentExpression &&
                 return AstExpression.ComputedDocument(computedFields);
             }
 
+            // { $size : { $map : { input : <array>, as : <var>, in : <expr> } } } => { $size : <array> }
+            // $map is always 1-to-1, so size of the mapped array equals size of the input array
+            if (node.Operator == AstUnaryOperator.Size &&
+                arg is AstMapExpression mapExpression)
+            {
+                return AstExpression.Size(mapExpression.Input);
+            }
+
             return node.Update(arg);
 
             static AstComputedField KeyValuePairDocumentToComputedField(AstExpression expression)

src/MongoDB.Driver/Linq/Linq3Implementation/Translators/ExpressionToAggregationExpressionTranslators/ExpressionToAggregationExpressionTranslator.cs

@@ -126,6 +126,21 @@ public static TranslatedExpression TranslateEnumerable(TranslationContext contex
                 aggregateExpression = new TranslatedExpression(expression, ast, arrayOfDocumentsDictionarySerializer);
             }
 
+            if (aggregateExpression.Serializer is IBsonArraySerializer arraySerializer &&
+                arraySerializer.TryGetItemSerializationInfo(out var itemSerializationInfo) &&
+                itemSerializationInfo.Serializer is IWrappedValueSerializer wrappedItemSerializer)
+            {
+                var itemVar = AstExpression.Var("item");
+                var unwrappedItemsAst = AstExpression.Map(
+                    input: aggregateExpression.Ast,
+                    @as: itemVar,
+                    @in: AstExpression.GetField(itemVar, wrappedItemSerializer.FieldName));
+                var unwrappedItemSerializer = wrappedItemSerializer.ValueSerializer;
+                var unwrappedItemsSerializer = ArraySerializerHelper.CreateSerializer(unwrappedItemSerializer);
+
+                aggregateExpression = new TranslatedExpression(expression, unwrappedItemsAst, unwrappedItemsSerializer);
+            }
+
             return aggregateExpression;
         }
 

src/MongoDB.Driver/Linq/Linq3Implementation/Translators/ExpressionToAggregationExpressionTranslators/MemberExpressionToAggregationExpressionTranslator.cs

@@ -14,7 +14,6 @@
 */
 
 using System;
-using System.Collections;
 using System.Collections.Generic;
 using System.Linq.Expressions;
 using System.Reflection;
@@ -71,11 +70,16 @@ public static TranslatedExpression Translate(TranslationContext context, MemberE
 
             if (!DocumentSerializerHelper.AreMembersRepresentedAsFields(containerTranslation.Serializer, out _))
             {
-                if (member is PropertyInfo propertyInfo  && propertyInfo.Name == "Length")
+                if (member is PropertyInfo propertyInfo && propertyInfo.Name == "Length")
                 {
                     return LengthPropertyToAggregationExpressionTranslator.Translate(context, expression);
                 }
 
+                if (TryTranslateKeyValuePairProperty(expression, containerTranslation, member, out var translatedKeyValuePairProperty))
+                {
+                    return translatedKeyValuePairProperty;
+                }
+
                 if (TryTranslateCollectionCountProperty(expression, containerTranslation, member, out var translatedCount))
                 {
                     return translatedCount;
@@ -126,11 +130,20 @@ private static bool TryTranslateCollectionCountProperty(MemberExpression express
         {
             if (EnumerableProperty.IsCountProperty(expression))
             {
-                SerializationHelper.EnsureRepresentationIsArray(expression, container.Serializer);
+                AstExpression ast;
 
-                var ast = AstExpression.Size(container.Ast);
-                var serializer = Int32Serializer.Instance;
+                if (container.Serializer is IBsonDictionarySerializer dictionarySerializer &&
+                    dictionarySerializer.DictionaryRepresentation == DictionaryRepresentation.Document)
+                {
+                    ast = AstExpression.Size(AstExpression.ObjectToArray(container.Ast));
+                }
+                else
+                {
+                    SerializationHelper.EnsureRepresentationIsArray(expression, container.Serializer);
+                    ast = AstExpression.Size(container.Ast);
+                }
 
+                var serializer = Int32Serializer.Instance;
                 result = new TranslatedExpression(expression, ast, serializer);
                 return true;
             }
@@ -213,6 +226,16 @@ private static bool TryTranslateDictionaryProperty(TranslationContext context, M
 
                         switch (propertyInfo.Name)
                         {
+                            case "Count":
+                                var countAst = dictionaryRepresentation switch
+                                {
+                                    DictionaryRepresentation.ArrayOfDocuments or DictionaryRepresentation.ArrayOfArrays => AstExpression.Size(containerAst),
+                                    _ => throw new ExpressionNotSupportedException(expression, $"Unexpected dictionary representation: {dictionaryRepresentation}")
+                                };
+                                var countSerializer = Int32Serializer.Instance;
+                                translatedDictionaryProperty = new TranslatedExpression(expression, countAst, countSerializer);
+                                return true;
+
                             case "Keys":
                                 var keysAst = dictionaryRepresentation switch
                                 {
@@ -261,5 +284,35 @@ private static bool TryTranslateDictionaryProperty(TranslationContext context, M
             translatedDictionaryProperty = null;
             return false;
         }
+
+        private static bool TryTranslateKeyValuePairProperty(MemberExpression expression, TranslatedExpression container, MemberInfo memberInfo, out TranslatedExpression result)
+        {
+            if (container.Expression.Type.IsGenericType &&
+                container.Expression.Type.GetGenericTypeDefinition() == typeof(KeyValuePair<,>) &&
+                container.Serializer is IKeyValuePairSerializerV2 { Representation: BsonType.Array } kvpSerializer)
+            {
+                AstExpression ast;
+                IBsonSerializer serializer;
+
+                switch (memberInfo.Name)
+                {
+                    case "Key":
+                        ast = AstExpression.ArrayElemAt(container.Ast, 0);
+                        serializer = kvpSerializer.KeySerializer;
+                        break;
+                    case "Value":
+                        ast = AstExpression.ArrayElemAt(container.Ast, 1);
+                        serializer = kvpSerializer.ValueSerializer;
+                        break;
+                    default:
+                        throw new ExpressionNotSupportedException(expression);
+                }
+                result = new TranslatedExpression(expression, ast, serializer);
+                return true;
+            }
+
+            result = null;
+            return false;
+        }
     }
 }

src/MongoDB.Driver/Linq/Linq3Implementation/Translators/ExpressionToAggregationExpressionTranslators/MethodTranslators/ContainsKeyMethodToAggregationExpressionTranslator.cs

@@ -36,27 +36,57 @@ public static TranslatedExpression Translate(TranslationContext context, MethodC
             {
                 var dictionaryExpression = expression.Object;
                 var keyExpression = arguments[0];
+                return TranslateContainsKey(context, expression, dictionaryExpression, keyExpression);
+            }
+
+            throw new ExpressionNotSupportedException(expression);
+        }
 
-                var dictionaryTranslation = ExpressionToAggregationExpressionTranslator.Translate(context, dictionaryExpression);
-                var dictionarySerializer = GetDictionarySerializer(expression, dictionaryTranslation);
-                var dictionaryRepresentation = dictionarySerializer.DictionaryRepresentation;
+        public static TranslatedExpression TranslateContainsKey(TranslationContext context, Expression expression, Expression dictionaryExpression, Expression keyExpression)
+        {
+            var dictionaryTranslation = ExpressionToAggregationExpressionTranslator.Translate(context, dictionaryExpression);
+            var dictionarySerializer = GetDictionarySerializer(expression, dictionaryTranslation);
+            var dictionaryRepresentation = dictionarySerializer.DictionaryRepresentation;
 
-                AstExpression ast;
-                switch (dictionaryRepresentation)
-                {
-                    case DictionaryRepresentation.Document:
+            AstExpression ast;
+            switch (dictionaryRepresentation)
+            {
+                case DictionaryRepresentation.Document:
+                    {
                         var keyFieldName = GetKeyFieldName(context, expression, keyExpression, dictionarySerializer.KeySerializer);
                         ast = AstExpression.IsNotMissing(AstExpression.GetField(dictionaryTranslation.Ast, keyFieldName));
                         break;
+                    }
 
-                    default:
-                        throw new ExpressionNotSupportedException(expression, because: $"ContainsKey is not supported when DictionaryRepresentation is: {dictionaryRepresentation}");
-                }
+                case DictionaryRepresentation.ArrayOfDocuments:
+                    {
+                        var keyFieldName = GetKeyFieldName(context, expression, keyExpression, dictionarySerializer.KeySerializer);
+                        var kvpVar = AstExpression.Var("kvp");
+                        var keysArray = AstExpression.Map(
+                            input: dictionaryTranslation.Ast,
+                            @as: kvpVar,
+                            @in: AstExpression.GetField(kvpVar, "k"));
+                        ast = AstExpression.In(keyFieldName, keysArray);
+                        break;
+                    }
+
+                case DictionaryRepresentation.ArrayOfArrays:
+                    {
+                        var keyFieldName = GetKeyFieldName(context, expression, keyExpression, dictionarySerializer.KeySerializer);
+                        var kvpVar = AstExpression.Var("kvp");
+                        var keysArray = AstExpression.Map(
+                            input: dictionaryTranslation.Ast,
+                            @as: kvpVar,
+                            @in: AstExpression.ArrayElemAt(kvpVar, 0));
+                        ast = AstExpression.In(keyFieldName, keysArray);
+                        break;
+                    }
 
-                return new TranslatedExpression(expression, ast, BooleanSerializer.Instance);
+                default:
+                    throw new ExpressionNotSupportedException(expression, because: $"DictionaryRepresentation: {dictionaryRepresentation} is not supported.");
             }
 
-            throw new ExpressionNotSupportedException(expression);
+            return new TranslatedExpression(expression, ast, BooleanSerializer.Instance);
         }
 
         private static AstExpression GetKeyFieldName(TranslationContext context, Expression expression, Expression keyExpression, IBsonSerializer keySerializer)