rewritten specialized GenCodec instances for BsonValue and its subclasses

* BsonInput and BsonOutput impls support BsonValue natively via BsonValueMarker
* codecs for BsonValue and its subclasses always yield binary representation for non-BSON Input & Output impls

Author: Roman Janusz

commons-mongo/jvm/src/main/scala/com/avsystem/commons/mongo/BsonGenCodecs.scala

@@ -1,18 +1,17 @@
 package com.avsystem.commons
 package mongo
 
-import java.nio.ByteBuffer
-import com.avsystem.commons.serialization.GenCodec.{ReadFailure, WriteFailure}
-import com.avsystem.commons.serialization.{GenCodec, GenKeyCodec, TransparentWrapping}
-import org.bson.codecs.{BsonValueCodec, DecoderContext, EncoderContext}
+import com.avsystem.commons.serialization.GenCodec.ReadFailure
+import com.avsystem.commons.serialization._
+import org.bson._
 import org.bson.io.BasicOutputBuffer
 import org.bson.types.{Decimal128, ObjectId}
-import org.bson.{BsonArray, BsonBinary, BsonBinaryReader, BsonBinaryWriter, BsonBoolean, BsonDateTime, BsonDecimal128, BsonDocument, BsonDouble, BsonElement, BsonInt32, BsonInt64, BsonNull, BsonObjectId, BsonString, BsonType, BsonValue}
+
+import java.nio.ByteBuffer
 
 trait BsonGenCodecs {
   implicit def objectIdCodec: GenCodec[ObjectId] = BsonGenCodecs.objectIdCodec
   implicit def objectIdKeyCodec: GenKeyCodec[ObjectId] = BsonGenCodecs.objectIdKeyCodec
-
   implicit def decimal128Codec: GenCodec[Decimal128] = BsonGenCodecs.decimal128Codec
 
   implicit def bsonArrayCodec: GenCodec[BsonArray] = BsonGenCodecs.bsonArrayCodec
@@ -33,13 +32,13 @@ trait BsonGenCodecs {
 object BsonGenCodecs {
   // needed so that ObjectId can be used as ID type in AutoIdMongoEntity
   // (TransparentWrapping is used in EntityIdMode)
-  implicit val objectIdIdentityWrapping: TransparentWrapping[ObjectId, ObjectId] =
-    TransparentWrapping.identity
+  implicit val objectIdIdentityWrapping: TransparentWrapping[ObjectId, ObjectId] = TransparentWrapping.identity
 
   implicit val objectIdCodec: GenCodec[ObjectId] = GenCodec.nullable(
     i => i.readCustom(ObjectIdMarker).getOrElse(new ObjectId(i.readSimple().readString())),
     (o, v) => if (!o.writeCustom(ObjectIdMarker, v)) o.writeSimple().writeString(v.toHexString)
   )
+
   implicit val objectIdKeyCodec: GenKeyCodec[ObjectId] =
     GenKeyCodec.create(new ObjectId(_), _.toHexString)
 
@@ -48,104 +47,43 @@ object BsonGenCodecs {
     (o, v) => if (!o.writeCustom(Decimal128Marker, v)) o.writeSimple().writeBigDecimal(v.bigDecimalValue())
   )
 
-  implicit val bsonArrayCodec: GenCodec[BsonArray] = GenCodec.nullableList(
-    li => new BsonArray(li.iterator(bsonValueCodec.read).to(JList)),
-    (lo, ba) => ba.asScala.foreach(bsonValueCodec.write(lo.writeElement(), _))
-  )
-
-  implicit val bsonBinaryCodec: GenCodec[BsonBinary] =
-    GenCodec.ByteArrayCodec.transform[BsonBinary](_.getData, new BsonBinary(_))
-
-  implicit val bsonBooleanCodec: GenCodec[BsonBoolean] =
-    GenCodec.BooleanCodec.transform[BsonBoolean](_.getValue, new BsonBoolean(_))
-
-  implicit val bsonDateTimeCodec: GenCodec[BsonDateTime] = GenCodec.nullableSimple(
-    i => new BsonDateTime(i.readTimestamp()),
-    (o, v) => o.writeTimestamp(v.getValue)
-  )
-
-  implicit val bsonDocumentCodec: GenCodec[BsonDocument] = GenCodec.nullableObject(
-    oi => new BsonDocument(oi.iterator(bsonValueCodec.read).map {
-      case (k, v) => new BsonElement(k, v)
-    }.to(JList)),
-    (oo, bd) => bd.asScala.foreach { case (key, value) =>
-      bsonValueCodec.write(oo.writeField(key), value)
+  implicit val bsonValueCodec: GenCodec[BsonValue] = GenCodec.create(
+    i => i.readCustom(BsonValueMarker).getOrElse {
+      val reader = new BsonBinaryReader(ByteBuffer.wrap(i.readSimple().readBinary()))
+      BsonValueUtils.decode(reader)
+    },
+    (o, bv) => if (!o.writeCustom(BsonValueMarker, bv)) {
+      val buffer = new BasicOutputBuffer()
+      val writer = new BsonBinaryWriter(buffer)
+      BsonValueUtils.encode(writer, bv)
+      writer.flush()
+      writer.close()
+      o.writeSimple().writeBinary(buffer.toByteArray)
     }
   )
 
-  implicit val bsonDecimal128Codec: GenCodec[BsonDecimal128] =
-    decimal128Codec.transform[BsonDecimal128](_.getValue, new BsonDecimal128(_))
-
-  implicit val bsonDoubleCodec: GenCodec[BsonDouble] =
-    GenCodec.DoubleCodec.transform(_.getValue, new BsonDouble(_))
-
-  implicit val bsonInt32Codec: GenCodec[BsonInt32] =
-    GenCodec.IntCodec.transform(_.getValue, new BsonInt32(_))
+  private def bsonValueSubCodec[T <: BsonValue](fromBsonValue: BsonValue => T): GenCodec[T] =
+    bsonValueCodec.transform(identity, fromBsonValue)
 
-  implicit val bsonInt64Codec: GenCodec[BsonInt64] =
-    GenCodec.LongCodec.transform(_.getValue, new BsonInt64(_))
+  implicit val bsonArrayCodec: GenCodec[BsonArray] = bsonValueSubCodec(_.asArray())
+  implicit val bsonBinaryCodec: GenCodec[BsonBinary] = bsonValueSubCodec(_.asBinary())
+  implicit val bsonBooleanCodec: GenCodec[BsonBoolean] = bsonValueSubCodec(_.asBoolean())
+  implicit val bsonDateTimeCodec: GenCodec[BsonDateTime] = bsonValueSubCodec(_.asDateTime())
+  implicit val bsonDocumentCodec: GenCodec[BsonDocument] = bsonValueSubCodec(_.asDocument())
+  implicit val bsonDecimal128Codec: GenCodec[BsonDecimal128] = bsonValueSubCodec(_.asDecimal128())
+  implicit val bsonDoubleCodec: GenCodec[BsonDouble] = bsonValueSubCodec(_.asDouble())
+  implicit val bsonInt32Codec: GenCodec[BsonInt32] = bsonValueSubCodec(_.asInt32())
+  implicit val bsonInt64Codec: GenCodec[BsonInt64] = bsonValueSubCodec(_.asInt64())
 
   implicit val bsonNullCodec: GenCodec[BsonNull] =
-    GenCodec.create(i => {
-      if (!i.readNull()) throw new ReadFailure("Input did not contain expected null value")
-      BsonNull.VALUE
-    }, (o, _) => o.writeNull())
+    bsonValueSubCodec { bv =>
+      if (bv.isNull) BsonNull.VALUE
+      else throw new ReadFailure("Input did not contain expected null value")
+    }
 
   implicit val bsonObjectIdCodec: GenCodec[BsonObjectId] =
     objectIdCodec.transform(_.getValue, new BsonObjectId(_))
 
   implicit val bsonStringCodec: GenCodec[BsonString] =
     GenCodec.StringCodec.transform(_.getValue, new BsonString(_))
-
-  private val mongoBsonValueCodec: BsonValueCodec = new BsonValueCodec()
-  private val mongoDecoderContext: DecoderContext = DecoderContext.builder().build()
-  private val mongoEncoderContext: EncoderContext = EncoderContext.builder().build()
-
-  implicit val bsonValueCodec: GenCodec[BsonValue] = GenCodec.create(
-    i => {
-      val bvOpt = i.readMetadata(BsonTypeMetadata) map {
-        case BsonType.ARRAY => bsonArrayCodec.read(i)
-        case BsonType.BINARY => bsonBinaryCodec.read(i)
-        case BsonType.BOOLEAN => bsonBooleanCodec.read(i)
-        case BsonType.DATE_TIME => bsonDateTimeCodec.read(i)
-        case BsonType.DECIMAL128 => bsonDecimal128Codec.read(i)
-        case BsonType.DOCUMENT => bsonDocumentCodec.read(i)
-        case BsonType.DOUBLE => bsonDoubleCodec.read(i)
-        case BsonType.INT32 => bsonInt32Codec.read(i)
-        case BsonType.INT64 => bsonInt64Codec.read(i)
-        case BsonType.NULL => bsonNullCodec.read(i)
-        case BsonType.OBJECT_ID => bsonObjectIdCodec.read(i)
-        case BsonType.STRING => bsonStringCodec.read(i)
-        case other => throw new ReadFailure(s"Unsupported Bson type: $other")
-      }
-      bvOpt.getOrElse {
-        val reader = new BsonBinaryReader(ByteBuffer.wrap(i.readSimple().readBinary()))
-        mongoBsonValueCodec.decode(reader, mongoDecoderContext)
-      }
-    },
-    (o, bv) => if (o.keepsMetadata(BsonTypeMetadata)) {
-      bv match {
-        case array: BsonArray => bsonArrayCodec.write(o, array)
-        case binary: BsonBinary => bsonBinaryCodec.write(o, binary)
-        case boolean: BsonBoolean => bsonBooleanCodec.write(o, boolean)
-        case dateTime: BsonDateTime => bsonDateTimeCodec.write(o, dateTime)
-        case decimal128: BsonDecimal128 => bsonDecimal128Codec.write(o, decimal128)
-        case document: BsonDocument => bsonDocumentCodec.write(o, document)
-        case double: BsonDouble => bsonDoubleCodec.write(o, double)
-        case int32: BsonInt32 => bsonInt32Codec.write(o, int32)
-        case int64: BsonInt64 => bsonInt64Codec.write(o, int64)
-        case bNull: BsonNull => bsonNullCodec.write(o, bNull)
-        case objectId: BsonObjectId => bsonObjectIdCodec.write(o, objectId)
-        case string: BsonString => bsonStringCodec.write(o, string)
-        case other => throw new WriteFailure(s"Unsupported value: $other")
-      }
-    } else {
-      val buffer = new BasicOutputBuffer()
-      val writer = new BsonBinaryWriter(buffer)
-      mongoBsonValueCodec.encode(writer, bv, mongoEncoderContext)
-      writer.flush()
-      writer.close()
-      o.writeSimple().writeBinary(buffer.toByteArray)
-    }
-  )
 }

commons-mongo/jvm/src/main/scala/com/avsystem/commons/mongo/BsonInputOutput.scala

@@ -3,19 +3,21 @@ package mongo
 
 import com.avsystem.commons.serialization.GenCodec.ReadFailure
 import com.avsystem.commons.serialization.{InputAndSimpleInput, InputMetadata, OutputAndSimpleOutput, TypeMarker}
-import org.bson.{BsonInvalidOperationException, BsonType}
+import org.bson.{BsonInvalidOperationException, BsonType, BsonValue}
 import org.bson.types.{Decimal128, ObjectId}
 
 import java.nio.ByteBuffer
 
 object ObjectIdMarker extends TypeMarker[ObjectId]
 object Decimal128Marker extends TypeMarker[Decimal128]
+object BsonValueMarker extends TypeMarker[BsonValue]
 
 object BsonTypeMetadata extends InputMetadata[BsonType]
 
 trait BsonInput extends Any with InputAndSimpleInput {
   def readObjectId(): ObjectId
   def readDecimal128(): Decimal128
+  def readBsonValue(): BsonValue
 
   protected def bsonType: BsonType
 
@@ -71,6 +73,7 @@ trait BsonInput extends Any with InputAndSimpleInput {
     typeMarker match {
       case ObjectIdMarker => readObjectId().opt
       case Decimal128Marker => readDecimal128().opt
+      case BsonValueMarker => readBsonValue().opt
       case _ => Opt.Empty
     }
 }
@@ -89,6 +92,7 @@ object BsonInput {
 trait BsonOutput extends Any with OutputAndSimpleOutput {
   def writeObjectId(objectId: ObjectId): Unit
   def writeDecimal128(decimal128: Decimal128): Unit
+  def writeBsonValue(bsonValue: BsonValue): Unit
 
   override def writeBigInt(bigInt: BigInt): Unit =
     if (bigInt.isValidLong) writeLong(bigInt.longValue)
@@ -110,6 +114,7 @@ trait BsonOutput extends Any with OutputAndSimpleOutput {
     typeMarker match {
       case ObjectIdMarker => writeObjectId(value); true
       case Decimal128Marker => writeDecimal128(value); true
+      case BsonValueMarker => writeBsonValue(value); true
       case _ => false
     }
 }

commons-mongo/jvm/src/main/scala/com/avsystem/commons/mongo/BsonReaderInput.scala

@@ -4,7 +4,7 @@ package mongo
 import com.avsystem.commons.serialization.{FieldInput, ListInput, ObjectInput}
 import com.google.common.collect.AbstractIterator
 import org.bson.types.{Decimal128, ObjectId}
-import org.bson.{BsonReader, BsonType}
+import org.bson.{BsonReader, BsonType, BsonValue}
 
 class BsonReaderInput(br: BsonReader, override val legacyOptionEncoding: Boolean = false)
   extends BsonInput {
@@ -57,11 +57,16 @@ class BsonReaderInput(br: BsonReader, override val legacyOptionEncoding: Boolean
   override def readDecimal128(): Decimal128 =
     expect(BsonType.DECIMAL128, br.readDecimal128())
 
+  override def readBsonValue(): BsonValue =
+    BsonValueUtils.decode(br)
+
   override def skip(): Unit =
     br.skipValue()
 
-  override protected final def bsonType: BsonType =
-    br.getCurrentBsonType
+  override protected final def bsonType: BsonType = br.getCurrentBsonType match {
+    case null => br.readBsonType() // reader may be in a state where the type hasn't been read yet
+    case bsonType => bsonType
+  }
 }
 
 final class BsonReaderFieldInput(name: String, br: BsonReader, legacyOptionEncoding: Boolean)

commons-mongo/jvm/src/main/scala/com/avsystem/commons/mongo/BsonValueInput.scala

@@ -1,7 +1,6 @@
 package com.avsystem.commons
 package mongo
 
-import com.avsystem.commons.serialization.GenCodec.ReadFailure
 import com.avsystem.commons.serialization._
 import org.bson._
 import org.bson.types.{Decimal128, ObjectId}
@@ -46,6 +45,9 @@ class BsonValueInput(bsonValue: BsonValue, override val legacyOptionEncoding: Bo
   def readDecimal128(): Decimal128 =
     expect(BsonType.DECIMAL128, bsonValue.asDecimal128().getValue)
 
+  def readBsonValue(): BsonValue =
+    bsonValue
+
   def readNull(): Boolean =
     bsonValue == BsonNull.VALUE
 

commons-mongo/jvm/src/main/scala/com/avsystem/commons/mongo/BsonValueOutput.scala

@@ -66,6 +66,9 @@ final class BsonValueOutput(
 
   override def writeDecimal128(decimal128: Decimal128): Unit =
     setValue(new BsonDecimal128(decimal128))
+
+  override def writeBsonValue(bsonValue: BsonValue): Unit =
+    setValue(bsonValue)
 }
 
 final class BsonValueListOutput(receiver: BsonArray => Unit, legacyOptionEncoding: Boolean)

commons-mongo/jvm/src/main/scala/com/avsystem/commons/mongo/BsonValueUtils.scala

@@ -0,0 +1,19 @@
+package com.avsystem.commons
+package mongo
+
+import org.bson.codecs.{BsonValueCodec, DecoderContext, EncoderContext}
+import org.bson.{BsonReader, BsonValue, BsonWriter}
+
+object BsonValueUtils {
+  private val bsonValueCodec = new BsonValueCodec
+  private val encoderContext = EncoderContext.builder.build
+  private val decoderContext = DecoderContext.builder.build
+
+  def encode(bw: BsonWriter, bv: BsonValue): Unit =
+    bsonValueCodec.encode(bw, bv, encoderContext)
+
+  def decode(br: BsonReader): BsonValue = {
+    br.readBsonType() // without this, `br.getCurrentBsonType` may return null (why?)
+    bsonValueCodec.decode(br, decoderContext)
+  }
+}

commons-mongo/jvm/src/main/scala/com/avsystem/commons/mongo/BsonWriterOutput.scala

@@ -3,7 +3,7 @@ package mongo
 
 import com.avsystem.commons.serialization.{ListOutput, ObjectOutput}
 import org.bson.types.{Decimal128, ObjectId}
-import org.bson.{BsonBinary, BsonWriter}
+import org.bson.{BsonBinary, BsonValue, BsonWriter}
 
 final class BsonWriterOutput(bw: BsonWriter, override val legacyOptionEncoding: Boolean = false)
   extends BsonOutput {
@@ -48,65 +48,24 @@ final class BsonWriterOutput(bw: BsonWriter, override val legacyOptionEncoding:
 
   override def writeDecimal128(decimal128: Decimal128): Unit =
     bw.writeDecimal128(decimal128)
-}
-
-final class BsonWriterNamedOutput(escapedName: String, bw: BsonWriter, override val legacyOptionEncoding: Boolean)
-  extends BsonOutput {
-
-  override def writeNull(): Unit =
-    bw.writeNull(escapedName)
-
-  override def writeString(str: String): Unit =
-    bw.writeString(escapedName, str)
-
-  override def writeBoolean(boolean: Boolean): Unit =
-    bw.writeBoolean(escapedName, boolean)
-
-  override def writeInt(int: Int): Unit =
-    bw.writeInt32(escapedName, int)
-
-  override def writeLong(long: Long): Unit =
-    if (long.isValidInt) writeInt(long.toInt)
-    else bw.writeInt64(escapedName, long)
-
-  override def writeTimestamp(millis: Long): Unit =
-    bw.writeDateTime(escapedName, millis)
-
-  override def writeDouble(double: Double): Unit =
-    bw.writeDouble(escapedName, double)
-
-  override def writeBinary(binary: Array[Byte]): Unit =
-    bw.writeBinaryData(escapedName, new BsonBinary(binary))
 
-  override def writeList(): BsonWriterListOutput = {
-    bw.writeStartArray(escapedName)
-    new BsonWriterListOutput(bw, legacyOptionEncoding)
-  }
-
-  override def writeObject(): BsonWriterObjectOutput = {
-    bw.writeName(escapedName) // org.bson.BsonWriter.writeStartDocument(java.lang.String) fails when writing _id in 4.0 driver
-    bw.writeStartDocument()
-    new BsonWriterObjectOutput(bw, legacyOptionEncoding)
-  }
-
-  override def writeObjectId(objectId: ObjectId): Unit =
-    bw.writeObjectId(escapedName, objectId)
-
-  override def writeDecimal128(decimal128: Decimal128): Unit =
-    bw.writeDecimal128(escapedName, decimal128)
+  override def writeBsonValue(bsonValue: BsonValue): Unit =
+    BsonValueUtils.encode(bw, bsonValue)
 }
 
 final class BsonWriterListOutput(bw: BsonWriter, legacyOptionEncoding: Boolean) extends ListOutput {
-  override def writeElement() =
+  override def writeElement(): BsonWriterOutput =
     new BsonWriterOutput(bw, legacyOptionEncoding)
 
   override def finish(): Unit =
     bw.writeEndArray()
 }
 
 final class BsonWriterObjectOutput(bw: BsonWriter, legacyOptionEncoding: Boolean) extends ObjectOutput {
-  override def writeField(key: String) =
-    new BsonWriterNamedOutput(KeyEscaper.escape(key), bw, legacyOptionEncoding)
+  override def writeField(key: String): BsonWriterOutput = {
+    bw.writeName(KeyEscaper.escape(key))
+    new BsonWriterOutput(bw, legacyOptionEncoding)
+  }
 
   override def finish(): Unit =
     bw.writeEndDocument()

commons-mongo/jvm/src/test/scala/com/avsystem/commons/mongo/BsonInputOutputTest.scala

@@ -167,7 +167,7 @@ class BsonInputOutputTest extends AnyFunSuite with ScalaCheckPropertyChecks {
       fieldInput.skip()
     }
 
-    assert(input.readMetadata(BsonTypeMetadata).isEmpty)
+    assert(input.readMetadata(BsonTypeMetadata).contains(BsonType.DOCUMENT))
 
     val objectInput = input.readObject()
     assert(input.readMetadata(BsonTypeMetadata).contains(BsonType.DOCUMENT))