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,36 @@ 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)
+            {
+                _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 +1374,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("_t", this);
+                }
+                else
+                {
+                    var discriminatorConvention = _baseClassMap.GetDiscriminatorConvention();
+                    if (discriminatorConvention is BsonClassMapScalarDiscriminatorConvention)
                     {
-                        // in this case LookupDiscriminatorConvention below will find it
-                        break;
+                        return new BsonClassMapScalarDiscriminatorConvention(discriminatorConvention.ElementName, this);
                     }
-
-                    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,122 @@
+/* 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.Serialization.Conventions;
+
+namespace MongoDB.Bson.Serialization
+{
+    /// <summary>
+    /// A scalar discriminator convention for class mapped types.
+    /// </summary>
+    public class BsonClassMapScalarDiscriminatorConvention : StandardDiscriminatorConvention, IScalarDiscriminatorConvention
+    {
+        private readonly BsonClassMap _classMap;
+
+        // cached map
+        private readonly ConcurrentDictionary<Type, BsonValue[]> _typeToDiscriminatorsForTypeAndSubTypesMap = new();
+
+        /// <summary>
+        /// Gets the class map.
+        /// </summary>
+        public BsonClassMap ClassMap => _classMap;
+
+        /// <summary>
+        /// Initializes a new instance of BsonClassMapScalarDiscriminatorConvention.
+        /// </summary>
+        /// <param name="elementName">The discriminator element name.</param>
+        /// <param name="classMap">The class map.</param>
+        public BsonClassMapScalarDiscriminatorConvention(string elementName, BsonClassMap classMap)
+            : base(elementName)
+        {
+            _classMap = classMap ?? throw new ArgumentNullException(nameof(classMap));
+        }
+
+        /// <summary>
+        /// Gets the actual type.
+        /// </summary>
+        /// <param name="nominalType">The nominal type.</param>
+        /// <param name="discriminator">The discriminator value.</param>
+        /// <returns>The actual type.</returns>
+        protected override Type GetActualType(Type nominalType, BsonValue discriminator)
+        {
+            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>
+        /// Gets the discriminator value.
+        /// </summary>
+        /// <param name="nominalType">The nominal type.</param>
+        /// <param name="actualType">The actual type.</param>
+        /// <returns>The discriminator value.</returns>
+        public override 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>
+        /// Gets the discriminator values for a type and all of its sub types.
+        /// </summary>
+        /// <param name="type">The type.</param>
+        /// <returns>The discriminator values for a type and all of its sub types.</returns>
+        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;
+                if (type.IsAssignableFrom(discriminatedType))
+                {
+                    var discriminator = entry.Value;
+                    discriminators.Add(discriminator);
+                }
+            }
+
+            return discriminators.OrderBy(x => x).ToArray();
+        }
+    }
+}

src/MongoDB.Bson/Serialization/Conventions/StandardDiscriminatorConvention.cs

@@ -96,35 +96,27 @@ obj is StandardDiscriminatorConvention other &&
         /// <returns>The actual type.</returns>
         public Type GetActualType(IBsonReader bsonReader, Type nominalType)
         {
-            // the BsonReader is sitting at the value whose actual type needs to be found
-            var bsonType = bsonReader.GetCurrentBsonType();
-            if (bsonType == BsonType.Document)
+            // ensure KnownTypes of nominalType are registered (so IsTypeDiscriminated returns correct answer)
+            BsonSerializer.EnsureKnownTypesAreRegistered(nominalType);
+
+            // we can skip looking for a discriminator if nominalType has no discriminated sub types
+            if (!BsonSerializer.IsTypeDiscriminated(nominalType))
             {
-                // ensure KnownTypes of nominalType are registered (so IsTypeDiscriminated returns correct answer)
-                BsonSerializer.EnsureKnownTypesAreRegistered(nominalType);
+                return nominalType;
+            }
 
-                // we can skip looking for a discriminator if nominalType has no discriminated sub types
-                if (BsonSerializer.IsTypeDiscriminated(nominalType))
-                {
-                    var bookmark = bsonReader.GetBookmark();
-                    bsonReader.ReadStartDocument();
-                    var actualType = nominalType;
-                    if (bsonReader.FindElement(_elementName))
-                    {
-                        var context = BsonDeserializationContext.CreateRoot(bsonReader);
-                        var discriminator = BsonValueSerializer.Instance.Deserialize(context);
-                        if (discriminator.IsBsonArray)
-                        {
-                            discriminator = discriminator.AsBsonArray.Last(); // last item is leaf class discriminator
-                        }
-                        actualType = BsonSerializer.LookupActualType(nominalType, discriminator);
-                    }
-                    bsonReader.ReturnToBookmark(bookmark);
-                    return actualType;
-                }
+            var discriminator = ReadDiscriminator(bsonReader);
+            if (discriminator == null)
+            {
+                return nominalType;
+            }
+
+            if (discriminator.IsBsonArray)
+            {
+                discriminator = discriminator.AsBsonArray.Last(); // last item is leaf class discriminator
             }
 
-            return nominalType;
+            return GetActualType(nominalType, discriminator);
         }
 
         /// <summary>
@@ -137,5 +129,47 @@ public Type GetActualType(IBsonReader bsonReader, Type nominalType)
 
         /// <inheritdoc/>
         public override int GetHashCode() => 0;
+
+        // protected methods
+        /// <summary>
+        /// Gets the actual type.
+        /// </summary>
+        /// <param name="nominalType">The nominal type.</param>
+        /// <param name="discriminator">The discriminator.</param>
+        /// <returns>The actual type.</returns>
+        protected virtual Type GetActualType(Type nominalType, BsonValue discriminator)
+        {
+            return BsonSerializer.LookupActualType(nominalType, discriminator);
+        }
+
+        /// <summary>
+        /// Reads the discriminator.
+        /// </summary>
+        /// <param name="bsonReader">The bsonReader.</param>
+        /// <returns>The discriminator, or null if no discriminator was found.</returns>
+        protected 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;
+        }
     }
 }