CSHARP-5356: Support OfType and is with scalar discriminators when using class mapped serializers.

src/MongoDB.Bson/Serialization/BsonClassMap.cs

@@ -14,6 +14,7 @@
 */
 
 using System;
+using System.Collections.Concurrent;
 using System.Collections.Generic;
 using System.Collections.ObjectModel;
 using System.Linq;
@@ -61,6 +62,9 @@ public class BsonClassMap
         private int _extraElementsMemberIndex = -1;
         private List<Type> _knownTypes = new List<Type>();
 
+        private ConcurrentDictionary<BsonValue, Type> _discriminatorToTypeMap = new();
+        private ConcurrentDictionary<Type, BsonValue> _typeToDiscriminatorMap = new();
+
         // constructors
         /// <summary>
         /// Initializes a new instance of the BsonClassMap class.
@@ -228,6 +232,10 @@ public bool IsRootClass
             get { return _isRootClass; }
         }
 
+        internal ConcurrentDictionary<BsonValue, Type> DiscriminatorToTypeMap => _discriminatorToTypeMap;
+
+        internal ConcurrentDictionary<Type, BsonValue> TypeToDiscriminatorMap => _typeToDiscriminatorMap;
+
         /// <summary>
         /// Gets the known types of this class.
         /// </summary>
@@ -695,6 +703,11 @@ public BsonClassMap Freeze()
                         }
                         _frozen = true;
 
+                        if (_discriminator != null)
+                        {
+                            AddKnownDiscriminator(_discriminator, _classType);
+                        }
+
                         // use a queue to postpone processing of known types until we get back to the first level call to Freeze
                         // this avoids infinite recursion when going back down the inheritance tree while processing known types
                         foreach (var knownType in _knownTypes)
@@ -1140,6 +1153,37 @@ public void SetExtraElementsMember(BsonMemberMap memberMap)
             _extraElementsMemberMap = memberMap;
         }
 
+        internal void AddKnownDiscriminator(BsonValue discriminator, Type type)
+        {
+            if (!_classType.IsAssignableFrom(type))
+            {
+                throw new ArgumentException($"Type {type} is not assignable to {_classType}.", nameof(type));
+            }
+
+            if (_classType == typeof(object) || _classType == typeof(ValueType) || _classType.IsInterface)
+            {
+                return;
+            }
+
+            if (_baseClassMap == null)
+            {
+                return;
+            }
+
+            _baseClassMap.AddKnownDiscriminator(discriminator, type);
+            var knownType = _discriminatorToTypeMap.GetOrAdd(discriminator, type);
+            if (knownType != type)
+            {
+                throw new ArgumentException($"Duplicate discriminator value \"{discriminator}\".", nameof(discriminator));
+            }
+
+            var knownDiscriminator = _typeToDiscriminatorMap.GetOrAdd(type, discriminator);
+            if (!knownDiscriminator.Equals(discriminator))
+            {
+                throw new ArgumentException($"Duplicate derived type \"{type}\".", nameof(type));
+            }
+        }
+
         /// <summary>
         /// Adds a known type to the class map.
         /// </summary>
@@ -1331,31 +1375,37 @@ internal IDiscriminatorConvention GetDiscriminatorConvention()
 
             IDiscriminatorConvention LookupDiscriminatorConvention()
             {
-                var classMap = this;
-                while (classMap != null)
+                if (_discriminatorConvention != null)
                 {
-                    if (classMap._discriminatorConvention != null)
-                    {
-                        return classMap._discriminatorConvention;
-                    }
+                    return _discriminatorConvention;
+                }
 
-                    if (BsonSerializer.IsDiscriminatorConventionRegisteredAtThisLevel(classMap._classType))
+                if (BsonSerializer.IsDiscriminatorConventionRegisteredAtThisLevel(_classType))
+                {
+                    return BsonSerializer.LookupDiscriminatorConvention(_classType);
+                }
+
+                if (_isRootClass)
+                {
+                    return StandardDiscriminatorConvention.Hierarchical;
+                }
+
+                if (_baseClassMap.ClassType == typeof(object))
+                {
+                    return new BsonClassMapScalarDiscriminatorConvention(this, "_t");
+                }
+                else
+                {
+                    var discriminatorConvention = _baseClassMap.GetDiscriminatorConvention();
+                    if (discriminatorConvention is BsonClassMapScalarDiscriminatorConvention)
                     {
-                        // in this case LookupDiscriminatorConvention below will find it
-                        break;
+                        return new BsonClassMapScalarDiscriminatorConvention(this, discriminatorConvention.ElementName);
                     }
-
-                    if (classMap._isRootClass)
+                    else
                     {
-                        // in this case auto-register a hierarchical convention for the root class and look it up as usual below
-                        BsonSerializer.GetOrRegisterDiscriminatorConvention(classMap._classType, StandardDiscriminatorConvention.Hierarchical);
-                        break;
+                        return discriminatorConvention;
                     }
-
-                    classMap = classMap._baseClassMap;
                 }
-
-                return BsonSerializer.LookupDiscriminatorConvention(_classType);
             }
         }
 

