From 6e060d278d233e9f98267be8f1437788663a7a2c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=C3=98yvind=20Raddum=20Berg?= <elacin@gmail.com>
Date: Thu, 4 Jul 2024 11:09:29 +0200
Subject: [PATCH] Handle multiple user-defined rows of same type in a row

---
 typo/src/scala/typo/internal/IdComputed.scala           | 5 +++--
 typo/src/scala/typo/internal/codegen/DbLibAnorm.scala   | 6 +++---
 typo/src/scala/typo/internal/codegen/DbLibDoobie.scala  | 8 +++-----
 typo/src/scala/typo/internal/codegen/DbLibZioJdbc.scala | 7 +++----
 4 files changed, 12 insertions(+), 14 deletions(-)

diff --git a/typo/src/scala/typo/internal/IdComputed.scala b/typo/src/scala/typo/internal/IdComputed.scala
index 5b35e202e5..186eccabca 100644
--- a/typo/src/scala/typo/internal/IdComputed.scala
+++ b/typo/src/scala/typo/internal/IdComputed.scala
@@ -6,8 +6,9 @@ sealed trait IdComputed {
   def cols: NonEmptyList[ComputedColumn]
   def tpe: sc.Type
   final def param: sc.Param = sc.Param(paramName, tpe, None)
-  final def userDefinedCols: List[ComputedColumn] =
-    cols.toList.collect { case x if sc.Type.containsUserDefined(x.tpe) => x }
+
+  final lazy val userDefinedColTypes: List[sc.Type] =
+    cols.toList.collect { case x if sc.Type.containsUserDefined(x.tpe) => x }.map(_.tpe).distinctBy(sc.Type.base)
 }
 
 object IdComputed {
diff --git a/typo/src/scala/typo/internal/codegen/DbLibAnorm.scala b/typo/src/scala/typo/internal/codegen/DbLibAnorm.scala
index c320f99726..03196afb12 100644
--- a/typo/src/scala/typo/internal/codegen/DbLibAnorm.scala
+++ b/typo/src/scala/typo/internal/codegen/DbLibAnorm.scala
@@ -168,11 +168,11 @@ class DbLibAnorm(pkg: sc.QIdent, inlineImplicits: Boolean, default: ComputedDefa
       case RepoMethod.SelectById(_, _, id, rowType) =>
         code"def $name(${id.param})(implicit c: ${TypesJava.Connection}): ${TypesScala.Option.of(rowType)}"
       case RepoMethod.SelectByIds(_, _, idComputed, idsParam, rowType) =>
-        val usedDefineds = idComputed.userDefinedCols.zipWithIndex.map { case (col, i) => sc.Param(sc.Ident(s"toStatement$i"), ToStatement.of(sc.Type.ArrayOf(col.tpe)), None) }
+        val usedDefineds = idComputed.userDefinedColTypes.zipWithIndex.map { case (colType, i) => sc.Param(sc.Ident(s"toStatement$i"), ToStatement.of(sc.Type.ArrayOf(colType)), None) }
         val params = sc.Param(sc.Ident("c"), TypesJava.Connection, None) :: usedDefineds
         code"def $name($idsParam)(implicit ${params.map(_.code).mkCode(", ")}): ${TypesScala.List.of(rowType)}"
       case RepoMethod.SelectByIdsTracked(x) =>
-        val usedDefineds = x.idComputed.userDefinedCols.zipWithIndex.map { case (col, i) => sc.Param(sc.Ident(s"toStatement$i"), ToStatement.of(sc.Type.ArrayOf(col.tpe)), None) }
+        val usedDefineds = x.idComputed.userDefinedColTypes.zipWithIndex.map { case (colType, i) => sc.Param(sc.Ident(s"toStatement$i"), ToStatement.of(sc.Type.ArrayOf(colType)), None) }
         val params = sc.Param(sc.Ident("c"), TypesJava.Connection, None) :: usedDefineds
         code"def $name(${x.idsParam})(implicit ${params.map(_.code).mkCode(", ")}): ${TypesScala.Map.of(x.idComputed.tpe, x.rowType)}"
       case RepoMethod.SelectByUnique(_, keyColumns, _, rowType) =>
@@ -200,7 +200,7 @@ class DbLibAnorm(pkg: sc.QIdent, inlineImplicits: Boolean, default: ComputedDefa
       case RepoMethod.Delete(_, id) =>
         code"def $name(${id.param})(implicit c: ${TypesJava.Connection}): ${TypesScala.Boolean}"
       case RepoMethod.DeleteByIds(_, idComputed, idsParam) =>
-        val usedDefineds = idComputed.userDefinedCols.zipWithIndex.map { case (col, i) => sc.Param(sc.Ident(s"toStatement$i"), ToStatement.of(sc.Type.ArrayOf(col.tpe)), None) }
+        val usedDefineds = idComputed.userDefinedColTypes.zipWithIndex.map { case (colType, i) => sc.Param(sc.Ident(s"toStatement$i"), ToStatement.of(sc.Type.ArrayOf(colType)), None) }
         val params = sc.Param(sc.Ident("c"), TypesJava.Connection, None) :: usedDefineds
         code"def $name($idsParam)(implicit ${params.map(_.code).mkCode(", ")}): ${TypesScala.Int}"
       case RepoMethod.SqlFile(sqlScript) =>
diff --git a/typo/src/scala/typo/internal/codegen/DbLibDoobie.scala b/typo/src/scala/typo/internal/codegen/DbLibDoobie.scala
index 38006c7fbe..7579232a75 100644
--- a/typo/src/scala/typo/internal/codegen/DbLibDoobie.scala
+++ b/typo/src/scala/typo/internal/codegen/DbLibDoobie.scala
@@ -71,7 +71,7 @@ class DbLibDoobie(pkg: sc.QIdent, inlineImplicits: Boolean, default: ComputedDef
       case RepoMethod.SelectById(_, _, id, rowType) =>
         code"def $name(${id.param}): ${ConnectionIO.of(TypesScala.Option.of(rowType))}"
       case RepoMethod.SelectByIds(_, _, idComputed, idsParam, rowType) =>
-        val usedDefineds = idComputed.userDefinedCols.zipWithIndex.map { case (col, i) => sc.Param(sc.Ident(s"puts$i"), Put.of(sc.Type.ArrayOf(col.tpe)), None) }
+        val usedDefineds = idComputed.userDefinedColTypes.zipWithIndex.map { case (colType, i) => sc.Param(sc.Ident(s"puts$i"), Put.of(sc.Type.ArrayOf(colType)), None) }
         usedDefineds match {
           case Nil =>
             code"def $name($idsParam): ${fs2Stream.of(ConnectionIO, rowType)}"
@@ -79,7 +79,7 @@ class DbLibDoobie(pkg: sc.QIdent, inlineImplicits: Boolean, default: ComputedDef
             code"def $name($idsParam)(implicit ${nonEmpty.map(_.code).mkCode(", ")}): ${fs2Stream.of(ConnectionIO, rowType)}"
         }
       case RepoMethod.SelectByIdsTracked(x) =>
-        val usedDefineds = x.idComputed.userDefinedCols.zipWithIndex.map { case (col, i) => sc.Param(sc.Ident(s"puts$i"), Put.of(sc.Type.ArrayOf(col.tpe)), None) }
+        val usedDefineds = x.idComputed.userDefinedColTypes.zipWithIndex.map { case (colType, i) => sc.Param(sc.Ident(s"puts$i"), Put.of(sc.Type.ArrayOf(colType)), None) }
         val returnType = ConnectionIO.of(TypesScala.Map.of(x.idComputed.tpe, x.rowType))
         usedDefineds match {
           case Nil =>
@@ -113,9 +113,7 @@ class DbLibDoobie(pkg: sc.QIdent, inlineImplicits: Boolean, default: ComputedDef
       case RepoMethod.Delete(_, id) =>
         code"def $name(${id.param}): ${ConnectionIO.of(TypesScala.Boolean)}"
       case RepoMethod.DeleteByIds(_, idComputed, idsParam) =>
-        val usedDefineds =
-          idComputed.userDefinedCols.zipWithIndex
-            .map { case (col, i) => sc.Param(sc.Ident(s"put$i"), Put.of(sc.Type.ArrayOf(col.tpe)), None) }
+        val usedDefineds = idComputed.userDefinedColTypes.zipWithIndex.map { case (colType, i) => sc.Param(sc.Ident(s"put$i"), Put.of(sc.Type.ArrayOf(colType)), None) }
         usedDefineds match {
           case Nil =>
             code"def $name(${idsParam}): ${ConnectionIO.of(TypesScala.Int)}"
diff --git a/typo/src/scala/typo/internal/codegen/DbLibZioJdbc.scala b/typo/src/scala/typo/internal/codegen/DbLibZioJdbc.scala
index ebbd9c94e8..1f9594d729 100644
--- a/typo/src/scala/typo/internal/codegen/DbLibZioJdbc.scala
+++ b/typo/src/scala/typo/internal/codegen/DbLibZioJdbc.scala
@@ -188,7 +188,7 @@ class DbLibZioJdbc(pkg: sc.QIdent, inlineImplicits: Boolean, dslEnabled: Boolean
       case RepoMethod.SelectById(_, _, id, rowType) =>
         code"def $name(${id.param}): ${ZIO.of(ZConnection, Throwable, TypesScala.Option.of(rowType))}"
       case RepoMethod.SelectByIds(_, _, idComputed, idsParam, rowType) =>
-        val usedDefineds = idComputed.userDefinedCols.zipWithIndex.map { case (col, i) => sc.Param(sc.Ident(s"encoder$i"), JdbcEncoder.of(sc.Type.ArrayOf(col.tpe)), None) }
+        val usedDefineds = idComputed.userDefinedColTypes.zipWithIndex.map { case (colType, i) => sc.Param(sc.Ident(s"encoder$i"), JdbcEncoder.of(sc.Type.ArrayOf(colType)), None) }
 
         usedDefineds match {
           case Nil =>
@@ -197,7 +197,7 @@ class DbLibZioJdbc(pkg: sc.QIdent, inlineImplicits: Boolean, dslEnabled: Boolean
             code"def $name($idsParam)(implicit ${nonEmpty.map(_.code).mkCode(", ")}): ${ZStream.of(ZConnection, Throwable, rowType)}"
         }
       case RepoMethod.SelectByIdsTracked(x) =>
-        val usedDefineds = x.idComputed.userDefinedCols.zipWithIndex.map { case (col, i) => sc.Param(sc.Ident(s"encoder$i"), JdbcEncoder.of(sc.Type.ArrayOf(col.tpe)), None) }
+        val usedDefineds = x.idComputed.userDefinedColTypes.zipWithIndex.map { case (colType, i) => sc.Param(sc.Ident(s"encoder$i"), JdbcEncoder.of(sc.Type.ArrayOf(colType)), None) }
         val returnType = ZIO.of(ZConnection, Throwable, TypesScala.Map.of(x.idComputed.tpe, x.rowType))
         usedDefineds match {
           case Nil =>
@@ -235,8 +235,7 @@ class DbLibZioJdbc(pkg: sc.QIdent, inlineImplicits: Boolean, dslEnabled: Boolean
       case RepoMethod.Delete(_, id) =>
         code"def $name(${id.param}): ${ZIO.of(ZConnection, Throwable, TypesScala.Boolean)}"
       case RepoMethod.DeleteByIds(_, idComputed, idsParam) =>
-        val usedDefineds = idComputed.userDefinedCols.zipWithIndex
-          .map { case (col, i) => sc.Param(sc.Ident(s"encoder$i"), JdbcEncoder.of(sc.Type.ArrayOf(col.tpe)), None) }
+        val usedDefineds = idComputed.userDefinedColTypes.zipWithIndex.map { case (colType, i) => sc.Param(sc.Ident(s"encoder$i"), JdbcEncoder.of(sc.Type.ArrayOf(colType)), None) }
         usedDefineds match {
           case Nil =>
             code"def $name(${idsParam}): ${ZIO.of(ZConnection, Throwable, TypesScala.Long)}"