implemented peekField for BsonReaderInput

Author: Roman Janusz

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

@@ -2,7 +2,7 @@ package com.avsystem.commons
 package mongo
 
 import com.avsystem.commons.serialization.GenCodec.ReadFailure
-import com.avsystem.commons.serialization.{InputAndSimpleInput, InputMetadata, OutputAndSimpleOutput, TypeMarker}
+import com.avsystem.commons.serialization.{FieldInput, InputAndSimpleInput, InputMetadata, OutputAndSimpleOutput, TypeMarker}
 import org.bson.types.{Decimal128, ObjectId}
 import org.bson.{BsonInvalidOperationException, BsonType, BsonValue}
 
@@ -89,6 +89,8 @@ object BsonInput {
   }
 }
 
+trait BsonFieldInput extends BsonInput with FieldInput
+
 trait BsonOutput extends Any with OutputAndSimpleOutput {
   def writeObjectId(objectId: ObjectId): Unit
   def writeDecimal128(decimal128: Decimal128): Unit

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

@@ -1,10 +1,12 @@
 package com.avsystem.commons
 package mongo
 
-import com.avsystem.commons.serialization.{FieldInput, ListInput, ObjectInput}
+import com.avsystem.commons.serialization.{ListInput, ObjectInput}
 import com.google.common.collect.AbstractIterator
+import org.bson._
 import org.bson.types.{Decimal128, ObjectId}
-import org.bson.{BsonReader, BsonType, BsonValue}
+
+import _root_.scala.annotation.tailrec
 
 class BsonReaderInput(br: BsonReader, override val legacyOptionEncoding: Boolean = false)
   extends BsonInput {
@@ -46,9 +48,7 @@ class BsonReaderInput(br: BsonReader, override val legacyOptionEncoding: Boolean
 
   override def readObject(): BsonReaderObjectInput = {
     br.readStartDocument()
-    new BsonReaderObjectInput(new BsonReaderIterator(br, _.readEndDocument(),
-      br => new BsonReaderFieldInput(KeyEscaper.unescape(br.readName()), br, legacyOptionEncoding)
-    ))
+    new BsonReaderObjectInput(br, legacyOptionEncoding)
   }
 
   override def readObjectId(): ObjectId =
@@ -70,7 +70,7 @@ class BsonReaderInput(br: BsonReader, override val legacyOptionEncoding: Boolean
 }
 
 final class BsonReaderFieldInput(name: String, br: BsonReader, legacyOptionEncoding: Boolean)
-  extends BsonReaderInput(br, legacyOptionEncoding) with FieldInput {
+  extends BsonReaderInput(br, legacyOptionEncoding) with BsonFieldInput {
   override def fieldName: String = name
 }
 
@@ -91,7 +91,48 @@ final class BsonReaderListInput(it: BsonReaderIterator[BsonReaderInput]) extends
   override def nextElement(): BsonReaderInput = it.next()
 }
 
-final class BsonReaderObjectInput(it: BsonReaderIterator[BsonReaderFieldInput]) extends ObjectInput {
+final class BsonReaderObjectInput(br: BsonReader, legacyOptionEncoding: Boolean) extends ObjectInput {
+  private[this] val it = new BsonReaderIterator(br, _.readEndDocument(),
+    br => new BsonReaderFieldInput(
+      KeyEscaper.unescape(br.readName()),
+      br,
+      legacyOptionEncoding
+    )
+  )
+
+  private[this] var peekMark: BsonReaderMark = br.getMark
+  private[this] var peekedFields: MHashMap[String, BsonValue] = _
+
+  override def peekField(name: String): Opt[BsonFieldInput] =
+    br match {
+      case _: BsonDocumentReader =>
+        // Looks like there's a bug in BsonDocumentReader.Mark implementation, tests don't pass
+        // TODO: find and report/fix this bug
+        Opt.Empty
+      case _ =>
+        val peekedValue = peekedFields.opt.flatMap(_.getOpt(name)).orElse {
+          val savedMark = br.getMark
+          peekMark.reset()
+
+          @tailrec def loop(): Opt[BsonValue] =
+            if (br.readBsonType() == BsonType.END_OF_DOCUMENT) Opt.Empty
+            else KeyEscaper.unescape(br.readName()) match {
+              case `name` => BsonValueUtils.decode(br).opt
+              case otherName =>
+                if (peekedFields eq null) {
+                  peekedFields = new MHashMap
+                }
+                peekedFields(otherName) = BsonValueUtils.decode(br)
+                peekMark = br.getMark
+                loop()
+            }
+
+          try loop() finally savedMark.reset()
+        }
+
+        peekedValue.map(new BsonValueFieldInput(name, _, legacyOptionEncoding))
+    }
+
   override def hasNext: Boolean = it.hasNext
   override def nextField(): BsonReaderFieldInput = it.next()
 }

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

@@ -61,7 +61,7 @@ class BsonValueInput(bsonValue: BsonValue, override val legacyOptionEncoding: Bo
 }
 
 class BsonValueFieldInput(val fieldName: String, bsonValue: BsonValue, legacyOptionEncoding: Boolean)
-  extends BsonValueInput(bsonValue, legacyOptionEncoding) with FieldInput
+  extends BsonValueInput(bsonValue, legacyOptionEncoding) with BsonFieldInput
 
 class BsonValueListInput(bsonArray: BsonArray, legacyOptionEncoding: Boolean) extends ListInput {
   private val it = bsonArray.iterator

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

@@ -13,7 +13,9 @@ object BsonValueUtils {
     bsonValueCodec.encode(bw, bv, encoderContext)
 
   def decode(br: BsonReader): BsonValue = {
-    br.readBsonType() // without this, `br.getCurrentBsonType` may return null (why?)
+    if (br.getCurrentBsonType eq null) {
+      br.readBsonType()
+    }
     bsonValueCodec.decode(br, decoderContext)
   }
 }

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

@@ -8,13 +8,11 @@ import org.bson.{BsonReader, BsonWriter}
 class GenCodecBasedBsonCodec[T](legacyOptionEncoding: Boolean)(
   implicit ct: ClassTag[T], genCodec: GenCodec[T]
 ) extends Codec[T] {
-  override def getEncoderClass = ct.runtimeClass.asInstanceOf[Class[T]]
+  override def getEncoderClass: Class[T] = ct.runtimeClass.asInstanceOf[Class[T]]
 
-  override def decode(reader: BsonReader, decoderContext: DecoderContext) = {
+  override def decode(reader: BsonReader, decoderContext: DecoderContext): T =
     genCodec.read(new BsonReaderInput(reader, legacyOptionEncoding))
-  }
 
-  override def encode(writer: BsonWriter, value: T, encoderContext: EncoderContext) = {
+  override def encode(writer: BsonWriter, value: T, encoderContext: EncoderContext): Unit =
     genCodec.write(new BsonWriterOutput(writer, legacyOptionEncoding), value)
-  }
 }

commons-mongo/jvm/src/main/scala/com/avsystem/commons/mongo/core/GenCodecRegistry.scala

@@ -8,9 +8,10 @@ object GenCodecRegistry {
   final val LegacyOptionEncoding: Boolean =
     System.getProperty("commons.mongo.legacyOptionEncoding").opt.fold(false)(_.toBoolean)
 
-  def create[T: ClassTag : GenCodec](baseRegistry: CodecRegistry,
-    legacyOptionEncoding: Boolean = LegacyOptionEncoding): CodecRegistry = {
-
+  def create[T: ClassTag : GenCodec](
+    baseRegistry: CodecRegistry,
+    legacyOptionEncoding: Boolean = LegacyOptionEncoding
+  ): CodecRegistry = {
     val genProvider = new GenCodecProvider[T](legacyOptionEncoding)
     val genRegistry = CodecRegistries.fromProviders(genProvider)
     CodecRegistries.fromRegistries(genRegistry, baseRegistry)

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

@@ -244,6 +244,54 @@ class BsonInputOutputTest extends AnyFunSuite with ScalaCheckPropertyChecks {
     }
   }
 
+  test("BsonBinaryReader peekField") {
+    val doc = new BsonDocument()
+      .append("str", new BsonString("str"))
+      .append("int", new BsonInt32(32))
+      .append("boo", new BsonBoolean(false))
+      .append("arr", new BsonArray(JList(new BsonInt32(42), new BsonString("foo"))))
+      .append("dbl", new BsonDouble(3.14))
+
+    val bof = new BasicOutputBuffer
+    val bw = new BsonBinaryWriter(bof)
+    BsonValueUtils.encode(bw, doc)
+    bw.flush()
+    bw.close()
+
+    val br = new BsonBinaryReader(ByteBuffer.wrap(bof.toByteArray))
+    val input = new BsonReaderInput(br).readObject()
+
+    locally {
+      val pf = input.peekField("str")
+      assert(pf.exists(fi => fi.fieldName == "str" && fi.readBsonValue() == new BsonString("str")))
+    }
+    locally {
+      val fi = input.nextField()
+      assert(fi.fieldName == "str")
+      fi.skip()
+    }
+    locally {
+      val pf = input.peekField("boo")
+      assert(pf.exists(fi => fi.fieldName == "boo" && fi.readBsonValue() == new BsonBoolean(false)))
+    }
+    locally {
+      assert(input.peekField("not").isEmpty)
+    }
+    locally {
+      val pf = input.peekField("dbl")
+      assert(pf.exists(fi => fi.fieldName == "dbl" && fi.readBsonValue() == new BsonDouble(3.14)))
+    }
+    locally {
+      val fi = input.nextField()
+      assert(fi.fieldName == "int")
+      fi.skip()
+    }
+    locally {
+      val pf = input.peekField("str")
+      assert(pf.exists(fi => fi.fieldName == "str" && fi.readBsonValue() == new BsonString("str")))
+    }
+  }
+
   def listToBson[T](list: List[T])(converter: T => BsonValue) = new BsonArray(list.map(converter).asJava)
 
   def mapToBson[T](map: Map[String, T])(valueConverter: T => BsonValue): BsonDocument = {