CSHARP-4876: Sort OfType in nested expressions.

Author: rstam

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

@@ -63,6 +63,7 @@ public static AggregationExpression Translate(TranslationContext context, Method
                 case "IsNullOrWhiteSpace": return IsNullOrWhiteSpaceMethodToAggregationExpressionTranslator.Translate(context, expression);
                 case "IsSubsetOf": return IsSubsetOfMethodToAggregationExpressionTranslator.Translate(context, expression);
                 case "Locf": return LocfMethodToAggregationExpressionTranslator.Translate(context, expression);
+                case "OfType": return OfTypeMethodToAggregationExpressionTranslator.Translate(context, expression);
                 case "Parse": return ParseMethodToAggregationExpressionTranslator.Translate(context, expression);
                 case "Pow": return PowMethodToAggregationExpressionTranslator.Translate(context, expression);
                 case "Push": return PushMethodToAggregationExpressionTranslator.Translate(context, expression);

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

@@ -0,0 +1,80 @@
+/* Copyright 2010-present MongoDB Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+using System.Linq.Expressions;
+using System.Reflection;
+using MongoDB.Bson.Serialization;
+using MongoDB.Bson.Serialization.Conventions;
+using MongoDB.Driver.Linq.Linq3Implementation.Ast.Expressions;
+using MongoDB.Driver.Linq.Linq3Implementation.Misc;
+using MongoDB.Driver.Linq.Linq3Implementation.Reflection;
+using MongoDB.Driver.Linq.Linq3Implementation.Serializers;
+
+namespace MongoDB.Driver.Linq.Linq3Implementation.Translators.ExpressionToAggregationExpressionTranslators.MethodTranslators
+{
+    internal static class OfTypeMethodToAggregationExpressionTranslator
+    {
+        private static readonly MethodInfo[] __ofTypeMethods =
+        {
+            EnumerableMethod.OfType,
+            QueryableMethod.OfType
+        };
+
+        public static AggregationExpression Translate(TranslationContext context, MethodCallExpression expression)
+        {
+            var method = expression.Method;
+            var arguments = expression.Arguments;
+
+            if (method.IsOneOf(__ofTypeMethods))
+            {
+                var sourceExpression = arguments[0];
+                var sourceTranslation = ExpressionToAggregationExpressionTranslator.TranslateEnumerable(context, sourceExpression);
+                NestedAsQueryableHelper.EnsureQueryableMethodHasNestedAsQueryableSource(expression, sourceTranslation);
+                var itemSerializer = ArraySerializerHelper.GetItemSerializer(sourceTranslation.Serializer);
+
+                var nominalType = sourceTranslation.Serializer.ValueType;
+                var actualType = method.GetGenericArguments()[0];
+
+                if (nominalType == actualType)
+                {
+                    return sourceTranslation;
+                }
+
+                var discriminatorConvention = itemSerializer.GetDiscriminatorConvention();
+                var discriminatorElementName = discriminatorConvention.ElementName;
+
+                var itemVar = AstExpression.Var("this");
+                var discriminatorField = AstExpression.GetField(itemVar, discriminatorElementName);
+                var ofTypePredicate = discriminatorConvention switch
+                {
+                    IHierarchicalDiscriminatorConvention hierarchicalDiscriminatorConvention => DiscriminatorAstExpression.TypeIs(discriminatorField, hierarchicalDiscriminatorConvention, nominalType, actualType),
+                    IScalarDiscriminatorConvention scalarDiscriminatorConvention => DiscriminatorAstExpression.TypeIs(discriminatorField, scalarDiscriminatorConvention, nominalType, actualType),
+                    _ => throw new ExpressionNotSupportedException(expression, because: "is operator is not supported with the configured discriminator convention")
+                };
+
+                var ast = AstExpression.Filter(
+                    input: sourceTranslation.Ast,
+                    @as: itemVar.Name,
+                    cond: ofTypePredicate);
+                var actualTypeSerializer = BsonSerializer.LookupSerializer(actualType);
+                var resultSerializer = NestedAsQueryableSerializer.CreateIEnumerableOrNestedAsQueryableSerializer(expression.Type, actualTypeSerializer);
+
+                return new AggregationExpression(expression, ast, resultSerializer);
+            }
+
+            throw new ExpressionNotSupportedException(expression);
+        }
+    }
+}

tests/MongoDB.Driver.Tests/Linq/Linq3Implementation/Jira/CSharp4876Tests.cs

@@ -0,0 +1,80 @@
+/* Copyright 2010-present MongoDB Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+using System.Linq;
+using FluentAssertions;
+using MongoDB.Bson.Serialization.Attributes;
+using MongoDB.TestHelpers.XunitExtensions;
+using Xunit;
+
+namespace MongoDB.Driver.Tests.Linq.Linq3Implementation.Jira
+{
+    public class CSharp4876Tests : Linq3IntegrationTest
+    {
+        [Theory]
+        [ParameterAttributeData]
+        public void OfType_should_work(
+            [Values(false, true)] bool withNestedAsQueryable)
+        {
+            var collection = GetCollection();
+
+            var queryable = withNestedAsQueryable ?
+                collection.AsQueryable().Select(x => x.A.AsQueryable().OfType<B2>().ToArray()) :
+                collection.AsQueryable().Select(x => x.A.OfType<B2>().ToArray());
+
+            var stages = Translate(collection, queryable);
+            AssertStages(stages, "{ $project : { _v : { $filter : { input : '$A', as : 'this', cond : { $cond : { if : { $eq : [{ $type : '$$this._t' }, 'array'] }, then : { $in : ['B2', '$$this._t'] }, else : { $eq : ['$$this._t', 'B2'] } } } } }, _id : 0 } }");
+
+            var result = queryable.Single();
+            result.Select(x => x.Id).Should().Equal(2);
+        }
+
+        private IMongoCollection<C> GetCollection()
+        {
+            var collection = GetCollection<C>("test");
+            CreateCollection(
+                collection,
+                new C
+                {
+                    Id = 1, A =
+                        [
+                            new B1 { Id = 1 },
+                            new B2 { Id = 2 }
+                        ]
+                });
+            return collection;
+        }
+
+        private class C
+        {
+            public int Id { get; set; }
+            public B[] A { get; set; }
+        }
+
+        [BsonDiscriminator(RootClass = true)]
+        public class B
+        {
+            public int Id { get; set; }
+        }
+
+        public class B1 : B
+        {
+        }
+
+        public class B2 : B
+        {
+        }
+    }
+}