Merge pull request #412 from AVSystem/optimized-bson

Changed `BsonInput`/`BsonOutput` implementations to use optimized value representations

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

@@ -1,23 +1,68 @@
 package com.avsystem.commons
 package mongo
 
-import java.nio.ByteBuffer
-
-import com.avsystem.commons.serialization.{InputAndSimpleInput, InputMetadata, OutputAndSimpleOutput, TypeMarker}
-import org.bson.BsonType
+import com.avsystem.commons.serialization.GenCodec.ReadFailure
+import com.avsystem.commons.serialization.{FieldInput, InputAndSimpleInput, InputMetadata, OutputAndSimpleOutput, TypeMarker}
 import org.bson.types.{Decimal128, ObjectId}
+import org.bson.{BsonInvalidOperationException, BsonType, BsonValue}
+
+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
 
+  protected def handleFailures[T](expr: => T): T =
+    try expr catch {
+      case e: BsonInvalidOperationException => throw new ReadFailure(e.getMessage, e)
+    }
+
+  protected def wrongType(expected: BsonType*): Nothing =
+    throw new ReadFailure(s"Encountered BsonValue of type $bsonType, expected ${expected.mkString(" or ")}")
+
+  protected def expect[T](tpe: BsonType, value: => T): T =
+    if (bsonType == tpe) handleFailures(value)
+    else wrongType(tpe)
+
+  override def readBigInt(): BigInt = handleFailures {
+    bsonType match {
+      case BsonType.INT32 | BsonType.INT64 =>
+        BigInt(readLong())
+      case BsonType.DECIMAL128 =>
+        val bigDec = BigDecimal(readDecimal128().bigDecimalValue())
+        bigDec.toBigIntExact.getOrElse(throw new ReadFailure(s"whole number expected but got $bigDec"))
+      case BsonType.BINARY =>
+        BigInt(readBinary())
+      case _ =>
+        wrongType(BsonType.INT32, BsonType.INT64, BsonType.DECIMAL128, BsonType.BINARY)
+    }
+  }
+
+  override def readBigDecimal(): BigDecimal =
+    handleFailures {
+      bsonType match {
+        case BsonType.INT32 | BsonType.INT64 =>
+          BigDecimal(readLong())
+        case BsonType.DOUBLE =>
+          BigDecimal(readDouble())
+        case BsonType.DECIMAL128 =>
+          BigDecimal(readDecimal128().bigDecimalValue)
+        case BsonType.BINARY =>
+          BsonInput.bigDecimalFromBytes(readBinary())
+        case _ =>
+          wrongType(BsonType.INT32, BsonType.INT64, BsonType.DOUBLE, BsonType.DECIMAL128, BsonType.BINARY)
+      }
+    }
+
   override def readMetadata[T](metadata: InputMetadata[T]): Opt[T] =
     metadata match {
       case BsonTypeMetadata => bsonType.opt
@@ -28,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
     }
 }
@@ -43,9 +89,25 @@ object BsonInput {
   }
 }
 
+trait BsonFieldInput extends BsonInput with FieldInput
+
 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)
+    else Decimal128Utils.fromBigDecimal(BigDecimal(bigInt)) match {
+      case Opt(dec128) => writeDecimal128(dec128)
+      case Opt.Empty => writeBinary(bigInt.toByteArray)
+    }
+
+  override def writeBigDecimal(bigDecimal: BigDecimal): Unit =
+    Decimal128Utils.fromBigDecimal(bigDecimal) match {
+      case Opt(dec128) => writeDecimal128(dec128)
+      case Opt.Empty => writeBinary(BsonOutput.bigDecimalBytes(bigDecimal))
+    }
 
   override def keepsMetadata(metadata: InputMetadata[_]): Boolean =
     BsonTypeMetadata == metadata
@@ -54,13 +116,15 @@ 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
     }
 }
 
 object BsonOutput {
   def bigDecimalBytes(bigDecimal: BigDecimal): Array[Byte] = {
     val unscaledBytes = bigDecimal.bigDecimal.unscaledValue.toByteArray
-    ByteBuffer.allocate(unscaledBytes.length + Integer.BYTES).put(unscaledBytes).putInt(bigDecimal.scale).array
+    ByteBuffer.allocate(unscaledBytes.length + Integer.BYTES)
+      .put(unscaledBytes).putInt(bigDecimal.scale).array
   }
 }