From dc4823cd9817769d84070f66adfe3e8b8fa9a2d1 Mon Sep 17 00:00:00 2001
From: Valdemar Grange <randomvald0069@gmail.com>
Date: Fri, 12 Jan 2024 05:08:43 +0100
Subject: [PATCH] work

---
 .../server/interpreter/AlphaRenaming.scala    |  3 +-
 .../server/interpreter/SubgraphBatches.scala  |  6 +-
 .../server/src/test/scala/gql/BatchTest.scala | 77 +++++++++++++++++++
 3 files changed, 84 insertions(+), 2 deletions(-)
 create mode 100644 modules/server/src/test/scala/gql/BatchTest.scala

diff --git a/modules/server/src/main/scala/gql/server/interpreter/AlphaRenaming.scala b/modules/server/src/main/scala/gql/server/interpreter/AlphaRenaming.scala
index a64291683..99d11d930 100644
--- a/modules/server/src/main/scala/gql/server/interpreter/AlphaRenaming.scala
+++ b/modules/server/src/main/scala/gql/server/interpreter/AlphaRenaming.scala
@@ -19,7 +19,8 @@ object AlphaRenaming {
         case alg: InlineBatch[F, k, v] => now(InlineBatch[F, k, v](alg.run, alg.sei.alpha(scope)))
         case alg: Batch[F, k, v]       => now(Batch[F, k, v](alg.id, alg.ubi.alpha(scope)))
         case alg: Compose[F, A, a, B] =>
-          now(Compose[F, A, a, B](alg.nodeId.alpha(scope), alg.left, alg.right))
+          (alphaStep(scope, alg.left), alphaStep(scope, alg.right))
+            .mapN(Compose[F, A, a, B](alg.nodeId.alpha(scope), _, _))
         case alg: Choose[f, i, a, b, c] =>
           (
             alphaStep(scope, alg.fac),
diff --git a/modules/server/src/main/scala/gql/server/interpreter/SubgraphBatches.scala b/modules/server/src/main/scala/gql/server/interpreter/SubgraphBatches.scala
index d68381866..339acf59c 100644
--- a/modules/server/src/main/scala/gql/server/interpreter/SubgraphBatches.scala
+++ b/modules/server/src/main/scala/gql/server/interpreter/SubgraphBatches.scala
@@ -137,7 +137,11 @@ object SubgraphBatches {
         .map(_.toMap)
 
     val batchLookup: Map[UniqueBatchInstance[?, ?], (SchemaState.BatchFunction[F, ?, ?], Int)] =
-      groups.map { case (k, _) => (k.uniqueNodeId, (schemaState.batchFunctions(k.batcherId), k.batcherId.id)) }.toMap
+      plan.tree.all.mapFilter { n =>
+        n.batchId.map { br =>
+          (br.uniqueNodeId, (schemaState.batchFunctions(br.batcherId), br.batcherId.id))
+        }
+      }.toMap
     def getBatchImpl[K, V](id: UniqueBatchInstance[K, V]) = {
       val (res, id0) = batchLookup(id)
       (res.asInstanceOf[SchemaState.BatchFunction[F, K, V]], id0)
diff --git a/modules/server/src/test/scala/gql/BatchTest.scala b/modules/server/src/test/scala/gql/BatchTest.scala
new file mode 100644
index 000000000..9d52139b8
--- /dev/null
+++ b/modules/server/src/test/scala/gql/BatchTest.scala
@@ -0,0 +1,77 @@
+package gql
+
+import munit.CatsEffectSuite
+import gql.ast._
+import gql.dsl.all._
+import cats.effect._
+import cats.implicits._
+import gql.resolver.Resolver
+
+class BatchTest extends CatsEffectSuite {
+  case class Batchers(
+      getString: Resolver[IO, Set[String], Map[String, String]],
+      getInt: Resolver[IO, Set[Int], Map[Int, Int]]
+  )
+
+  lazy val intInputs = IO.ref(List.empty[Set[Int]]).unsafeRunSync()
+  lazy val stringInputs = IO.ref(List.empty[Set[String]]).unsafeRunSync()
+
+  lazy val batchersF = (
+    Resolver.batch[IO, String, String](xs => IO(xs.map(x => (x, x)).toMap) <* stringInputs.update(xs :: _)),
+    Resolver.batch[IO, Int, Int](xs => IO(xs.map(x => (x, x)).toMap) <* intInputs.update(xs :: _))
+  ).mapN(Batchers.apply)
+
+  lazy val schema: Schema[IO, Unit, Unit, Unit] = Schema
+    .stateful {
+      batchersF.map { bs =>
+        case object OpaqueTpe
+        implicit lazy val opaqueTpe: Type[IO, OpaqueTpe.type] = builder[IO, OpaqueTpe.type] { b =>
+          b.tpe(
+            "Opaque",
+            "b1" -> b(_.as("hello") andThen bs.getString.opt),
+            "b2" -> b(_.as(1) andThen bs.getInt.opt)
+          )
+        }
+
+        SchemaShape.unit[IO](
+          builder[IO, Unit] { b =>
+            b.fields(
+              "opaque" -> b(_.as(OpaqueTpe)),
+              "b2" -> b(_.as(2) andThen bs.getInt.opt)
+            )
+          }
+        )
+      }
+    }
+    .unsafeRunSync()
+
+  test("should batch the two int resolvers") {
+    val query = """
+      query {
+        opaque {
+          b1
+          b2
+        }
+        b2
+      }
+    """
+
+    Compiler[IO]
+      .compile(schema, query)
+      .traverse {
+        case Application.Query(fa) =>
+          fa.flatMap { qr =>
+            assert(clue(qr.errors).isEmpty)
+            assertEquals(qr.data("b2").get.asNumber.get.toInt.get, 2)
+            val opq = qr.data("opaque").get.asObject.get
+            assertEquals(opq("b1").get.asString.get, "hello")
+            assertEquals(opq("b2").get.asNumber.get.toInt.get, 1)
+            (intInputs.get, stringInputs.get).mapN { (ii, si) =>
+              assertEquals(ii, List(Set(1, 2)))
+              assertEquals(si, List(Set("hello")))
+            }
+          }
+        case x => fail(s"unexpected program ${x.toString()}")
+      }
+  }
+}