From 047ce84454698f84d747d058419015259513ad70 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Olivier=20M=C3=A9lois?= <olivier.melois@disneystreaming.com>
Date: Thu, 14 Mar 2024 14:55:17 +0100
Subject: [PATCH] Update translation of protoTimestampFormat

See https://github.com/disneystreaming/alloy/pull/150
---
 buildDeps.sc                                     |  2 +-
 buildSetup.sc                                    |  8 ++++++++
 .../proto3/internals/Compiler.scala              | 16 +++++++++-------
 .../proto3/internals/ProtoIR.scala               |  8 ++++++++
 .../proto3/internals/CompilerRendererSuite.scala |  6 ++++--
 5 files changed, 30 insertions(+), 10 deletions(-)

diff --git a/buildDeps.sc b/buildDeps.sc
index a44ebf9..b37f67f 100644
--- a/buildDeps.sc
+++ b/buildDeps.sc
@@ -3,7 +3,7 @@ import mill.define._
 import mill.scalalib._
 
 object alloy {
-  val alloyVersion = "0.3.3"
+  val alloyVersion = "0.3.3-1-88e46a-SNAPSHOT"
   val core =
     ivy"com.disneystreaming.alloy:alloy-core:$alloyVersion"
   val protobuf =
diff --git a/buildSetup.sc b/buildSetup.sc
index 4e8bae8..ce1532d 100644
--- a/buildSetup.sc
+++ b/buildSetup.sc
@@ -115,6 +115,14 @@ trait BaseScala213Module extends BaseScalaModule with ScalafmtModule {
 
 trait BaseScalaModule extends ScalaModule with BaseModule {
 
+  override def repositoriesTask = T.task {
+    super.repositoriesTask() ++ Seq(
+      coursier.maven.MavenRepository(
+        "https://s01.oss.sonatype.org/content/repositories/snapshots/"
+      )
+    )
+  }
+
   override def scalacPluginIvyDeps = T {
     val sv = scalaVersion()
     val plugins =
diff --git a/modules/proto/src/smithytranslate/proto3/internals/Compiler.scala b/modules/proto/src/smithytranslate/proto3/internals/Compiler.scala
index cd61eb6..2227f00 100644
--- a/modules/proto/src/smithytranslate/proto3/internals/Compiler.scala
+++ b/modules/proto/src/smithytranslate/proto3/internals/Compiler.scala
@@ -32,6 +32,7 @@ import software.amazon.smithy.model.neighbor.Walker
 import alloy.OpenEnumTrait
 import software.amazon.smithy.model.traits.EnumValueTrait
 import alloy.proto.ProtoTimestampFormatTrait
+import alloy.proto.ProtoTimestampFormatTrait.TimestampFormat
 
 private[proto3] class Compiler(model: Model, allShapes: Boolean) {
 
@@ -659,13 +660,14 @@ private[proto3] class Compiler(model: Model, allShapes: Boolean) {
           extractTimestampFormat(shape)
             .orElse(timestampFormat)
             .getOrElse(ProtoTimestampFormatTrait.TimestampFormat.PROTOBUF)
-        val isEpochMillis =
-          (format == ProtoTimestampFormatTrait.TimestampFormat.EPOCH_MILLIS)
-        if (!isWrapped) {
-          if (isEpochMillis) Type.Int64 else Type.GoogleTimestamp
-        } else {
-          if (isEpochMillis) Type.GoogleWrappers.Int64
-          else Type.AlloyWrappers.Timestamp
+
+        format match {
+          case TimestampFormat.PROTOBUF | TimestampFormat.UNKNOWN =>
+            if (isWrapped) Type.AlloyWrappers.Timestamp
+            else Type.GoogleTimestamp
+          case TimestampFormat.EPOCH_MILLIS =>
+            if (isWrapped) Type.AlloyWrappers.EpochMillisTimestamp
+            else Type.AlloyTypes.EpochMillisTimestamp
         }
       }
 
diff --git a/modules/proto/src/smithytranslate/proto3/internals/ProtoIR.scala b/modules/proto/src/smithytranslate/proto3/internals/ProtoIR.scala
index 271956e..0b171c7 100644
--- a/modules/proto/src/smithytranslate/proto3/internals/ProtoIR.scala
+++ b/modules/proto/src/smithytranslate/proto3/internals/ProtoIR.scala
@@ -157,6 +157,10 @@ private[internals] object ProtoIR {
         alloyFqn("CompactUUID"),
         alloyTypesImport
       )
+      val EpochMillisTimestamp = RefType(
+        alloyFqn("EpochMillisTimestamp"),
+        alloyTypesImport
+      )
     }
 
     object AlloyWrappers {
@@ -204,6 +208,10 @@ private[internals] object ProtoIR {
         alloyFqn("TimestampValue"),
         alloyWrappersImport
       )
+      val EpochMillisTimestamp = RefType(
+        alloyFqn("EpochMillisTimestampValue"),
+        alloyWrappersImport
+      )
       val CompactUUID = RefType(
         alloyFqn("CompactUUIDValue"),
         alloyWrappersImport
diff --git a/modules/proto/tests/src/smithytranslate/proto3/internals/CompilerRendererSuite.scala b/modules/proto/tests/src/smithytranslate/proto3/internals/CompilerRendererSuite.scala
index 5ba7877..76577a4 100644
--- a/modules/proto/tests/src/smithytranslate/proto3/internals/CompilerRendererSuite.scala
+++ b/modules/proto/tests/src/smithytranslate/proto3/internals/CompilerRendererSuite.scala
@@ -207,6 +207,8 @@ class CompilerRendererSuite extends FunSuite {
                       |
                       |import "google/protobuf/timestamp.proto";
                       |
+                      |import "alloy/protobuf/types.proto";
+                      |
                       |message Struct {
                       |  bool boolean = 1;
                       |  int32 int = 2;
@@ -221,7 +223,7 @@ class CompilerRendererSuite extends FunSuite {
                       |  google.protobuf.Value document = 11;
                       |  string string = 12;
                       |  google.protobuf.Timestamp timestamp = 13;
-                      |  int64 epoch = 14;
+                      |  alloy.protobuf.EpochMillisTimestamp epoch = 14;
                       |}
                       |""".stripMargin
     convertCheck(source, Map("com/example/example.proto" -> expected))
@@ -288,7 +290,7 @@ class CompilerRendererSuite extends FunSuite {
                       |  alloy.protobuf.DocumentValue document = 11;
                       |  google.protobuf.StringValue string = 12;
                       |  alloy.protobuf.TimestampValue timestamp = 13;
-                      |  google.protobuf.Int64Value epoch = 14;
+                      |  alloy.protobuf.EpochMillisTimestampValue epoch = 14;
                       |}
                       |""".stripMargin
     convertCheck(source, Map("com/example/example.proto" -> expected))