add support for jsonUnknown trait (#1574)

Implement support for `jsonUnknown` trait defined in disneystreaming/alloy#180

modules/bootstrapped/resources/smithy4s.example.DiscriminatedService.json

@@ -54,52 +54,51 @@
             "TestBiggerUnion": {
                 "oneOf": [
                     {
-                        "allOf": [
-                            {
-                                "$ref": "#/components/schemas/One"
-                            },
-                            {
-                                "type": "object",
-                                "properties": {
-                                    "tpe": {
-                                        "type": "string",
-                                        "enum": [
-                                            "one"
-                                        ]
-                                    }
-                                },
-                                "required": [
-                                    "tpe"
-                                ]
-                            }
-                        ]
+                        "$ref": "#/components/schemas/TestBiggerUnionOne"
                     },
                     {
-                        "allOf": [
-                            {
-                                "$ref": "#/components/schemas/Two"
-                            },
-                            {
-                                "type": "object",
-                                "properties": {
-                                    "tpe": {
-                                        "type": "string",
-                                        "enum": [
-                                            "two"
-                                        ]
-                                    }
-                                },
-                                "required": [
-                                    "tpe"
-                                ]
-                            }
-                        ]
+                        "$ref": "#/components/schemas/TestBiggerUnionTwo"
                     }
                 ],
                 "discriminator": {
-                    "propertyName": "tpe"
+                    "propertyName": "tpe",
+                    "mapping": {
+                        "one": "#/components/schemas/TestBiggerUnionOne",
+                        "two": "#/components/schemas/TestBiggerUnionTwo"
+                    }
                 }
             },
+            "TestBiggerUnionMixin": {
+                "type": "object",
+                "properties": {
+                    "tpe": {
+                        "type": "string"
+                    }
+                },
+                "required": [
+                    "tpe"
+                ]
+            },
+            "TestBiggerUnionOne": {
+                "allOf": [
+                    {
+                        "$ref": "#/components/schemas/One"
+                    },
+                    {
+                        "$ref": "#/components/schemas/TestBiggerUnionMixin"
+                    }
+                ]
+            },
+            "TestBiggerUnionTwo": {
+                "allOf": [
+                    {
+                        "$ref": "#/components/schemas/Two"
+                    },
+                    {
+                        "$ref": "#/components/schemas/TestBiggerUnionMixin"
+                    }
+                ]
+            },
             "Two": {
                 "type": "object",
                 "properties": {

modules/bootstrapped/src/generated/smithy4s/example/AdditionalProperties.scala

@@ -0,0 +1,18 @@
+package smithy4s.example
+
+import smithy4s.Document
+import smithy4s.Hints
+import smithy4s.Newtype
+import smithy4s.Schema
+import smithy4s.ShapeId
+import smithy4s.schema.Schema.bijection
+import smithy4s.schema.Schema.document
+import smithy4s.schema.Schema.map
+import smithy4s.schema.Schema.string
+
+object AdditionalProperties extends Newtype[Map[String, Document]] {
+  val id: ShapeId = ShapeId("smithy4s.example", "AdditionalProperties")
+  val hints: Hints = Hints.empty
+  val underlyingSchema: Schema[Map[String, Document]] = map(string, document).withId(id).addHints(hints)
+  implicit val schema: Schema[AdditionalProperties] = bijection(underlyingSchema, asBijection)
+}

modules/bootstrapped/src/generated/smithy4s/example/JsonUnknownExample.scala

@@ -0,0 +1,27 @@
+package smithy4s.example
+
+import smithy4s.Document
+import smithy4s.Hints
+import smithy4s.Schema
+import smithy4s.ShapeId
+import smithy4s.ShapeTag
+import smithy4s.schema.Schema.int
+import smithy4s.schema.Schema.string
+import smithy4s.schema.Schema.struct
+
+final case class JsonUnknownExample(s: Option[String] = None, i: Option[Int] = None, additionalProperties: Option[Map[String, Document]] = None)
+
+object JsonUnknownExample extends ShapeTag.Companion[JsonUnknownExample] {
+  val id: ShapeId = ShapeId("smithy4s.example", "JsonUnknownExample")
+
+  val hints: Hints = Hints.empty
+
+  // constructor using the original order from the spec
+  private def make(s: Option[String], i: Option[Int], additionalProperties: Option[Map[String, Document]]): JsonUnknownExample = JsonUnknownExample(s, i, additionalProperties)
+
+  implicit val schema: Schema[JsonUnknownExample] = struct(
+    string.optional[JsonUnknownExample]("s", _.s),
+    int.optional[JsonUnknownExample]("i", _.i),
+    AdditionalProperties.underlyingSchema.optional[JsonUnknownExample]("additionalProperties", _.additionalProperties).addHints(alloy.JsonUnknown()),
+  )(make).withId(id).addHints(hints)
+}

modules/bootstrapped/src/generated/smithy4s/example/package.scala

@@ -43,6 +43,7 @@ package object example {
 
   /** This is a simple example of a "quoted string" */
   type AString = smithy4s.example.AString.Type
+  type AdditionalProperties = smithy4s.example.AdditionalProperties.Type
   type Age = smithy4s.example.Age.Type
   /** Multiple line doc comment for another string
     * Containing a random \*\/ here.

modules/codegen-plugin/src/sbt-test/codegen-plugin/update-lsp-config/expected.json

@@ -6,7 +6,7 @@
     ],
     "maven" : {
         "dependencies" : [
-            "com.disneystreaming.alloy:alloy-core:0.3.9"
+            "com.disneystreaming.alloy:alloy-core:0.3.13"
         ],
         "repositories" : [
             {

modules/json/src/smithy4s/json/JsoniterCodecCompiler.scala

@@ -94,7 +94,8 @@ object JsoniterCodecCompiler {
       DiscriminatedUnionMember,
       Default,
       Required,
-      Nullable
+      Nullable,
+      JsonUnknown
     )
 
 }

modules/json/src/smithy4s/json/internals/SchemaVisitorJCodec.scala

@@ -26,6 +26,7 @@ import com.github.plokhotnyuk.jsoniter_scala.core.JsonWriter
 import smithy.api.JsonName
 import smithy.api.TimestampFormat
 import alloy.Discriminated
+import alloy.JsonUnknown
 import alloy.Nullable
 import alloy.Untagged
 import smithy4s.internals.DiscriminatedUnionMember
@@ -1366,6 +1367,9 @@ private[smithy4s] class SchemaVisitorJCodec(
       case Some(x) => x.value
     }
 
+  private def isForJsonUnknown[Z, A](field: Field[Z, A]): Boolean =
+    field.hints.has(JsonUnknown)
+
   private type Handler = (Cursor, JsonReader, util.HashMap[String, Any]) => Unit
 
   private def fieldHandler[Z, A](
@@ -1384,27 +1388,47 @@ private[smithy4s] class SchemaVisitorJCodec(
       )
   }
 
+  private def writeLabel(label: String, out: JsonWriter): Unit =
+    if (label.forall(JsonWriter.isNonEscapedAscii)) {
+      out.writeNonEscapedAsciiKey(label)
+    } else out.writeKey(label)
+
   private def fieldEncoder[Z, A](
       field: Field[Z, A]
   ): (Z, JsonWriter) => Unit = {
     val codec = apply(field.schema)
     val jLabel = jsonLabel(field)
-    val writeLabel: JsonWriter => Unit =
-      if (jLabel.forall(JsonWriter.isNonEscapedAscii)) {
-        _.writeNonEscapedAsciiKey(jLabel)
-      } else _.writeKey(jLabel)
-
     if (explicitDefaultsEncoding) { (z: Z, out: JsonWriter) =>
-      writeLabel(out)
+      writeLabel(jLabel, out)
       codec.encodeValue(field.get(z), out)
     } else { (z: Z, out: JsonWriter) =>
       field.foreachUnlessDefault(z) { (a: A) =>
-        writeLabel(out)
+        writeLabel(jLabel, out)
         codec.encodeValue(a, out)
       }
     }
   }
 
+  private def jsonUnknownFieldEncoder[Z, A](
+      field: Field[Z, A]
+  ): (Z, JsonWriter) => Unit = {
+    val docEncoder = Document.Encoder.fromSchema(field.schema)
+    (z: Z, out: JsonWriter) =>
+      field.foreachUnlessDefault(z) { a =>
+        docEncoder.encode(a) match {
+          case Document.DObject(value) =>
+            value.foreach { case (label: String, value: Document) =>
+              writeLabel(label, out)
+              documentJCodec.encodeValue(value, out)
+            }
+          case _ =>
+            out.encodeError(
+              s"Failed encoding field ${field.label} because it cannot be converted to a JSON object"
+            )
+        }
+      }
+  }
+
   private type Fields[Z] = Vector[Field[Z, _]]
   private type LabelledFields[Z] = Vector[(Field[Z, _], String, Any)]
   private def labelledFields[Z](fields: Fields[Z]): LabelledFields[Z] =
@@ -1415,7 +1439,124 @@ private[smithy4s] class SchemaVisitorJCodec(
       (field, jLabel, default)
     }
 
-  private def nonPayloadStruct[Z](
+  private def structRetainUnknownFields[Z](
+      allFields: LabelledFields[Z],
+      knownFields: LabelledFields[Z],
+      fieldsForUnknown: LabelledFields[Z],
+      structHints: Hints
+  )(
+      const: Vector[Any] => Z,
+      encode: (Z, JsonWriter, Vector[(Z, JsonWriter) => Unit]) => Unit
+  ): JCodec[Z] =
+    new JCodec[Z] {
+
+      private val fieldForUnknownDocumentDecoders = fieldsForUnknown.map {
+        case (field, label, _) =>
+          label -> Document.Decoder
+            .fromSchema(field.schema)
+            .asInstanceOf[Document.Decoder[Any]]
+      }.toMap
+
+      private[this] val handlers =
+        new util.HashMap[String, Handler](knownFields.length << 1, 0.5f) {
+          knownFields.foreach { case (field, jLabel, _) =>
+            put(jLabel, fieldHandler(field))
+          }
+        }
+
+      private[this] val documentEncoders =
+        knownFields.map(labelledField => fieldEncoder(labelledField._1)) ++
+          fieldsForUnknown.map(f => jsonUnknownFieldEncoder(f._1))
+
+      def expecting: String = "object"
+
+      override def canBeKey = false
+
+      def decodeValue(cursor: Cursor, in: JsonReader): Z =
+        decodeValue_(cursor, in)(emptyMetadata)
+
+      private def decodeValue_(
+          cursor: Cursor,
+          in: JsonReader
+      ): scala.collection.Map[String, Any] => Z = {
+        val unknownValues = ListBuffer[(String, Document)]()
+        val buffer = new util.HashMap[String, Any](handlers.size << 1, 0.5f)
+        if (in.isNextToken('{')) {
+          if (!in.isNextToken('}')) {
+            in.rollbackToken()
+            while ({
+              val key = in.readKeyAsString()
+              val handler = handlers.get(key)
+              if (handler eq null) {
+                val value = documentJCodec.decodeValue(cursor, in)
+                unknownValues += (key -> value)
+              } else handler(cursor, in, buffer)
+              in.isNextToken(',')
+            }) ()
+            if (!in.isCurrentToken('}')) in.objectEndOrCommaError()
+          }
+        } else in.decodeError("Expected JSON object")
+
+        // At this point, we have parsed the json and retrieved
+        // all the values that interest us for the construction
+        // of our domain object.
+        // We re-order the values following the order of the schema
+        // fields before calling the constructor.
+        { (meta: scala.collection.Map[String, Any]) =>
+          meta.foreach(kv => buffer.put(kv._1, kv._2))
+          val stage2 = new VectorBuilder[Any]
+          val unknownValue =
+            if (unknownValues.nonEmpty) Document.obj(unknownValues) else null
+
+          allFields.foreach { case (f, jsonLabel, default) =>
+            stage2 += {
+              fieldForUnknownDocumentDecoders.get(jsonLabel) match {
+                case None =>
+                  val value = buffer.get(f.label)
+                  if (value == null) {
+                    if (default == null)
+                      cursor.requiredFieldError(jsonLabel, jsonLabel)
+                    else default
+                  } else value
+
+                case Some(docDecoder) =>
+                  if (unknownValue == null) {
+                    if (default == null) {
+                      docDecoder
+                        .decode(Document.obj())
+                        .getOrElse(
+                          in.decodeError(
+                            s"${cursor.getPath(Nil)} Failed translating a Document.DObject to the type targeted by ${f.label}."
+                          )
+                        )
+                    } else default
+                  } else {
+                    docDecoder
+                      .decode(unknownValue)
+                      .getOrElse(
+                        in.decodeError(
+                          s"${cursor.getPath(Nil)} Failed translating a Document.DObject to the type targeted by ${f.label}."
+                        )
+                      )
+                  }
+              }
+            }
+          }
+          const(stage2.result())
+        }
+      }
+
+      def encodeValue(z: Z, out: JsonWriter): Unit =
+        encode(z, out, documentEncoders)
+
+      def decodeKey(in: JsonReader): Z =
+        in.decodeError("Cannot use products as keys")
+
+      def encodeKey(x: Z, out: JsonWriter): Unit =
+        out.encodeError("Cannot use products as keys")
+    }
+
+  private def structIgnoreUnknownFields[Z](
       fields: LabelledFields[Z],
       structHints: Hints
   )(
@@ -1491,6 +1632,28 @@ private[smithy4s] class SchemaVisitorJCodec(
         out.encodeError("Cannot use products as keys")
     }
 
+  private def nonPayloadStruct[Z](
+      fields: LabelledFields[Z],
+      structHints: Hints
+  )(
+      const: Vector[Any] => Z,
+      encode: (Z, JsonWriter, Vector[(Z, JsonWriter) => Unit]) => Unit
+  ): JCodec[Z] = {
+    val (fieldsForUnknown, knownFields) = fields.partition {
+      case (field, _, _) => isForJsonUnknown(field)
+    }
+
+    if (fieldsForUnknown.isEmpty)
+      structIgnoreUnknownFields(fields, structHints)(const, encode)
+    else
+      structRetainUnknownFields(
+        fields,
+        knownFields,
+        fieldsForUnknown,
+        structHints
+      )(const, encode)
+  }
+
   private def basicStruct[A, S](
       fields: LabelledFields[S],
       structHints: Hints