src/MongoDB.Bson/Serialization/BsonClassMapScalarDiscriminatorConvention.cs

@@ -0,0 +1,167 @@
+/* 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;
+using System.Collections.Concurrent;
+using System.Collections.Generic;
+using System.Linq;
+using MongoDB.Bson.IO;
+using MongoDB.Bson.Serialization.Conventions;
+using MongoDB.Bson.Serialization.Serializers;
+
+namespace MongoDB.Bson.Serialization
+{
+    /// <summary>
+    ///
+    /// </summary>
+    public class BsonClassMapScalarDiscriminatorConvention : IScalarDiscriminatorConvention
+    {
+        private readonly BsonClassMap _classMap;
+        private readonly string _elementName;
+
+        // cached map
+        private readonly ConcurrentDictionary<Type, BsonValue[]> _typeToDiscriminatorsForTypeAndSubTypesMap = new();
+
+        /// <summary>
+        ///
+        /// </summary>
+        public BsonClassMap ClassMap => _classMap;
+
+        /// <summary>
+        ///
+        /// </summary>
+        public string ElementName => _elementName;
+
+        /// <summary>
+        ///
+        /// </summary>
+        /// <param name="elementName"></param>
+        /// <param name="classMap"></param>
+        /// <exception cref="ArgumentNullException"></exception>
+        public BsonClassMapScalarDiscriminatorConvention(BsonClassMap classMap, string elementName)
+        {
+            _classMap = classMap ?? throw new ArgumentNullException(nameof(classMap));
+            _elementName = elementName ?? throw new ArgumentNullException(nameof(elementName));
+        }
+
+        /// <summary>
+        ///
+        /// </summary>
+        /// <param name="bsonReader"></param>
+        /// <param name="nominalType"></param>
+        /// <returns></returns>
+        /// <exception cref="NotImplementedException"></exception>
+        public Type GetActualType(IBsonReader bsonReader, Type nominalType)
+        {
+            BsonSerializer.EnsureKnownTypesAreRegistered(nominalType);
+
+            var discriminator = ReadDiscriminator(bsonReader);
+            if (discriminator == null)
+            {
+                return nominalType;
+            }
+
+            if (_classMap.DiscriminatorToTypeMap.TryGetValue(discriminator, out Type actualType))
+            {
+                return actualType;
+            }
+
+            if (nominalType == typeof(object) && discriminator.IsString)
+            {
+                actualType = TypeNameDiscriminator.GetActualType(discriminator.AsString);
+                if (actualType != null)
+                {
+                    return actualType;
+                }
+            }
+
+            throw new BsonSerializationException($"No type found for discriminator value: {discriminator}.");
+        }
+
+        /// <summary>
+        ///
+        /// </summary>
+        /// <param name="nominalType"></param>
+        /// <param name="actualType"></param>
+        /// <returns></returns>
+        /// <exception cref="NotImplementedException"></exception>
+        public BsonValue GetDiscriminator(Type nominalType, Type actualType)
+        {
+            if (actualType == nominalType && !_classMap.DiscriminatorIsRequired)
+            {
+                return null;
+            }
+
+            if (_classMap.TypeToDiscriminatorMap.TryGetValue(actualType, out BsonValue discriminator))
+            {
+                return discriminator;
+            }
+
+            throw new BsonSerializationException($"No discriminator value found for type: {actualType}.");
+        }
+
+        /// <summary>
+        ///
+        /// </summary>
+        /// <param name="type"></param>
+        /// <returns></returns>
+        /// <exception cref="NotImplementedException"></exception>
+        public BsonValue[] GetDiscriminatorsForTypeAndSubTypes(Type type)
+        {
+            return _typeToDiscriminatorsForTypeAndSubTypesMap.GetOrAdd(type, MapTypeToDiscriminatorsForTypeAndSubTypes);
+        }
+
+        private BsonValue[] MapTypeToDiscriminatorsForTypeAndSubTypes(Type type)
+        {
+            var discriminators = new List<BsonValue>();
+            foreach (var entry in _classMap.TypeToDiscriminatorMap)
+            {
+                var discriminatedType = entry.Key;
+                var discriminator = entry.Value;
+                if (type.IsAssignableFrom(discriminatedType))
+                {
+                    discriminators.Add(discriminator);
+                }
+            }
+
+            return discriminators.OrderBy(x => x).ToArray();
+        }
+
+        private BsonValue ReadDiscriminator(IBsonReader bsonReader)
+        {
+            // the BsonReader is sitting at the value whose actual type needs to be found
+            var bsonType = bsonReader.GetCurrentBsonType();
+            if (bsonType == BsonType.Document)
+            {
+                var bookmark = bsonReader.GetBookmark();
+                try
+                {
+                    bsonReader.ReadStartDocument();
+                    if (bsonReader.FindElement(_elementName))
+                    {
+                        var context = BsonDeserializationContext.CreateRoot(bsonReader);
+                        return BsonValueSerializer.Instance.Deserialize(context);
+                    }
+                }
+                finally
+                {
+                    bsonReader.ReturnToBookmark(bookmark);
+                }
+            }
+
+            return null;
+        }
+    }
+}