Merge pull request #403 from AVSystem/mongo-typed-map

Reworked TypedMap and added support for it in typed Mongo API

Author: ghik

commons-core/src/main/scala/com/avsystem/commons/misc/SealedUtils.scala

@@ -94,6 +94,13 @@ trait AutoNamedEnum extends NamedEnum with Product {
   def name: String = productPrefix
 }
 
+/**
+  * Like [[AutoNamedEnum]] but derived names are uncapitalized (first letter lowercased).
+  */
+trait LowerCaseAutoNamedEnum extends AutoNamedEnum {
+  override def name: String = super.name.uncapitalize
+}
+
 /**
   * Base trait for companion objects of sealed traits that serve as named enums. `NamedEnumCompanion` is an
   * extension of [[SealedEnumCompanion]] which additionally requires that every enum value has distinct string

commons-core/src/main/scala/com/avsystem/commons/misc/TypedMap.scala

@@ -1,16 +1,74 @@
-package com.avsystem.commons
-package misc
+package com.avsystem.commons.misc
 
+import com.avsystem.commons.SharedExtensions._
 import com.avsystem.commons.misc.TypedMap.GenCodecMapping
-import com.avsystem.commons.serialization.{GenCodec, GenKeyCodec, ObjectInput, ObjectOutput}
+import com.avsystem.commons.serialization._
 
+/**
+  * A map whose keys are parameterized with value type.
+  * This makes it possible to associate different value type with each key, in a type-safe way.
+  *
+  * [[TypedMap[K]]] has a [[GenCodec]] instance as long as there is a `GenCodec[K[_]]` instance for the key
+  * type and a [[GenCodecMapping[K]]] instance that determines the codec for the value type associated with given
+  * key.
+  *
+  * Example:
+  * {{{
+  *   sealed abstract class AttributeKey[T](implicit val valueCodec: GenCodec[T])
+  *     extends TypedKey[T] with AutoNamedEnum
+  *
+  *   object AttributeKey extends NamedEnumCompanion[AttributeKey[_]] {
+  *     object StringKey extends AttributeKey[String]
+  *     object IntKey extends AttributeKey[Int]
+  *
+  *     val values: List[AttributeKey[_]] = caseObjects
+  *   }
+  *
+  *   val attributes = TypedMap[AttributeKey](
+  *     AttributeKey.StringKey -> "foo",
+  *     AttributeKey.IntKey -> 42,
+  *   )
+  * }}}
+  *
+  * Note that since all keys and value types are known statically,
+  * the map above is somewhat equivalent to a case class:
+  *
+  * {{{
+  *   case class Attributes(
+  *     string: Opt[String],
+  *     int: Opt[Int]
+  *   )
+  * }}}
+  *
+  * [[TypedMap]] might be a good choice if there is a lot of attribute keys, they aren't statically known or some
+  * collection-like behaviour is necessary (e.g. computing the size, iterating over all elements). A [[TypedMap]]
+  * is also easier to evolve than a case class (e.g. because of binary compatibility issues).
+  */
 class TypedMap[K[_]](val raw: Map[K[_], Any]) extends AnyVal {
-  def apply[T](key: K[T]): Option[T] =
-    raw.get(key).map(_.asInstanceOf[T])
+  def apply[T](key: K[T]): T =
+    raw(key).asInstanceOf[T]
+
+  def get[T](key: K[T]): Option[T] =
+    raw.get(key).asInstanceOf[Option[T]]
+
+  def getOpt[T](key: K[T]): Opt[T] =
+    raw.getOpt(key).asInstanceOf[Opt[T]]
+
+  def getOrElse[T](key: K[T], defaultValue: => T): T =
+    get(key).getOrElse(defaultValue)
 
   def updated[T](key: K[T], value: T): TypedMap[K] =
     new TypedMap[K](raw.updated(key, value))
 
+  def ++(other: TypedMap[K]): TypedMap[K] =
+    new TypedMap[K](raw ++ other.raw)
+
+  def keys: Iterable[K[_]] = raw.keys
+  def keysIterator: Iterator[K[_]] = raw.keysIterator
+  def keySet: Set[K[_]] = raw.keySet
+
+  def size: Int = raw.size
+
   override def toString = s"TypedMap($raw)"
 }
 
