getkyo · ahoy-jon · May 19, 2025
diff --git a/kyo-direct/shared/src/main/scala/kyo/Validate.scala b/kyo-direct/shared/src/main/scala/kyo/Validate.scala
@@ -16,6 +16,14 @@ private[kyo] object Validate:
                 case _                                    => false
             }
 
+        def validShift(tree: Tree): Boolean =
+            Trees.exists(expr.asTerm) {
+                case Apply(TypeApply(Select(s, methodName), _), List(Block(List(DefDef(_, _, _, Some(body))), _))) if body eq tree =>
+                    s.tpe <:< TypeRepr.of[Seq[?]] && methodName == "map"
+            }
+
+        end validShift
+
         Trees.traverse(expr.asTerm) {
             case Apply(TypeApply(Ident("now" | "later"), _), List(qual)) =>
                 Trees.traverse(qual) {
@@ -126,7 +134,7 @@ private[kyo] object Validate:
                     |}""".stripMargin)}""".stripMargin.stripMargin
                 )
 
-            case tree @ DefDef(_, _, _, Some(body)) if !pure(body) =>
+            case tree @ DefDef(_, _, _, Some(body)) if !pure(body) && !validShift(body) =>
                 fail(
                     tree,
                     s"""Method definitions containing ${".now".cyan} are not supported inside ${"`defer`".yellow} blocks.

diff --git a/kyo-direct/shared/src/test/scala/kyo/ShiftTest.scala b/kyo-direct/shared/src/test/scala/kyo/ShiftTest.scala
@@ -0,0 +1,24 @@
+package kyo
+
+import org.scalatest.Assertions
+import org.scalatest.freespec.AnyFreeSpec
+
+class ShiftTest extends AnyFreeSpec with Assertions:
+    "basic" in {
+        val x1: Seq[Int < Any] = Seq[Int < Any](1, 2, 3)
+        val x2: Seq[Int] < Any = defer:
+            x1.map(_.now)
+
+        assert(x2.eval == Seq(1, 2, 3))
+    }
+
+    "error" in {
+        val x1: Seq[Int < Abort[String]] = Seq(1, Abort.fail("error"), 3)
+        val x2: Seq[Int] < Abort[String] = defer:
+            x1.map(_.now)
+
+        Abort.run(x2).map: result =>
+            assert(result == Result.fail("error"))
+    }
+
+end ShiftTest