Let map be a macro

Dsl/build.sbt

@@ -1,3 +1,5 @@
 libraryDependencies += "com.github.sbt" % "junit-interface" % "0.13.2" % Test
 
 libraryDependencies += "junit" % "junit" % "4.13.2" % Test
+
+libraryDependencies += "com.thoughtworks.dsl" %%% "bangnotation" % "2.0.0-M2+100-66bf7c43"

Dsl/src/main/scala/com/thoughtworks/dsl/Dsl.scala

@@ -481,6 +481,22 @@ object Dsl extends LowPriorityDsl0 {
       run(asKeyword(keyword))
     }
 
+  /** A marker trait that denotes a keyword class, enabling extension method
+    * defined in [[Dsl]] for subclasses of [[Keyword]].
+    */
+  trait Keyword extends Any
+
+  /** A marker trait that denotes a keyword opaque type, enabling extension
+    * method defined in [[Dsl]] for its subtypes of [[OpaqueKeyword]].
+    */
+  opaque type OpaqueKeyword = Any
+  object OpaqueKeyword {
+    opaque type Of[Self] <: Self & OpaqueKeyword = Self
+    object Of {
+      def apply[Self]: Self =:= Of[Self] = summon
+    }
+  }
+
   trait AsKeyword[From, Keyword, Value] extends (From => Keyword)
 
   object AsKeyword {
@@ -513,4 +529,23 @@ object Dsl extends LowPriorityDsl0 {
       dsl.cpsApply(keyword, handler)
     }
 
+  extension [From, Keyword, Value](inline from: From)(using inline asKeyword: Dsl.AsKeyword.SearchIsKeywordFirst[From, Keyword, Value])
+
+    inline def map[MappedValue](
+        mapper: Value => MappedValue
+    ): keywords.FlatMap[Keyword, Value, keywords.Pure[MappedValue]] =
+      keywords.FlatMap(asKeyword(from), keywords.Pure.apply.liftCo(mapper))
+
+    inline def flatMap[Mapped, MappedValue](
+        flatMapper: Value => Mapped
+    )(
+        using /*erased*/ AsKeyword.IsKeyword[Mapped, MappedValue]
+    ): keywords.FlatMap[Keyword, Value, Mapped] =
+      keywords.FlatMap(asKeyword(from), flatMapper)
+
+    inline def withFilter[Mapped, MappedValue](
+        filter: Value => Boolean
+    ) =
+      keywords.WithFilter(asKeyword(from), filter)
+
 }

bangnotation/src/main/scala/com/thoughtworks/dsl/bangnotation.scala

@@ -67,10 +67,31 @@ object bangnotation {
       }
     }
 
