From f80c0244016fd61e030649d7033d54e26dbd2c6e Mon Sep 17 00:00:00 2001
From: Rob Norris <rob_norris@mac.com>
Date: Thu, 29 Aug 2024 12:03:12 -0500
Subject: [PATCH] make effect elaborator respect predicated mappings

---
 modules/core/src/main/scala/compiler.scala | 15 +++++++--------
 modules/core/src/main/scala/mapping.scala  | 17 +++++------------
 2 files changed, 12 insertions(+), 20 deletions(-)

diff --git a/modules/core/src/main/scala/compiler.scala b/modules/core/src/main/scala/compiler.scala
index 32fbca3d..d3332fa3 100644
--- a/modules/core/src/main/scala/compiler.scala
+++ b/modules/core/src/main/scala/compiler.scala
@@ -1243,7 +1243,8 @@ object QueryCompiler {
    * multiple effect handlers.
    *
    * This phase transforms the input query by assigning subtrees to effect
-   * handlers as specified by the supplied `emapping`.
+   * handlers as specified by the supplied `effects`, which takes a `Context` and
+   * fieldName, retuning an `EffectHandler` (if any).
    *
    * The mapping has `Type` and field name pairs as keys and effect handlers
    * as values. When the traversal of the input query visits a `Select` node
@@ -1254,21 +1255,18 @@ object QueryCompiler {
    *    and evaluating the subquery against its result.
    * 2. the subquery which will be evaluated by the effect handler.
    */
-  class EffectElaborator[F[_]] private (emapping: Map[(Type, String), EffectHandler[F]]) extends Phase {
+  class EffectElaborator[F[_]] private (effects: (Context, String) => Option[EffectHandler[F]]) extends Phase {
     override def transform(query: Query): Elab[Query] =
       query match {
         case s@Select(fieldName, resultName, child) =>
           for {
             c        <- Elab.context
-            obj      <- Elab.liftR(c.tpe.underlyingObject.toResultOrError(s"Type ${c.tpe} is not an object or interface type"))
             childCtx =  c.forFieldOrAttribute(fieldName, resultName)
             _        <- Elab.push(childCtx, child)
             ec       <- transform(child)
             _        <- Elab.pop
-            schema   <- Elab.schema
-            ref      =  schema.uncheckedRef(obj)
           } yield
-            emapping.get((ref, fieldName)) match {
+            effects(c, fieldName) match {
               case Some(handler) =>
                 Select(fieldName, resultName, Effect(handler, s.copy(child = ec)))
               case None =>
@@ -1282,8 +1280,9 @@ object QueryCompiler {
   object EffectElaborator {
     case class EffectMapping[F[_]](tpe: TypeRef, fieldName: String, handler: EffectHandler[F])
 
-    def apply[F[_]](mappings: Seq[EffectMapping[F]]): EffectElaborator[F] =
-      new EffectElaborator(mappings.map(m => ((m.tpe, m.fieldName), m.handler)).toMap)
+    def apply[F[_]](effects: (Context, String) => Option[EffectHandler[F]]): EffectElaborator[F] =
+      new EffectElaborator(effects)
+
   }
 
   /**
diff --git a/modules/core/src/main/scala/mapping.scala b/modules/core/src/main/scala/mapping.scala
index 6e610235..c1b552a4 100644
--- a/modules/core/src/main/scala/mapping.scala
+++ b/modules/core/src/main/scala/mapping.scala
@@ -1188,19 +1188,12 @@ abstract class Mapping[F[_]] {
     ComponentElaborator(componentMappings)
   }
 
-  lazy val effectElaborator = {
-    val effectMappings =
-      typeMappings.mappings.flatMap {
-        case om: ObjectMapping =>
-          om.fieldMappings.collect {
-            case EffectField(fieldName, handler, _, _) =>
-              EffectElaborator.EffectMapping(schema.uncheckedRef(om.tpe), fieldName, handler)
-          }
-        case _ => Seq.empty
+  lazy val effectElaborator: EffectElaborator[F] =
+    EffectElaborator { (ctx, fieldName) =>
+      typeMappings.fieldMapping(ctx, fieldName).collect {
+        case e: EffectField => e.handler
       }
-
-    EffectElaborator(effectMappings)
-  }
+    }
 
   def compilerPhases: List[QueryCompiler.Phase] = List(selectElaborator, componentElaborator, effectElaborator)