@@ -33,39 +91,46 @@ object TypedMap {
     def valueCodec[T](key: K[T]): GenCodec[T]
   }
 
-  implicit def typedMapCodec[K[_]](implicit keyCodec: GenKeyCodec[K[_]], codecMapping: GenCodecMapping[K]): GenCodec[TypedMap[K]] =
+  implicit def typedMapCodec[K[_]](implicit
+    keyCodec: GenKeyCodec[K[_]],
+    codecMapping: GenCodecMapping[K]
+  ): GenObjectCodec[TypedMap[K]] =
     new GenCodec.ObjectCodec[TypedMap[K]] {
       def nullable = false
       def readObject(input: ObjectInput): TypedMap[K] = {
         val rawBuilder = Map.newBuilder[K[_], Any]
+        input.knownSize match {
+          case -1 =>
+          case size => rawBuilder.sizeHint(size)
+        }
         while (input.hasNext) {
           val fieldInput = input.nextField()
           val key = keyCodec.read(fieldInput.fieldName)
           rawBuilder += ((key, codecMapping.valueCodec(key).read(fieldInput)))
         }
         new TypedMap[K](rawBuilder.result())
       }
-      def writeObject(output: ObjectOutput, typedMap: TypedMap[K]): Unit =
+      def writeObject(output: ObjectOutput, typedMap: TypedMap[K]): Unit = {
+        output.declareSizeOf(typedMap.raw)
         typedMap.raw.foreach { case (key, value) =>
-          codecMapping.valueCodec(key.asInstanceOf[K[Any]]).write(output.writeField(keyCodec.write(key)), value)
+          val valueCodec = codecMapping.valueCodec(key.asInstanceOf[K[Any]])
+          valueCodec.write(output.writeField(keyCodec.write(key)), value)
         }
+      }
     }
 }
 
 /**
-  * Base class for sealed enums which can be used as key type for a [[TypedMap]]. It also ensures that
-  * the [[TypedMap]] using [[TypedKey]] as a key type will always have a `GenCodec`.
+  * Base class for key types of [[TypedMap]] (typically enums parameterized by value type).
+  * Provides an instance of [[GenCodecMapping]] which is necessary for [[GenCodec]] instance for [[TypedMap]] that
+  * uses this key type.
   */
