Enable codecs re-using or discarding parameters

Author: s5bug

core/shared/src/main/scala/porcupine/codec.scala

@@ -27,7 +27,7 @@ import scala.deriving.Mirror
 trait Encoder[A]:
   outer =>
 
-  def parameters: Int
+  def parameters: State[Int, List[Int]]
 
   def encode(a: A): List[LiteValue]
 
@@ -53,7 +53,7 @@ object Encoder:
 
     def product[A, B](fa: Encoder[A], fb: Encoder[B]) = new:
       def parameters =
-        fa.parameters + fb.parameters
+        (fa.parameters, fb.parameters).mapN(_ ++ _)
 
       def encode(ab: (A, B)) =
         val (a, b) = ab
@@ -130,7 +130,7 @@ object Codec:
       apply: T => LiteValue,
       unapply: PartialFunction[LiteValue, T],
   ) extends Codec[T] {
-    override def parameters: Int = 1
+    override def parameters: State[Int, List[Int]] = State(idx => (idx + 1, List(idx)))
     override def encode(a: T): List[LiteValue] = apply(a) :: Nil
     override def decode: StateT[Either[Throwable, *], List[LiteValue], T] = StateT {
       case unapply(l) :: tail => Right((tail, l))
@@ -154,7 +154,7 @@ object Codec:
     new Simple("NULL", _ => LiteValue.Null, { case LiteValue.Null => None })
 
   def unit: Codec[Unit] = new:
-    def parameters: Int = 0
+    def parameters = State.pure(List.empty)
     def encode(u: Unit) = Nil
     def decode = StateT.pure(())
 
@@ -165,7 +165,7 @@ object Codec:
 
     def product[A, B](fa: Codec[A], fb: Codec[B]) = new:
       def parameters =
-        fa.parameters + fb.parameters
+        (fa.parameters, fb.parameters).mapN(_ ++ _)
 
       def encode(ab: (A, B)) =
         val (a, b) = ab

core/shared/src/main/scala/porcupine/sql.scala

@@ -35,41 +35,61 @@ object Query:
       Query(fab.sql, fab.encoder.contramap(f), fab.decoder.map(g))
 
 final class Fragment[A](
-    val parts: List[Either[String, Int]],
+    val part: Fragment.Part,
     val encoder: Encoder[A],
 ):
-  def sql: String = parts.foldMap {
-    case Left(s) => s
-    case Right(i) => ("?, " * (i - 1)) ++ "?"
-  }
+  def sql: String = part.compile.runA(1).value
 
   def command: Query[A, Unit] = Query(sql, encoder, Codec.unit)
 
   def query[B](decoder: Decoder[B]): Query[A, B] = Query(sql, encoder, decoder)
 
-  def apply(a: A): Fragment[Unit] = Fragment(parts, encoder.contramap(_ => a))
+  def apply(a: A): Fragment[Unit] = Fragment(part, encoder.contramap(_ => a))
 
   def stripMargin: Fragment[A] = stripMargin('|')
 
   def stripMargin(marginChar: Char): Fragment[A] =
-    val head = parts.headOption
-    val tail = parts.tail
-    val ps = head.map {
-      _.leftMap(_.stripMargin(marginChar))
-    }.toList ++ tail.map {
-      _.leftMap(str =>
-        str.takeWhile(_ != '\n') + str.dropWhile(_ != '\n').stripMargin(marginChar),
-      )
-    }
-    Fragment(ps, encoder)
+    Fragment(part.stripMargin(true, marginChar), encoder)
 
 object Fragment:
+  sealed trait Part:
+    def compile: State[Int, String]
+    def concatenate(other: Part): Part = other match {
+      case Part.Concatenate(values) => Part.Concatenate(this :: values)
+      case _ => Part.Concatenate(List(this, other))
+    }
+    def stripMargin(head: Boolean, marginChar: Char): Part
+
+  object Part:
+    final case class Literal(x: String) extends Part:
+      def compile = State.pure(x)
+      def stripMargin(head: Boolean, marginChar: Char) =
+        if (head) Literal(x.stripMargin(marginChar))
+        else Literal(x.takeWhile(_ != '\n') ++ x.dropWhile(_ != '\n').stripMargin(marginChar))
+    final case class Concatenate(values: List[Part]) extends Part:
+      def compile = values.traverse(_.compile).map(_.combineAll)
+      override def concatenate(other: Part) = other match {
+        case Concatenate(values) => Concatenate(this.values ++ values)
+        case _ => Concatenate(this.values :+ other)
+      }
+      def stripMargin(head: Boolean, marginChar: Char) =
+        values match {
+          case h :: t =>
+            Concatenate(
+              h.stripMargin(head, marginChar) :: t.map(_.stripMargin(false, marginChar)),
+            )
+          case other => this
+        }
+    final case class Parameters(advance: State[Int, List[Int]]) extends Part:
+      def compile = advance.map(_.map(idx => s"?$idx").mkString(", "))
+      def stripMargin(head: Boolean, marginChar: Char) = this
+
   given ContravariantMonoidal[Fragment] = new:
-    val unit = Fragment(List.empty, Codec.unit)
+    val unit = Fragment(Part.Concatenate(List.empty), Codec.unit)
     def product[A, B](fa: Fragment[A], fb: Fragment[B]) =
-      Fragment(fa.parts ++ fb.parts, (fa.encoder, fb.encoder).tupled)
+      Fragment(fa.part.concatenate(fb.part), (fa.encoder, fb.encoder).tupled)
     def contramap[A, B](fa: Fragment[A])(f: B => A) =
-      Fragment(fa.parts, fa.encoder.contramap(f))
+      Fragment(fa.part, fa.encoder.contramap(f))
 
   given Monoid[Fragment[Unit]] = new:
     def empty = ContravariantMonoidal[Fragment].unit
@@ -93,13 +113,13 @@ private def sqlImpl(
 
   // TODO appending to `List` is slow
   val fragment =
-    parts.zipAll(args, '{ "" }, '{ "" }).foldLeft('{ List.empty[Either[String, Int]] }) {
-      case ('{ $acc: List[Either[String, Int]] }, ('{ $p: String }, '{ $s: String })) =>
-        '{ $acc :+ Left($p) :+ Left($s) }
-      case ('{ $acc: List[Either[String, Int]] }, ('{ $p: String }, '{ $e: Encoder[t] })) =>
-        '{ $acc :+ Left($p) :+ Right($e.parameters) }
-      case ('{ $acc: List[Either[String, Int]] }, ('{ $p: String }, '{ $f: Fragment[t] })) =>
-        '{ $acc :+ Left($p) :++ $f.parts }
+    parts.zipAll(args, '{ "" }, '{ "" }).foldLeft('{ List.empty[Fragment.Part] }) {
+      case ('{ $acc: List[Fragment.Part] }, ('{ $p: String }, '{ $s: String })) =>
+        '{ $acc :+ Fragment.Part.Literal($p) :+ Fragment.Part.Literal($s) }
+      case ('{ $acc: List[Fragment.Part] }, ('{ $p: String }, '{ $e: Encoder[t] })) =>
+        '{ $acc :+ Fragment.Part.Literal($p) :+ Fragment.Part.Parameters($e.parameters) }
+      case ('{ $acc: List[Fragment.Part] }, ('{ $p: String }, '{ $f: Fragment[t] })) =>
+        '{ $acc :+ Fragment.Part.Literal($p) :+ $f.part }
     }
 
   val encoder = args.collect {
@@ -125,5 +145,6 @@ private def sqlImpl(
       }
 
   (fragment, encoder) match
-    case ('{ $s: List[Either[String, Int]] }, '{ $e: Encoder[a] }) => '{ Fragment[a]($s, $e) }
+    case ('{ $s: List[Fragment.Part] }, '{ $e: Encoder[a] }) =>
+      '{ Fragment[a](Fragment.Part.Concatenate($s), $e) }
     case _ => sys.error("porcupine pricked itself")