From 830017e50623936a7688746878beceb56a736aba Mon Sep 17 00:00:00 2001
From: Danny Mor <dannymor83@gmail.com>
Date: Wed, 27 Dec 2023 10:55:21 +0200
Subject: [PATCH] add support for transformed primitve map keys (#627)

* add support for transformed primitve map keys

* run sbt prepare

* restore formatting annotation

* rename for clarity
---
 .../scala/zio/schema/codec/JsonCodec.scala    | 21 ++++++++-
 .../zio/schema/codec/JsonCodecSpec.scala      | 45 +++++++++++++++++++
 2 files changed, 64 insertions(+), 2 deletions(-)

diff --git a/zio-schema-json/shared/src/main/scala/zio/schema/codec/JsonCodec.scala b/zio-schema-json/shared/src/main/scala/zio/schema/codec/JsonCodec.scala
index ac2e0dd4c..af8d3b99f 100644
--- a/zio-schema-json/shared/src/main/scala/zio/schema/codec/JsonCodec.scala
+++ b/zio-schema-json/shared/src/main/scala/zio/schema/codec/JsonCodec.scala
@@ -285,11 +285,26 @@ object JsonCodec {
       }
     //scalafmt: { maxColumn = 120, optIn.configStyleArguments = true }
 
+    private[codec] def transformFieldEncoder[A, B](
+      schema: Schema[A],
+      g: B => Either[String, A]
+    ): Option[JsonFieldEncoder[B]] =
+      jsonFieldEncoder(schema).map { fieldEncoder =>
+        new JsonFieldEncoder[B] {
+          override def unsafeEncodeField(b: B): String =
+            g(b) match {
+              case Left(_)  => throw new RuntimeException(s"Failed to encode field $b")
+              case Right(a) => fieldEncoder.unsafeEncodeField(a)
+            }
+        }
+      }
+
     private[codec] def jsonFieldEncoder[A](schema: Schema[A]): Option[JsonFieldEncoder[A]] =
       schema match {
         case Schema.Primitive(StandardType.StringType, _) => Option(JsonFieldEncoder.string)
         case Schema.Primitive(StandardType.LongType, _)   => Option(JsonFieldEncoder.long)
         case Schema.Primitive(StandardType.IntType, _)    => Option(JsonFieldEncoder.int)
+        case Schema.Transform(c, _, g, a, _)              => transformFieldEncoder(a.foldLeft(c)((s, a) => s.annotate(a)), g)
         case Schema.Lazy(inner)                           => jsonFieldEncoder(inner())
         case _                                            => None
       }
@@ -617,8 +632,10 @@ object JsonCodec {
         case Schema.Primitive(StandardType.StringType, _) => Option(JsonFieldDecoder.string)
         case Schema.Primitive(StandardType.LongType, _)   => Option(JsonFieldDecoder.long)
         case Schema.Primitive(StandardType.IntType, _)    => Option(JsonFieldDecoder.int)
-        case Schema.Lazy(inner)                           => jsonFieldDecoder(inner())
-        case _                                            => None
+        case Schema.Transform(c, f, _, a, _) =>
+          jsonFieldDecoder(a.foldLeft(c)((s, a) => s.annotate(a))).map(_.mapOrFail(f))
+        case Schema.Lazy(inner) => jsonFieldDecoder(inner())
+        case _                  => None
       }
 
     private def dynamicDecoder(schema: Schema.Dynamic): ZJsonDecoder[DynamicValue] = {
diff --git a/zio-schema-json/shared/src/test/scala-2/zio/schema/codec/JsonCodecSpec.scala b/zio-schema-json/shared/src/test/scala-2/zio/schema/codec/JsonCodecSpec.scala
index bc4b39b8e..5fb306e63 100644
--- a/zio-schema-json/shared/src/test/scala-2/zio/schema/codec/JsonCodecSpec.scala
+++ b/zio-schema-json/shared/src/test/scala-2/zio/schema/codec/JsonCodecSpec.scala
@@ -163,6 +163,19 @@ object JsonCodecSpec extends ZIOSpecDefault {
             """{"0":{"first":0,"second":true},"1":{"first":1,"second":false}}"""
           )
         )
+      },
+      test("of complex keys with transformation to primitive keys") {
+        assertEncodes(
+          Schema
+            .map[KeyWrapper, ValueWrapper],
+          Map(
+            KeyWrapper("wrapped_key_1") -> ValueWrapper(value = "some_value"),
+            KeyWrapper("wrapped_key_2") -> ValueWrapper(value = "some_other_value")
+          ),
+          charSequenceToByteChunk(
+            """{"wrapped_key_1":{"value":"some_value"},"wrapped_key_2":{"value":"some_other_value"}}"""
+          )
+        )
       }
     ),
     suite("Set")(
@@ -617,6 +630,19 @@ object JsonCodecSpec extends ZIOSpecDefault {
             """{"0":{"first":0,"second":true},"1":{"first":1,"second":false}}"""
           )
         )
+      },
+      test("of primitive keys with transformation to complex keys") {
+        assertDecodes(
+          Schema
+            .map[KeyWrapper, ValueWrapper],
+          Map(
+            KeyWrapper("wrapped_key_1") -> ValueWrapper(value = "some_value"),
+            KeyWrapper("wrapped_key_2") -> ValueWrapper(value = "some_other_value")
+          ),
+          charSequenceToByteChunk(
+            """{"wrapped_key_1":{"value":"some_value"},"wrapped_key_2":{"value":"some_other_value"}}"""
+          )
+        )
       }
     ),
     suite("zio.json.ast.Json decoding")(
@@ -1605,4 +1631,23 @@ object JsonCodecSpec extends ZIOSpecDefault {
   object ListAndMap {
     implicit lazy val schema: Schema[ListAndMap] = DeriveSchema.gen[ListAndMap]
   }
+
+  final case class KeyWrapper(key: String)
+
+  object KeyWrapper {
+    implicit lazy val schema: Schema[KeyWrapper] = Schema[String].transform(KeyWrapper.apply, _.key)
+  }
+
+  final case class ValueWrapper(value: String)
+
+  object ValueWrapper {
+    implicit lazy val schema: Schema[ValueWrapper] = DeriveSchema.gen[ValueWrapper]
+  }
+
+  final case class MapOfComplexKeysAndValues(map: Map[KeyWrapper, ValueWrapper])
+
+  object MapOfComplexKeysAndValues {
+    implicit lazy val mapSchema: Schema[Map[KeyWrapper, ValueWrapper]] = Schema.map[KeyWrapper, ValueWrapper]
+    implicit lazy val schema: Schema[MapOfComplexKeysAndValues]        = DeriveSchema.gen[MapOfComplexKeysAndValues]
+  }
 }