-abstract class TypedKey[T](implicit val valueCodec: GenCodec[T])
-trait TypedKeyCompanion[K[X] <: TypedKey[X]] extends SealedEnumCompanion[K[_]] {
-  /**
-    * `GenKeyCodec` for typed key `K`.
-    * You can implement this with `GenKeyCodec.forSealedEnum` macro
-    */
-  implicit def keyCodec: GenKeyCodec[K[_]]
-
-  implicit val codecMapping: GenCodecMapping[K] =
+trait TypedKey[T] {
+  def valueCodec: GenCodec[T]
+}
+object TypedKey {
+  implicit def codecMapping[K[X] <: TypedKey[X]]: GenCodecMapping[K] =
     new GenCodecMapping[K] {
-      def valueCodec[T](key: K[T]) = key.valueCodec
+      def valueCodec[T](key: K[T]): GenCodec[T] = key.valueCodec
     }
 }

commons-core/src/main/scala/com/avsystem/commons/serialization/GenCodec.scala

@@ -434,21 +434,9 @@ object GenCodec extends RecursiveAutoCodecs with TupleGenCodecs {
   implicit lazy val BytesCodec: GenCodec[Bytes] =
     GenCodec.nullableSimple(i => Bytes(i.readBinary()), (o, b) => o.writeBinary(b.bytes))
 
-  private def declareSize(o: SequentialOutput, coll: BIterable[_]): Unit =
-    o.sizePolicy match {
-      case SizePolicy.Ignored =>
-      case SizePolicy.Optional =>
-        coll.knownSize match {
-          case -1 =>
-          case size => o.declareSize(size)
-        }
-      case SizePolicy.Required =>
-        o.declareSize(coll.size)
-    }
-
   private implicit class IterableOps[A](private val coll: BIterable[A]) extends AnyVal {
     def writeToList(lo: ListOutput)(implicit writer: GenCodec[A]): Unit = {
-      declareSize(lo, coll)
+      lo.declareSizeOf(coll)
       coll.foreach(new (A => Unit) {
         private var idx = 0
         def apply(a: A): Unit = {
@@ -463,7 +451,7 @@ object GenCodec extends RecursiveAutoCodecs with TupleGenCodecs {
 
   private implicit class PairIterableOps[A, B](private val coll: BIterable[(A, B)]) extends AnyVal {
     def writeToObject(oo: ObjectOutput)(implicit keyWriter: GenKeyCodec[A], writer: GenCodec[B]): Unit = {
-      declareSize(oo, coll)
+      oo.declareSizeOf(coll)
       coll.foreach { case (key, value) =>
         val fieldName = keyWriter.write(key)
         try writer.write(oo.writeField(fieldName), value) catch {

commons-core/src/main/scala/com/avsystem/commons/serialization/InputOutput.scala

@@ -176,6 +176,21 @@ trait SequentialOutput extends Any with AcceptsCustomEvents {
     * have been written. This method MUST always be called after list/object writing has been finished.
     */
   def finish(): Unit
+
+  /**
+    * Based on given collection's `knownSize` and [[sizePolicy]], declares the size
+    * of this output as size of this collection if it is either cheap or required to do so.
+    */
+  final def declareSizeOf(coll: BIterable[_]): Unit = sizePolicy match {
+    case SizePolicy.Ignored =>
+    case SizePolicy.Optional =>
+      coll.knownSize match {
+        case -1 =>
+        case size => declareSize(size)
+      }
+    case SizePolicy.Required =>
+      declareSize(coll.size)
+  }
 }
 /**
   * Represents an abstract sink for serialization of sequences of values. Any [[ListOutput]] instance

commons-core/src/test/scala/com/avsystem/commons/serialization/CodecTestData.scala

@@ -2,9 +2,8 @@ package com.avsystem.commons
 package serialization
 
 import com.avsystem.commons.annotation.AnnotationAggregate
-import com.avsystem.commons.meta.MacroInstances
-import com.avsystem.commons.meta.AutoOptionalParams
-import com.avsystem.commons.misc.{TypedKey, TypedKeyCompanion}
+import com.avsystem.commons.meta.{AutoOptionalParams, MacroInstances}
+import com.avsystem.commons.misc.{AutoNamedEnum, NamedEnumCompanion, TypedKey}
 
 object CodecTestData {
   def col[T <: JCollection[Int]](col: T): T = {
@@ -300,15 +299,13 @@ object CodecTestData {
     implicit val codec: GenCodec[KeyEnumz] = GenCodec.forSealedEnum[KeyEnumz]
   }
 
-  sealed abstract class SealedKey[T: GenCodec] extends TypedKey[T]
-  object SealedKey extends TypedKeyCompanion[SealedKey] {
-    @name("StrKey")
+  sealed abstract class SealedKey[T](implicit val valueCodec: GenCodec[T]) extends TypedKey[T] with AutoNamedEnum
+  object SealedKey extends NamedEnumCompanion[SealedKey[_]] {
     case object StringKey extends SealedKey[String]
     case object IntKey extends SealedKey[Int]
     case object BooleanKey extends SealedKey[Boolean]
 
     val values: List[SealedKey[_]] = caseObjects
-    implicit def keyCodec: GenKeyCodec[SealedKey[_]] = GenKeyCodec.forSealedEnum[SealedKey[_]]
   }
 
   @flatten("kejs") sealed trait CustomizedSeal

commons-core/src/test/scala/com/avsystem/commons/serialization/SimpleGenCodecTest.scala

@@ -322,7 +322,7 @@ class SimpleGenCodecTest extends SimpleIOCodecTest {
 
     testWrite(
       TypedMap(StringKey -> "lol", IntKey -> 42, BooleanKey -> true),
-      Map[String, Any]("StrKey" -> "lol", "IntKey" -> 42, "BooleanKey" -> true)
+      Map[String, Any]("StringKey" -> "lol", "IntKey" -> 42, "BooleanKey" -> true)
     )
   }
 

commons-macros/src/main/scala/com/avsystem/commons/macros/serialization/MongoMacros.scala

@@ -12,6 +12,7 @@ class MongoMacros(ctx: blackbox.Context) extends CodecMacroCommons(ctx) {
   lazy val SeqApplySym: Symbol = typeOf[scala.collection.Seq[Any]].member(TermName("apply"))
   lazy val SeqHeadRef: Symbol = typeOf[scala.collection.Seq[Any]].member(TermName("head"))
   lazy val MapApplySym: Symbol = typeOf[scala.collection.Map[Any, Any]].member(TermName("apply"))
+  lazy val TypedMapApplySym: Symbol = getType(tq"$MiscPkg.TypedMap[$ScalaPkg.Any]").member(TermName("apply"))
   lazy val AdtAsSym: Symbol = getType(tq"$MongoTypedPkg.AbstractMongoDataCompanion[Any, Any]#macroDslExtensions").member(TermName("as"))
 
   // check if some symbol is an abstract method of a sealed trait/class implemented in every case class
@@ -74,6 +75,11 @@ class MongoMacros(ctx: blackbox.Context) extends CodecMacroCommons(ctx) {
 
           q"${extractRefStep(prefix)}.apply($argument)"
 
+        case Apply(TypeApply(Select(prefix, TermName("apply")), List(valueTpeTree)), List(argument))
+          if overridesAnyOf(body.symbol, TypedMapApplySym) =>
+
+          q"${extractRefStep(prefix)}.apply[$valueTpeTree]($argument)"
+
         case _ =>
           invalid(body)
       }

commons-mongo/jvm/src/main/scala/com/avsystem/commons/mongo/typed/MongoFormat.scala

@@ -3,7 +3,7 @@ package mongo.typed
 
 import com.avsystem.commons.annotation.positioned
 import com.avsystem.commons.meta._
-import com.avsystem.commons.misc.ValueOf
+import com.avsystem.commons.misc.{TypedMap, ValueOf}
 import com.avsystem.commons.mongo.{BsonValueInput, BsonValueOutput}
 import com.avsystem.commons.serialization._
 import org.bson.{BsonNull, BsonValue}
@@ -61,6 +61,12 @@ object MongoFormat extends MetadataCompanion[MongoFormat] with MongoFormatLowPri
     valueFormat: MongoFormat[V]
   ) extends MongoFormat[M[K, V]]
 
+  final case class TypedMapFormat[K[_]](
+    codec: GenCodec[TypedMap[K]],
+    keyCodec: GenKeyCodec[K[_]],
+    valueFormats: MongoFormatMapping[K]
+  ) extends MongoFormat[TypedMap[K]]
+
   final case class OptionalFormat[O, T](
     codec: GenCodec[O],
     optionLike: OptionLike.Aux[O, T],
@@ -75,6 +81,10 @@ object MongoFormat extends MetadataCompanion[MongoFormat] with MongoFormatLowPri
     implicit mapCodec: GenCodec[M[K, V]], keyCodec: GenKeyCodec[K], valueFormat: MongoFormat[V]
   ): MongoFormat[M[K, V]] = DictionaryFormat(mapCodec, keyCodec, valueFormat)
 
+  implicit def typedMapFormat[K[_]](
+    implicit keyCodec: GenKeyCodec[K[_]], valueFormats: MongoFormatMapping[K]
+  ): MongoFormat[TypedMap[K]] = TypedMapFormat[K](TypedMap.typedMapCodec, keyCodec, valueFormats)
+
   implicit def optionalFormat[O, T](
     implicit optionLike: OptionLike.Aux[O, T], optionCodec: GenCodec[O], wrappedFormat: MongoFormat[T]
   ): MongoFormat[O] = OptionalFormat(optionCodec, optionLike, wrappedFormat)
@@ -96,6 +106,15 @@ object MongoFormat extends MetadataCompanion[MongoFormat] with MongoFormatLowPri
           "do you have a custom implicit MongoFormat for that type?")
     }
   }
+
+  implicit class typedMapFormatOps[K[_]](private val format: MongoFormat[TypedMap[K]]) extends AnyVal {
+    def assumeTypedMap: TypedMapFormat[K] = format match {
+      case typedMap: TypedMapFormat[K] => typedMap
+      case _ => throw new IllegalArgumentException(
+        "Encountered a non-typed-map MongoFormat for a TypedMap type - " +
+          "do you have a custom implicit MongoFormat for that type?")
+    }
+  }
 }
 trait MongoFormatLowPriority { this: MongoFormat.type =>
   implicit def leafFormat[T: GenCodec]: MongoFormat[T] = Opaque(GenCodec[T])

commons-mongo/jvm/src/main/scala/com/avsystem/commons/mongo/typed/MongoRef.scala

@@ -3,6 +3,7 @@ package mongo.typed
 
 import com.avsystem.commons.annotation.macroPrivate
 import com.avsystem.commons.meta.OptionLike
+import com.avsystem.commons.misc.TypedMap
 import com.avsystem.commons.mongo.typed.MongoPropertyRef.Separator
 import com.avsystem.commons.mongo.{BsonValueInput, KeyEscaper}
 import com.avsystem.commons.serialization.GenCodec.ReadFailure
@@ -257,6 +258,13 @@ object MongoPropertyRef {
     }
   }
 
+  implicit class TypedMapRefOps[E, K[_]](private val ref: MongoPropertyRef[E, TypedMap[K]]) extends AnyVal {
+    def apply[T](key: K[T]): MongoPropertyRef[E, T] = {
+      val tmFormat = ref.format.assumeTypedMap
+      MongoRef.FieldRef(ref, tmFormat.keyCodec.write(key), tmFormat.valueFormats.valueFormat(key), Opt.Empty)
+    }
+  }
+
   implicit def optionalRefOps[E, O, T](ref: MongoPropertyRef[E, O])(implicit optionLike: OptionLike.Aux[O, T]): OptionalRefOps[E, O, T] =
     new OptionalRefOps[E, O, T](ref)
 

commons-mongo/jvm/src/main/scala/com/avsystem/commons/mongo/typed/MongoTypedKey.scala

@@ -0,0 +1,22 @@
+package com.avsystem.commons
+package mongo.typed
+
+import com.avsystem.commons.misc.TypedKey
+import com.avsystem.commons.misc.TypedMap.GenCodecMapping
+import com.avsystem.commons.serialization.GenCodec
+
+trait MongoFormatMapping[K[_]] extends GenCodecMapping[K] {
+  def valueFormat[T](key: K[T]): MongoFormat[T]
+  override def valueCodec[T](key: K[T]): GenCodec[T] = valueFormat(key).codec
+}
+
+trait MongoTypedKey[T] extends TypedKey[T] {
+  def valueFormat: MongoFormat[T]
+  override def valueCodec: GenCodec[T] = valueFormat.codec
+}
+object MongoTypedKey {
+  implicit def mongoFormatMapping[K[X] <: MongoTypedKey[X]]: MongoFormatMapping[K] =
+    new MongoFormatMapping[K] {
+      override def valueFormat[T](key: K[T]): MongoFormat[T] = key.valueFormat
+    }
+}