Merge pull request #16 from TinkoffCreditSystems/sealed-traits

Sealed traits encoding and decoding support

README.md

@@ -8,11 +8,13 @@ It depends on [cats-core](https://github.com/typelevel/cats) and
 
 Scala 2.12 and 2.13 are supported. Support for Scala 2.11 may be implemented on demand.
 
+:fire: Sealed traits enconding and decoding was implemented in 0.5.0
+
 ## QuickStart
 Add phobos-core to your dependencies:
 
 ```
-libraryDependencies += "ru.tinkoff" %% "phobos-core" % "0.4.0"
+libraryDependencies += "ru.tinkoff" %% "phobos-core" % "0.5.0"
 ```
 
 Then try this code out in `sbt console` or in a separate source file:
@@ -65,7 +67,7 @@ Performance details can be found out in [phobos-benchmark repository](https://gi
 There are several additional modules for some specific cases. 
 These modules could be added with command below:
 ```
-libraryDependencies += "ru.tinkoff" %% "phobos-<module>" % "0.4.0"
+libraryDependencies += "ru.tinkoff" %% "phobos-<module>" % "0.5.0"
 ```
 Where `<module>` is module name.
 

modules/core/src/main/scala/ru/tinkoff/phobos/configured/ElementCodecConfig.scala

@@ -1,12 +1,29 @@
 package ru.tinkoff.phobos.configured
 
-final case class ElementCodecConfig(transformAttributeNames: String => String,
-                                    transformElementNames: String => String) {
-  def withElementsRenamed(transform: String => String): ElementCodecConfig   = copy(transformElementNames = transform)
-  def withAttributesRenamed(transform: String => String): ElementCodecConfig = copy(transformAttributeNames = transform)
-  def withStyle(transform: String => String): ElementCodecConfig = copy(transformElementNames = transform, transformAttributeNames = transform)
+final case class ElementCodecConfig(
+    transformAttributeNames: String => String,
+    transformElementNames: String => String,
+    transformConstructorNames: String => String,
+    discriminatorLocalName: String,
+    discriminatorNamespace: Option[String]
+) {
+  def withElementsRenamed(transform: String => String): ElementCodecConfig =
+    copy(transformElementNames = transform)
+
+  def withAttributesRenamed(transform: String => String): ElementCodecConfig =
+    copy(transformAttributeNames = transform)
+
+  def withConstructorsRenamed(transform: String => String): ElementCodecConfig =
+    copy(transformConstructorNames = transform)
+
+  def withStyle(transform: String => String): ElementCodecConfig =
+    copy(transformElementNames = transform, transformAttributeNames = transform)
+
+  def withDiscriminator(localName: String, namespace: Option[String]): ElementCodecConfig =
+    copy(discriminatorLocalName = localName, discriminatorNamespace = namespace)
 }
 
 object ElementCodecConfig {
-  val default: ElementCodecConfig = ElementCodecConfig(identity, identity)
+  val default: ElementCodecConfig =
+    ElementCodecConfig(identity, identity, identity, "type", Some("http://www.w3.org/2001/XMLSchema-instance"))
 }

modules/core/src/main/scala/ru/tinkoff/phobos/derivation/DecoderDerivation.scala

@@ -3,7 +3,8 @@ package ru.tinkoff.phobos.derivation
 import ru.tinkoff.phobos.Namespace
 import ru.tinkoff.phobos.configured.ElementCodecConfig
 import ru.tinkoff.phobos.decoding.{AttributeDecoder, ElementDecoder, TextDecoder}
-import ru.tinkoff.phobos.derivation.CompileTimeState.{ProductType, Stack}
+import ru.tinkoff.phobos.derivation.CompileTimeState.{CoproductType, ProductType, Stack}
+
 import scala.collection.mutable
 import scala.collection.mutable.ListBuffer
 import scala.reflect.macros.blackbox
@@ -13,12 +14,77 @@ class DecoderDerivation(ctx: blackbox.Context) extends Derivation(ctx) {
 
   def searchType[T: c.WeakTypeTag]: Type = appliedType(c.typeOf[ElementDecoder[_]], c.weakTypeOf[T])
 
-  def deriveProductCodec[T: c.WeakTypeTag](stack: Stack[c.type])(params: IndexedSeq[CaseClassParam]): Tree = {
+  val derivationPkg = q"_root_.ru.tinkoff.phobos.derivation"
+  val decodingPkg   = q"_root_.ru.tinkoff.phobos.decoding"
+  val scalaPkg      = q"_root_.scala"
+  val javaPkg       = q"_root_.java.lang"
+
+  def deriveCoproductCodec[T: c.WeakTypeTag](stack: Stack[c.type])(
+      config: Expr[ElementCodecConfig],
+      subtypes: Iterable[SealedTraitSubtype]
+  ): Tree = {
+    val assignedName = TermName(c.freshName(s"ElementDecoderTypeclass")).encodedName.toTermName
+
+    val elementDecoderType = typeOf[ElementDecoder[_]]
+    val preAssignments     = new ListBuffer[Tree]
+    val classType          = c.weakTypeOf[T]
+
+    val alternatives = subtypes.map { subtype =>
+      val requiredImplicit = appliedType(elementDecoderType, subtype.subtypeType)
+      val path             = CoproductType(weakTypeOf[T].toString)
+      val frame            = stack.Frame(path, appliedType(elementDecoderType, weakTypeOf[T]), assignedName)
+      val derivedImplicit = stack.recurse(frame, requiredImplicit) {
+        typeclassTree(stack)(subtype.subtypeType, elementDecoderType)
+      }
+
+      val ref      = TermName(c.freshName("paramTypeclass"))
+      val assigned = deferredVal(ref, requiredImplicit, derivedImplicit)
 
-    val derivationPkg = q"_root_.ru.tinkoff.phobos.derivation"
-    val decodingPkg   = q"_root_.ru.tinkoff.phobos.decoding"
-    val scalaPkg      = q"_root_.scala"
-    val javaPkg       = q"_root_.java.lang"
+      preAssignments.append(assigned)
+
+      cq""" 
+        discriminator if discriminator == `${subtype.constructorName}` =>
+           $ref.decodeAsElement(cursor, localName, namespaceUri).map(d => d: $classType)
+      """
+    }.toBuffer :+
+      cq""" unknown =>
+          new $decodingPkg.ElementDecoder.FailedDecoder[$classType](
+            cursor.error(s"Unknown type discriminator value: '$${unknown}'")
+        )"""
+
+    q"""
+      ..$preAssignments
+      new $decodingPkg.ElementDecoder[$classType] {
+        def decodeAsElement(
+          cursor : $decodingPkg.Cursor,
+          localName: $javaPkg.String,
+          namespaceUri: $scalaPkg.Option[$javaPkg.String],
+        ): $decodingPkg.ElementDecoder[$classType] = {
+          if (cursor.getEventType == _root_.com.fasterxml.aalto.AsyncXMLStreamReader.EVENT_INCOMPLETE) {
+            this
+          } else {
+            val discriminatorIdx = cursor.getAttributeIndex($config.discriminatorNamespace.getOrElse(null), $config.discriminatorLocalName)
+            if (discriminatorIdx > -1) {
+              cursor.getAttributeValue(discriminatorIdx) match {
+                case ..$alternatives
+              }
+            } else {
+              new $decodingPkg.ElementDecoder.FailedDecoder[$classType](
+                cursor.error(s"No type discriminator '$${$config.discriminatorNamespace.fold("")(_ + ":")}$${$config.discriminatorLocalName}' found")
+              )
+            }
+          }
+        }
+        
+        val isCompleted: $scalaPkg.Boolean = false
+
+        def result(history: $scalaPkg.List[$javaPkg.String]): $scalaPkg.Either[$decodingPkg.DecodingError, $classType] =
+          $scalaPkg.Left($decodingPkg.DecodingError("Decoding not complete", history))
+      }
+    """
+  }
+
+  def deriveProductCodec[T: c.WeakTypeTag](stack: Stack[c.type])(params: IndexedSeq[CaseClassParam]): Tree = {
 
     val decoderStateObj = q"$derivationPkg.DecoderDerivation.DecoderState"
 

modules/core/src/main/scala/ru/tinkoff/phobos/derivation/Derivation.scala

@@ -5,21 +5,34 @@ import ru.tinkoff.phobos.configured.ElementCodecConfig
 import ru.tinkoff.phobos.derivation.CompileTimeState.{ChainedImplicit, Stack}
 import ru.tinkoff.phobos.derivation.Derivation.DirectlyReentrantException
 import ru.tinkoff.phobos.derivation.auto.Exported
-import ru.tinkoff.phobos.syntax.{attr, renamed, text, xmlns}
+import ru.tinkoff.phobos.syntax.{attr, discriminator, renamed, text, xmlns}
+
 import scala.reflect.macros.blackbox
 
 private[phobos] abstract class Derivation(val c: blackbox.Context) {
 
   import c.universe._
 
-  final case class CaseClassParam(localName: String,
-                                  xmlName: Tree,
-                                  namespaceUri: Tree,
-                                  paramType: Type,
-                                  category: ParamCategory)
+  final case class CaseClassParam(
+      localName: String,
+      xmlName: Tree,
+      namespaceUri: Tree,
+      paramType: Type,
+      category: ParamCategory
+  )
+
+  final case class SealedTraitSubtype(
+      constructorName: Tree,
+      subtypeType: Type
+  )
 
   def searchType[T: c.WeakTypeTag]: Type
 
+  def deriveCoproductCodec[T: c.WeakTypeTag](stack: Stack[c.type])(
+      config: Expr[ElementCodecConfig],
+      subtypes: Iterable[SealedTraitSubtype]
+  ): Tree
+
   def deriveProductCodec[T: c.WeakTypeTag](stack: Stack[c.type])(params: IndexedSeq[CaseClassParam]): Tree
 
   def error(msg: String): Nothing = c.abort(c.enclosingPosition, msg)
@@ -30,8 +43,8 @@ private[phobos] abstract class Derivation(val c: blackbox.Context) {
 
   def exportedTypecclass(searchType: Type): Option[Tree] =
     Option(c.inferImplicitValue(appliedType(typeOf[Exported[_]], searchType)))
-      .map(exported => q"$exported.value")
       .filterNot(_.isEmpty)
+      .map(exported => q"$exported.value")
 
   def typeclassTree(stack: Stack[c.type])(genericType: Type, typeConstructor: Type): Tree = {
     val prefixType   = c.prefix.tree.tpe
@@ -68,12 +81,13 @@ private[phobos] abstract class Derivation(val c: blackbox.Context) {
     val classType  = weakTypeOf[T]
     val typeSymbol = classType.typeSymbol
     if (!typeSymbol.isClass) error("Don't know how to work with not classes")
-    val classSymbol   = typeSymbol.asClass
-    val namespaceType = typeOf[Namespace[_]]
-    val attrType      = typeOf[attr]
-    val textType      = typeOf[text]
-    val xmlnsType     = weakTypeOf[xmlns[_]]
-    val renamedType   = typeOf[renamed]
+    val classSymbol       = typeSymbol.asClass
+    val namespaceType     = typeOf[Namespace[_]]
+    val attrType          = typeOf[attr]
+    val textType          = typeOf[text]
+    val xmlnsType         = weakTypeOf[xmlns[_]]
+    val renamedType       = typeOf[renamed]
+    val discriminatorType = typeOf[discriminator]
 
     val expandDeferred = new Transformer {
       override def transform(tree: Tree) = tree match {
@@ -86,7 +100,19 @@ private[phobos] abstract class Derivation(val c: blackbox.Context) {
 
     def inferCodec: Tree = {
       if (classSymbol.isSealed) {
-        error("Sealed traits support is not implemented yet")
+        val sealedTraitSubtypes = classType.typeSymbol.asClass.knownDirectSubclasses.map { symbol =>
+          val constructorName = q"""$config.transformConstructorNames(`${symbol.name.decodedName.toString}`)"""
+          val discriminatorValue = symbol.annotations.collectFirst {
+            case annot if annot.tree.tpe =:= discriminatorType =>
+              annot.tree.children.tail.collectFirst {
+                case t @ Literal(Constant(_: String)) => t
+              }.getOrElse {
+                error("@discriminator is only allowed to be used with string literals")
+              }
+          }.getOrElse(constructorName)
+          SealedTraitSubtype(discriminatorValue, symbol.asType.toType)
+        }
+        deriveCoproductCodec(stack)(config, sealedTraitSubtypes)
       } else if (classSymbol.isCaseClass) {
 
         def fetchGroup(param: TermSymbol): ParamCategory = {
@@ -143,7 +169,7 @@ private[phobos] abstract class Derivation(val c: blackbox.Context) {
             val xmlName: Tree = param.annotations.collectFirst {
               case annotation if annotation.tree.tpe =:= renamedType =>
                 annotation.tree.children.tail.collectFirst {
-                  case t@Literal(Constant(_: String)) => t
+                  case t @ Literal(Constant(_: String)) => t
                 }.getOrElse {
                   error("@renamed is only allowed to be used with string literals")
                 }

modules/core/src/main/scala/ru/tinkoff/phobos/derivation/EncoderDerivation.scala

@@ -2,8 +2,9 @@ package ru.tinkoff.phobos.derivation
 
 import ru.tinkoff.phobos.Namespace
 import ru.tinkoff.phobos.configured.ElementCodecConfig
-import ru.tinkoff.phobos.derivation.CompileTimeState.{ProductType, Stack}
+import ru.tinkoff.phobos.derivation.CompileTimeState.{CoproductType, ProductType, Stack}
 import ru.tinkoff.phobos.encoding.{AttributeEncoder, ElementEncoder, TextEncoder}
+
 import scala.collection.mutable.ListBuffer
 import scala.reflect.macros.blackbox
 
@@ -13,6 +14,55 @@ class EncoderDerivation(ctx: blackbox.Context) extends Derivation(ctx) {
 
   def searchType[T: c.WeakTypeTag]: Type = appliedType(c.typeOf[ElementEncoder[_]], c.weakTypeOf[T])
 
+  def deriveCoproductCodec[T: c.WeakTypeTag](stack: Stack[c.type])(
+      config: Expr[ElementCodecConfig],
+      subtypes: Iterable[SealedTraitSubtype]
+  ): Tree = {
+    val assignedName = TermName(c.freshName(s"ElementEncoderTypeclass")).encodedName.toTermName
+
+    val preAssignments     = new ListBuffer[Tree]
+    val classType          = c.weakTypeOf[T]
+    val scalaPkg           = q"_root_.scala"
+    val javaPkg            = q"_root_.java.lang"
+    val elementEncoderType = typeOf[ElementEncoder[_]]
+
+    val alternatives = subtypes.map { subtype =>
+      val requiredImplicit = appliedType(elementEncoderType, subtype.subtypeType)
+      val path             = CoproductType(weakTypeOf[T].toString)
+      val frame            = stack.Frame(path, appliedType(elementEncoderType, weakTypeOf[T]), assignedName)
+      val derivedImplicit = stack.recurse(frame, requiredImplicit) {
+        typeclassTree(stack)(subtype.subtypeType, elementEncoderType)
+      }
+
+      val ref      = TermName(c.freshName("paramTypeclass"))
+      val assigned = deferredVal(ref, requiredImplicit, derivedImplicit)
+
+      preAssignments.append(assigned)
+
+      cq"""sub: ${subtype.subtypeType.resultType} =>
+             sw.memorizeDiscriminator($config.discriminatorNamespace, $config.discriminatorLocalName, ${subtype.constructorName})
+             $ref.encodeAsElement(sub, sw, localName, namespaceUri)
+      """
+    }
+
+    q"""
+      ..$preAssignments
+
+      new _root_.ru.tinkoff.phobos.encoding.ElementEncoder[$classType] {
+        def encodeAsElement(
+          a: $classType,
+          sw: _root_.ru.tinkoff.phobos.encoding.PhobosStreamWriter,
+          localName: $javaPkg.String,
+          namespaceUri: $scalaPkg.Option[$javaPkg.String]
+        ): $scalaPkg.Unit = {
+          a match {
+            case ..$alternatives
+          }
+        }
+      }
+    """
+  }
+
   def deriveProductCodec[T: c.WeakTypeTag](stack: Stack[c.type])(params: IndexedSeq[CaseClassParam]): Tree = {
     val assignedName = TermName(c.freshName(s"ElementEncoderTypeclass")).encodedName.toTermName
 
@@ -71,7 +121,7 @@ class EncoderDerivation(ctx: blackbox.Context) extends Derivation(ctx) {
      new _root_.ru.tinkoff.phobos.encoding.ElementEncoder[$classType] {
        def encodeAsElement(
          a: $classType,
-         sw: _root_.org.codehaus.stax2.XMLStreamWriter2,
+         sw: _root_.ru.tinkoff.phobos.encoding.PhobosStreamWriter,
          localName: $javaPkg.String,
          namespaceUri: $scalaPkg.Option[$javaPkg.String]
        ): $scalaPkg.Unit = {
@@ -97,7 +147,9 @@ class EncoderDerivation(ctx: blackbox.Context) extends Derivation(ctx) {
   def xmlNs[T: c.WeakTypeTag, NS: c.WeakTypeTag](localName: Tree, ns: Tree): Tree =
     xmlNsConfigured[T, NS](localName, ns, defaultConfig)
 
-  def xmlNsConfigured[T: c.WeakTypeTag, NS: c.WeakTypeTag](localName: Tree, ns: Tree, config: Expr[ElementCodecConfig]): Tree = {
+  def xmlNsConfigured[T: c.WeakTypeTag, NS: c.WeakTypeTag](localName: Tree,
+                                                           ns: Tree,
+                                                           config: Expr[ElementCodecConfig]): Tree = {
     val nsInstance = Option(c.inferImplicitValue(appliedType(weakTypeOf[Namespace[_]], weakTypeOf[NS])))
       .filter(_.nonEmpty)
       .getOrElse(error(s"Could not find Namespace instance for $ns"))

modules/core/src/main/scala/ru/tinkoff/phobos/encoding/AttributeEncoder.scala

@@ -4,7 +4,6 @@ import java.time.{LocalDate, LocalDateTime, LocalTime, ZonedDateTime}
 import java.util.{Base64, UUID}
 
 import cats.Contravariant
-import org.codehaus.stax2.XMLStreamWriter2
 
 /**
  * Warning! This is an internal API which may change in future.
@@ -18,11 +17,11 @@ import org.codehaus.stax2.XMLStreamWriter2
  * To create new instance use .contramap method of existing instance.
  */
 trait AttributeEncoder[A] { self =>
-  def encodeAsAttribute(a: A, sw: XMLStreamWriter2, localName: String, namespaceUri: Option[String]): Unit
+  def encodeAsAttribute(a: A, sw: PhobosStreamWriter, localName: String, namespaceUri: Option[String]): Unit
 
   def contramap[B](f: B => A): AttributeEncoder[B] =
     new AttributeEncoder[B] {
-      def encodeAsAttribute(b: B, sw: XMLStreamWriter2, localName: String, namespaceUri: Option[String]): Unit =
+      def encodeAsAttribute(b: B, sw: PhobosStreamWriter, localName: String, namespaceUri: Option[String]): Unit =
         self.encodeAsAttribute(f(b), sw, localName, namespaceUri)
     }
 }
@@ -38,13 +37,13 @@ object AttributeEncoder {
     */
   implicit val stringEncoder: AttributeEncoder[String] =
     new AttributeEncoder[String] {
-      def encodeAsAttribute(a: String, sw: XMLStreamWriter2, localName: String, namespaceUri: Option[String]): Unit =
+      def encodeAsAttribute(a: String, sw: PhobosStreamWriter, localName: String, namespaceUri: Option[String]): Unit =
         namespaceUri.fold(sw.writeAttribute(localName, a))(ns => sw.writeAttribute(ns, localName, a))
     }
 
   implicit val unitEncoder: AttributeEncoder[Unit] =
     new AttributeEncoder[Unit] {
-      def encodeAsAttribute(a: Unit, sw: XMLStreamWriter2, localName: String, namespaceUri: Option[String]): Unit = ()
+      def encodeAsAttribute(a: Unit, sw: PhobosStreamWriter, localName: String, namespaceUri: Option[String]): Unit = ()
     }
 
   implicit val booleanEncoder: AttributeEncoder[Boolean]                     = stringEncoder.contramap(_.toString)
@@ -74,7 +73,7 @@ object AttributeEncoder {
 
   implicit def optionEncoder[A](implicit encoder: AttributeEncoder[A]): AttributeEncoder[Option[A]] =
     new AttributeEncoder[Option[A]] {
-      def encodeAsAttribute(a: Option[A], sw: XMLStreamWriter2, localName: String, namespaceUri: Option[String]): Unit =
+      def encodeAsAttribute(a: Option[A], sw: PhobosStreamWriter, localName: String, namespaceUri: Option[String]): Unit =
         a.foreach(encoder.encodeAsAttribute(_, sw, localName, namespaceUri))
     }