-    def resetDefDef(defDef: DefDef): DefDef = {
+    def reifyFunction[R, A](function: quoted.Expr[R => A])(using quoted.Type[R], quoted.Type[A]): quoted.Expr[R => Any] = {
+      val functionTerm = function.asTerm.underlyingArgument
+      functionTerm match {
+        case block @ qctx.reflect.Block(
+          List(
+            defDef: DefDef
+          ),
+          closure @ Closure(ident: Ident, _)
+        ) if (ident.name == defDef.name) =>
+          qctx.reflect.Block.copy(block)(
+            List(reifyDefDef(defDef)),
+            closure,
+          ).usingExpr { [F] => (f: quoted.Expr[F]) => (tf: quoted.Type[F]) =>
+            f.asInstanceOf[quoted.Expr[R => Any]]
+          }
+        case _ =>
+          report.error("Expect a function literal", functionTerm.underlyingArgument.pos)
+          '{???}
+      }
+    }
+
+    def reifyDefDef(defDef: DefDef): DefDef = {
       val DefDef(name, typeParamsAndParams, tpt, rhsOption) = defDef
       rhsOption match {
-        case Some(rhs) if resetDescendant =>
+        case Some(rhs) =>
           rhs match {
             case matchTree @ qctx.reflect.Match(scrutinee, cases) =>
               DefDef.copy(defDef)(
@@ -94,6 +115,40 @@ object bangnotation {
       }
     }
 
+    def resetDefDef(defDef: DefDef): DefDef = {
+      val DefDef(name, typeParamsAndParams, tpt, rhsOption) = defDef
+      rhsOption match {
+        case Some(rhs) if resetDescendant =>
+          rhs match {
+            case matchTree @ qctx.reflect.Match(scrutinee, cases) =>
+              DefDef.copy(defDef)(
+                name, typeParamsAndParams, tpt, Some(
+                  qctx.reflect.Match.copy(matchTree)(
+                    scrutinee,
+                    cases.map {
+                      case caseDef @ CaseDef(pattern, guard, caseRhs) =>
+                        CaseDef.copy(caseDef)(pattern, guard, reifyTerm(caseRhs))
+                    }
+                  )
+                )
+              )
+            case _ =>
+              DefDef.copy(defDef)(
+                name, typeParamsAndParams, tpt, Some(reifyTerm(rhs))
+              )
+          }
+        case _ =>
+          defDef
+      }
+    }
+
+    def reifyTerm(term: Term): Term = {
+      term.usingExpr { [Value] => (body: quoted.Expr[Value]) => (tv: quoted.Type[Value]) =>
+        given quoted.Type[Value] = tv
+        reify[Value](body).asTerm
+      }
+    }
+
     def resetTerm(term: Term): Term = {
       term.usingExpr { [Value] => (body: quoted.Expr[Value]) => (tv: quoted.Type[Value]) =>
         given quoted.Type[Value] = tv
@@ -686,7 +741,9 @@ object bangnotation {
     def reify[V](body: quoted.Expr[_])(using qctx: Quotes, tv: quoted.Type[V]): quoted.Expr[_] = {
       Macros[qctx.type](resetDescendant = false).reify[V](body/*.underlyingArgument*/)
     }
-
+    def reifyFunction[R, A](body: quoted.Expr[R => A])(using qctx: Quotes, tr: quoted.Type[R], ta: quoted.Type[A]): quoted.Expr[R => Any] = {
+      Macros[qctx.type](resetDescendant = false).reifyFunction[R, A](body/*.underlyingArgument*/)
+    }
     def reset[From, To](body: quoted.Expr[From])(using qctx: Quotes, fromType: quoted.Type[From], toType: quoted.Type[To]): quoted.Expr[To] = {
       import qctx.reflect.{_, given}
       val result: quoted.Expr[To] = Macros[qctx.type](resetDescendant = false).reset(body/*.underlyingArgument*/)
@@ -701,6 +758,10 @@ object bangnotation {
     Macros.reify[Value]('value)
   }
 
+  transparent inline def reifyFunction[State, Value](inline value: State => Value): State => Any = ${
+    Macros.reifyFunction[State, Value]('value)
+  }
+
   class *[Functor[_]] {
     inline def apply[Value](inline value: Value): Functor[Value] = ${
       Macros.reset[Value, Functor[Value]]('value)

bangnotation/src/test/scala/com/thoughtworks/dsl/keywords/ReturnSpec.scala

@@ -1,7 +1,7 @@
 package com.thoughtworks.dsl
 package keywords
 import com.thoughtworks.dsl.Dsl.!!
-import bangnotation.{_, given}
+import bangnotation._
 import utest.{TestSuite, Tests, given}
 import Dsl.Run
 import scala.language.implicitConversions

build.sbt

@@ -12,6 +12,7 @@ lazy val bangnotation =
       `keywords-Typed`,
       `keywords-Return`,
       `keywords-FlatMap`,
+      `keywords-WithFilter`,
       `keywords-Match`,
       `keywords-Suspend`,
       `keywords-Pure`,
@@ -39,7 +40,6 @@ lazy val `domains-Task` =
       `keywords-Each` % Test,
       `keywords-Using` % Test,
       `keywords-Yield` % Test,
-      comprehension % Test
     )
 
 lazy val `keywords-Fork` =
@@ -107,7 +107,7 @@ lazy val `keywords-Return` =
 lazy val `keywords-Continue` =
   crossProject(JSPlatform, JVMPlatform)
     .crossType(CrossType.Pure)
-    .dependsOn(Dsl, bangnotation % Test, `keywords-Each` % Test)
+    .dependsOn(Dsl)
 
 lazy val `keywords-Get` =
   crossProject(JSPlatform, JVMPlatform)
@@ -133,7 +133,6 @@ lazy val `keywords-AsynchronousIo` =
       `keywords-In` % Test,
       `keywords-Each` % Test,
       `keywords-Using` % Test,
-      comprehension % Test,
       `domains-Task` % Test
     )
 
@@ -183,7 +182,6 @@ lazy val `keywords-Await` =
     .dependsOn(
       Dsl,
       `domains-Continuation`,
-      comprehension % Test,
       bangnotation % Test,
       `domains-Task` % Test,
       `keywords-In` % Test,
@@ -214,21 +212,6 @@ lazy val `keywords-Monadic` =
     .crossType(CrossType.Pure)
     .dependsOn(Dsl, bangnotation % Test)
 
-lazy val comprehension =
-  crossProject(JSPlatform, JVMPlatform)
-    .crossType(CrossType.Pure)
-    .dependsOn(
-      `keywords-Map`,
-      `keywords-FlatMap`,
-      `keywords-WithFilter`,
-      `keywords-Return`,
-      `keywords-Yield` % Test,
-      `keywords-In` % Test,
-      `keywords-Using` % Test,
-      `keywords-Continue` % Test,
-      bangnotation % Test
-    )
-
 organization in ThisBuild := "com.thoughtworks.dsl"
 
 skip in publish := true

domains-Task/src/main/scala/com/thoughtworks/dsl/domains/Task.scala

@@ -25,7 +25,6 @@ import scala.util.control.TailCalls.TailRec
   *
   *          {{{
   *          import com.thoughtworks.dsl.Dsl._
-  *          import com.thoughtworks.dsl.comprehension._
   *          import com.thoughtworks.dsl.bangnotation._
   *          import com.thoughtworks.dsl._
   *          import com.thoughtworks.dsl.keywords._

keywords-AsynchronousIo/src/main/scala/com/thoughtworks/dsl/keywords/AsynchronousIo.scala

@@ -16,14 +16,13 @@ import scala.util.control.NonFatal
   *          {{{
   *          import java.nio._, file._, channels._
   *          import com.thoughtworks.dsl.bangnotation._
-  *          import com.thoughtworks.dsl.comprehension._
   *          import com.thoughtworks.dsl.domains.Task
   *          import com.thoughtworks.dsl.keywords._
   *          import com.thoughtworks.dsl.keywords.Shift._
   *          import com.thoughtworks.dsl.keywords.AsynchronousIo.ReadFile
   *          import scala.collection.mutable.ArrayBuffer
   *          import scala.io.Codec
-  *          def readAll(channel: AsynchronousFileChannel, temporaryBufferSize: Int = 4096): Task[ArrayBuffer[CharBuffer]] = *[Task] {
+  *          def readAll(channel: AsynchronousFileChannel, temporaryBufferSize: Int = 4096): Task[ArrayBuffer[CharBuffer]] = Task {
   *            val charBuffers = ArrayBuffer.empty[CharBuffer]
   *            val decoder = Codec.UTF8.decoder
   *            val byteBuffer = ByteBuffer.allocate(temporaryBufferSize)
@@ -37,19 +36,17 @@ import scala.util.control.NonFatal
   *            charBuffers
   *          }
   *          }}}
-  *
   *          `Task`s created from !-notation can be used in `for`-comprehension,
   *          and other keywords can be used together in the same `for` block.
   *
   *          For example, the following `cat` function contains a single `for` block to concatenate file contents.
   *          It asynchronously iterates elements `Seq`, `ArrayBuffer` and `String` with the help of [[keywords.Each]],
   *          managed native resources with the help of [[keywords.Using]],
   *          performs previously created `readAll` task with the help of [[keywords.Shift]],
-  *          and finally converts the return type [[comprehension.ComprehensionOps.as as]] a `Task[Vector[Char]]`.
+  *          and finally converts the return type [[bangnotation.Ops.as as]] a `Task[Vector[Char]]`.
   *
   *          {{{
   *          import com.thoughtworks.dsl._
-  *          import com.thoughtworks.dsl.bangnotation._
   *          import com.thoughtworks.dsl.keywords._
   *          import com.thoughtworks.dsl.keywords.Shift._
   *          import com.thoughtworks.dsl.domains.Task
@@ -81,7 +78,7 @@ import scala.util.control.NonFatal
   *          })
   *          }}}
   */
-trait AsynchronousIo[Value] extends Any {
+trait AsynchronousIo[Value] extends Any with Dsl.Keyword {
 
   /** Starts the asynchronous operations */
   protected def start[Attachment](attachment: Attachment, handler: CompletionHandler[Value, _ >: Attachment]): Unit

keywords-Await/src/main/scala/com/thoughtworks/dsl/keywords/Await.scala

@@ -106,9 +106,9 @@ import scala.language.implicitConversions
   * @author
   *   杨博 (Yang Bo)
   */
-opaque type Await[Result] = concurrent.Future[Result]
+opaque type Await[Result] <: Dsl.OpaqueKeyword = Dsl.OpaqueKeyword.Of[concurrent.Future[Result]]
 object Await {
-  @inline def apply[Result]: concurrent.Future[Result] =:= Await[Result] = summon
+  @inline def apply[Result]: concurrent.Future[Result] =:= Await[Result] = Dsl.OpaqueKeyword.Of.apply
   given [Result]: AsKeyword.IsKeyword[Await[Result], Result] with {}
 
   implicit def streamAwaitDsl[Value, That](implicit
@@ -143,6 +143,6 @@ object Await {
         !!.fromTryContinuation[Unit, Value](keyword.onComplete)(handler)
     }
 
-  given implicitAwait[Value]: AsKeyword.IsKeywordSubtype[Future[Value], Await[Value], Value] with {}
+  given implicitAwait[Value]: AsKeyword[Future[Value], Await[Value], Value] = Await(_)
 
 }

keywords-Await/src/test/scala/com/thoughtworks/dsl/keywords/AwaitTest.scala

@@ -1,6 +1,5 @@
 package com.thoughtworks.dsl
 
-import comprehension._
 import bangnotation._
 import Dsl._
 import keywords._, Match._

keywords-Continue/build.sbt

@@ -1,4 +1,4 @@
-enablePlugins(Example)
+// enablePlugins(Example)
 
 libraryDependencies += "org.scalatest" %%% "scalatest" % "3.2.10" % Test
 

keywords-Continue/src/main/scala/com/thoughtworks/dsl/keywords/Continue.scala

@@ -13,7 +13,7 @@ import scala.collection._
   *
   * @see The [[Continue$ Continue]] object, which is the only instance of this [[Continue]] class.
   */
-sealed class Continue
+sealed class Continue extends Dsl.Keyword
 
 /** A keyword to skip the current iteration in a collection comprehension block.
   *

keywords-Each/src/main/scala/com/thoughtworks/dsl/keywords/Each.scala

@@ -1,6 +1,6 @@
 package com.thoughtworks.dsl.keywords
 
-import com.thoughtworks.dsl.bangnotation.{reset, unary_!}
+import com.thoughtworks.dsl.bangnotation._
 import com.thoughtworks.dsl.Dsl
 import com.thoughtworks.dsl.Dsl.{!!, AsKeyword}
 
@@ -21,7 +21,7 @@ import scala.collection.mutable.Builder
   *          }}}
   * @see [[comprehension]] if you want to use traditional `for` comprehension instead of !-notation.
   */
-final case class Each[Element](elements: Traversable[Element])
+final case class Each[Element](elements: Traversable[Element]) extends Dsl.Keyword
 object Each {
   given [Element]: AsKeyword.IsKeyword[Each[Element], Element] with {}