Add wrapper types for map keys and values (#73)

open-toast · Jun 21, 2020 · 0215f0a · 0215f0a
1 parent fb529f5
commit 0215f0a
Show file tree

Hide file tree

Showing 18 changed files with 533 additions and 243 deletions.
diff --git a/protokt-codegen/src/main/kotlin/com/toasttab/protokt/codegen/impl/AnnotatorUtils.kt b/protokt-codegen/src/main/kotlin/com/toasttab/protokt/codegen/impl/AnnotatorUtils.kt
@@ -19,37 +19,35 @@ import arrow.core.None
 import arrow.core.Some
 import com.toasttab.protokt.codegen.algebra.AST
 import com.toasttab.protokt.codegen.impl.STAnnotator.Context
+import com.toasttab.protokt.codegen.impl.Wrapper.interceptMapKeyTypeName
+import com.toasttab.protokt.codegen.impl.Wrapper.interceptMapValueTypeName
+import com.toasttab.protokt.codegen.protoc.MapEntry
 import com.toasttab.protokt.codegen.protoc.Message
 import com.toasttab.protokt.codegen.protoc.Oneof
 import com.toasttab.protokt.codegen.protoc.StandardField
 import com.toasttab.protokt.codegen.protoc.TypeDesc
 import com.toasttab.protokt.codegen.template.Renderers.ConcatWithScope
 
-internal fun resolveMapEntry(m: Message) =
-    MapEntryInfo(
+fun resolveMapEntry(m: Message) =
+    MapEntry(
         (m.fields[0] as StandardField),
         (m.fields[1] as StandardField)
     )
 
-internal fun resolveMapEntryTypes(m: Message, ctx: Context) =
-    resolveMapEntry(m).let {
+fun resolveMapEntryTypes(f: StandardField, ctx: Context) =
+    f.mapEntry!!.let {
         MapTypeParams(
-            it.key.unqualifiedTypeName,
-            it.value.typePClass.renderName(ctx.pkg)
+            interceptMapKeyTypeName(f, it.key.unqualifiedTypeName, ctx),
+            interceptMapValueTypeName(f, it.value.typePClass.renderName(ctx.pkg), ctx)
         )
     }
 
-internal class MapTypeParams(
+class MapTypeParams(
     val kType: String,
     val vType: String
 )
 
-internal class MapEntryInfo(
-    val key: StandardField,
-    val value: StandardField
-)
-
-internal fun oneOfScope(f: Oneof, type: String, ctx: Context) =
+fun oneOfScope(f: Oneof, type: String, ctx: Context) =
     ctx.stripEnclosingMessageNamePrefix(
         ctx.stripRootMessageNamePrefix(
             ConcatWithScope.render(
@@ -59,14 +57,14 @@ internal fun oneOfScope(f: Oneof, type: String, ctx: Context) =
         )
     )
 
-internal fun String.emptyToNone() =
+fun String.emptyToNone() =
     if (isEmpty()) {
         None
     } else {
         Some(this)
     }
 
-internal fun kotlinPackage(ast: AST<TypeDesc>) =
+fun kotlinPackage(ast: AST<TypeDesc>) =
     resolvePackage(
         ast.data.desc.options,
         ast.data.desc.packageName,

diff --git a/protokt-codegen/src/main/kotlin/com/toasttab/protokt/codegen/impl/DeserializerAnnotator.kt b/protokt-codegen/src/main/kotlin/com/toasttab/protokt/codegen/impl/DeserializerAnnotator.kt
@@ -22,6 +22,10 @@ import arrow.core.getOrElse
 import com.toasttab.protokt.codegen.impl.MessageAnnotator.idealMaxWidth
 import com.toasttab.protokt.codegen.impl.STAnnotator.Context
 import com.toasttab.protokt.codegen.impl.Wrapper.interceptReadFn
+import com.toasttab.protokt.codegen.impl.Wrapper.keyWrapped
+import com.toasttab.protokt.codegen.impl.Wrapper.mapKeyConverter
+import com.toasttab.protokt.codegen.impl.Wrapper.mapValueConverter
+import com.toasttab.protokt.codegen.impl.Wrapper.valueWrapped
 import com.toasttab.protokt.codegen.impl.Wrapper.wrapped
 import com.toasttab.protokt.codegen.impl.Wrapper.wrapperName
 import com.toasttab.protokt.codegen.model.FieldType
@@ -99,9 +103,12 @@ private constructor(
             lhs = f.fieldName,
             packed = packed,
             options =
-                if (f.wrapped) {
+                if (f.wrapped || f.keyWrapped || f.valueWrapped) {
                     Options(
                         wrapName = wrapperName(f, ctx).getOrElse { "" },
+                        keyWrap = mapKeyConverter(f, ctx),
+                        valueWrap = mapValueConverter(f, ctx),
+                        valueType = f.mapEntry?.value?.type,
                         type = f.type.toString(),
                         oneof = true
                     )

diff --git a/protokt-codegen/src/main/kotlin/com/toasttab/protokt/codegen/impl/PropertyAnnotator.kt b/protokt-codegen/src/main/kotlin/com/toasttab/protokt/codegen/impl/PropertyAnnotator.kt
@@ -15,11 +15,6 @@
 
 package com.toasttab.protokt.codegen.impl
 
-import arrow.core.None
-import arrow.core.Option
-import arrow.core.Some
-import arrow.core.getOrElse
-import arrow.core.toOption
 import com.toasttab.protokt.codegen.impl.Deprecation.renderOptions
 import com.toasttab.protokt.codegen.impl.Implements.overrides
 import com.toasttab.protokt.codegen.impl.Nullability.deserializeType
@@ -101,7 +96,7 @@ private constructor(
             field = f,
             any =
                 if (f.map) {
-                    typeParams(f.protoTypeName)
+                    resolveMapEntryTypes(f, ctx)
                 } else {
                     interceptTypeName(
                         f,
@@ -111,25 +106,6 @@ private constructor(
                 }
         )
 
-    private fun typeParams(n: String) =
-        findType(n, msg)
-            .map { resolveMapEntryTypes(it, ctx) }
-            .getOrElse { error("missing type params") }
-
-    private fun findType(
-        tn: String,
-        msg: Message
-    ): Option<Message> {
-        val n = tn.split(".").let { if (it.isEmpty()) tn else it.last() }
-        return msg.nestedTypes.find {
-            when (it) {
-                is Message ->
-                    if (it.name == n) Some(it) else findType(n, it)
-                else -> None
-            }.isDefined()
-        }.toOption().map { f -> f as Message }
-    }
-
     private fun Field.defaultValue(ctx: Context) =
         when (this) {
             is StandardField ->

diff --git a/protokt-codegen/src/main/kotlin/com/toasttab/protokt/codegen/impl/SizeofAnnotator.kt b/protokt-codegen/src/main/kotlin/com/toasttab/protokt/codegen/impl/SizeofAnnotator.kt
@@ -23,6 +23,8 @@ import com.toasttab.protokt.codegen.impl.STAnnotator.Context
 import com.toasttab.protokt.codegen.impl.Wrapper.interceptFieldSizeof
 import com.toasttab.protokt.codegen.impl.Wrapper.interceptSizeof
 import com.toasttab.protokt.codegen.impl.Wrapper.interceptValueAccess
+import com.toasttab.protokt.codegen.impl.Wrapper.mapKeyConverter
+import com.toasttab.protokt.codegen.impl.Wrapper.mapValueConverter
 import com.toasttab.protokt.codegen.protoc.Message
 import com.toasttab.protokt.codegen.protoc.Oneof
 import com.toasttab.protokt.codegen.protoc.StandardField
@@ -87,7 +89,10 @@ private constructor(
                 Options(
                     fieldSizeof = interceptFieldSizeof(f, name, ctx),
                     fieldAccess =
-                        interceptValueAccess(f, ctx, IterationVar.render())
+                        interceptValueAccess(f, ctx, IterationVar.render()),
+                    keyAccess = mapKeyConverter(f, ctx),
+                    valueAccess = mapValueConverter(f, ctx),
+                    valueType = f.mapEntry?.value?.type
                 )
         )
     }

diff --git a/protokt-codegen/src/main/kotlin/com/toasttab/protokt/codegen/impl/StandardFieldExt.kt b/protokt-codegen/src/main/kotlin/com/toasttab/protokt/codegen/impl/StandardFieldExt.kt
@@ -17,6 +17,8 @@ package com.toasttab.protokt.codegen.impl
 
 import com.toasttab.protokt.codegen.impl.STAnnotator.Context
 import com.toasttab.protokt.codegen.impl.Wrapper.interceptValueAccess
+import com.toasttab.protokt.codegen.impl.Wrapper.mapKeyConverter
+import com.toasttab.protokt.codegen.impl.Wrapper.mapValueConverter
 import com.toasttab.protokt.codegen.protoc.StandardField
 import com.toasttab.protokt.codegen.protoc.Tag
 import com.toasttab.protokt.codegen.template.Renderers.Box
@@ -69,7 +71,12 @@ internal fun StandardField.nonDefault(ctx: Context) =
 internal fun StandardField.boxMap(ctx: Context) =
     BoxMap.render(
         type = type,
-        box = unqualifiedNestedTypeName(ctx)
+        box = unqualifiedNestedTypeName(ctx),
+        options = BoxMap.Options(
+            keyWrap = mapKeyConverter(this, ctx),
+            valueWrap = mapValueConverter(this, ctx),
+            valueType = mapEntry!!.value.type
+        )
     )
 
 internal fun StandardField.box(s: String) =

diff --git a/...-codegen/src/main/kotlin/com/toasttab/protokt/codegen/impl/StandardFieldImportResolver.kt b/...-codegen/src/main/kotlin/com/toasttab/protokt/codegen/impl/StandardFieldImportResolver.kt
@@ -18,7 +18,7 @@ package com.toasttab.protokt.codegen.impl
 import com.github.andrewoma.dexx.kollection.ImmutableSet
 import com.github.andrewoma.dexx.kollection.immutableSetOf
 import com.toasttab.protokt.codegen.impl.Wrapper.converter
-import com.toasttab.protokt.codegen.impl.Wrapper.foldWrap
+import com.toasttab.protokt.codegen.impl.Wrapper.foldFieldWrap
 import com.toasttab.protokt.codegen.model.Import
 import com.toasttab.protokt.codegen.model.PPackage
 import com.toasttab.protokt.codegen.model.pclass
@@ -81,7 +81,7 @@ class StandardFieldImportResolver(
             set.add(Import.Class(f.typePClass))
         }
 
-        f.foldWrap(
+        f.foldFieldWrap(
             pkg,
             ctx,
             { },