Part 8 - implement Cross.Factory macro (TODO split shims out)

- fix scanning of Cross modules in Discover macro

Author: bishabosha

example/depth/cross/6-axes-extension/build.mill

@@ -56,6 +56,9 @@ trait FooModule3 extends FooModule2 with Cross.Module3[String, Int, Boolean] {
 
 > mill show foo3[b,2,false].param3
 error: ...object foo3 extends Cross[FooModule3](("a", 1), ("b", 2))
-error: ...                                      ^
-error: ...value _3 is not a member of (String, Int)
-*/
+error: ...                                      ^^^^^^^^
+error: ...expected at least 3 elements, got 2...
+error: ...object foo3 extends Cross[FooModule3](("a", 1), ("b", 2))
+error: ...                                                ^^^^^^^^
+error: ...expected at least 3 elements, got 2...
+*/

main/define/src/mill/define/Cross.scala

@@ -2,10 +2,10 @@ package mill.define
 
 import mill.api.{BuildScriptException, Lazy}
 
-import language.experimental.macros
 import scala.collection.mutable
 import scala.reflect.ClassTag
-import scala.reflect.macros.blackbox
+
+import scala.quoted.*
 
 object Cross {
 
@@ -136,100 +136,264 @@ object Cross {
      * expression of type `Any`, but type-checking on the macro- expanded code
      * provides some degree of type-safety.
      */
-    implicit def make[M <: Module[_]](t: Any): Factory[M] = ??? // macro makeImpl[M]
-    // def makeImpl[T: c.WeakTypeTag](c: blackbox.Context)(t: c.Expr[Any]): c.Expr[Factory[T]] = {
-    //   import c.universe._
-    //   val tpe = weakTypeOf[T]
-
-    //   if (!tpe.typeSymbol.isClass) {
-    //     c.abort(c.enclosingPosition, s"Cross type $tpe must be trait")
-    //   }
-
-    //   if (!tpe.typeSymbol.asClass.isTrait) abortOldStyleClass(c)(tpe)
-
-    //   val wrappedT = if (t.tree.tpe <:< typeOf[Seq[_]]) t.tree else q"_root_.scala.Seq($t)"
-    //   val v1 = c.freshName(TermName("v1"))
-    //   val ctx0 = c.freshName(TermName("ctx0"))
-    //   val concreteCls = c.freshName(TypeName(tpe.typeSymbol.name.toString))
-
-    //   val newTrees = collection.mutable.Buffer.empty[Tree]
-    //   var valuesTree: Tree = null
-    //   var pathSegmentsTree: Tree = null
-
-    //   val segments = q"_root_.mill.define.Cross.ToSegments"
-    //   if (tpe <:< typeOf[Module[_]]) {
-    //     newTrees.append(q"override def crossValue = $v1")
-    //     pathSegmentsTree = q"$segments($v1)"
-    //     valuesTree = q"$wrappedT.map(List(_))"
-    //   } else c.abort(
-    //     c.enclosingPosition,
-    //     s"Cross type $tpe must implement Cross.Module[T]"
-    //   )
-
-    //   if (tpe <:< typeOf[Module2[_, _]]) {
-    //     // For `Module2` and above, `crossValue` is no longer the entire value,
-    //     // but instead is just the first element of a tuple
-    //     newTrees.clear()
-    //     newTrees.append(q"override def crossValue = $v1._1")
-    //     newTrees.append(q"override def crossValue2 = $v1._2")
-    //     pathSegmentsTree = q"$segments($v1._1) ++ $segments($v1._2)"
-    //     valuesTree = q"$wrappedT.map(_.productIterator.toList)"
-    //   }
-
-    //   if (tpe <:< typeOf[Module3[_, _, _]]) {
-    //     newTrees.append(q"override def crossValue3 = $v1._3")
-    //     pathSegmentsTree = q"$segments($v1._1) ++ $segments($v1._2) ++ $segments($v1._3)"
-    //   }
-
-    //   if (tpe <:< typeOf[Module4[_, _, _, _]]) {
-    //     newTrees.append(q"override def crossValue4 = $v1._4")
-    //     pathSegmentsTree =
-    //       q"$segments($v1._1) ++ $segments($v1._2) ++ $segments($v1._3) ++ $segments($v1._4)"
-    //   }
-
-    //   if (tpe <:< typeOf[Module5[_, _, _, _, _]]) {
-    //     newTrees.append(q"override def crossValue5 = $v1._5")
-    //     pathSegmentsTree =
-    //       q"$segments($v1._1) ++ $segments($v1._2) ++ $segments($v1._3) ++ $segments($v1._4) ++ $segments($v1._5)"
-    //   }
-
-    //   // We need to create a `class $concreteCls` here, rather than just
-    //   // creating an anonymous sub-type of $tpe, because our task resolution
-    //   // logic needs to use java reflection to identify sub-modules and java
-    //   // reflect can only properly identify nested `object`s inside Scala
-    //   // `object` and `class`es.
-    //   val tree = q"""
-    //     new mill.define.Cross.Factory[$tpe](
-    //       makeList = $wrappedT.map{($v1: ${tq""}) =>
-    //         class $concreteCls()(implicit ctx: mill.define.Ctx) extends $tpe{..$newTrees}
-    //         (classOf[$concreteCls], ($ctx0: ${tq""}) => new $concreteCls()($ctx0))
-    //       },
-    //       crossSegmentsList = $wrappedT.map(($v1: ${tq""}) => $pathSegmentsTree ),
-    //       crossValuesListLists = $valuesTree,
-    //       crossValuesRaw = $wrappedT
-    //    ).asInstanceOf[${weakTypeOf[Factory[T]]}]
-    //   """
-
-    //   c.Expr[Factory[T]](tree)
-    // }
-
-    def abortOldStyleClass(c: blackbox.Context)(tpe: c.Type): Nothing = {
+    implicit inline def make[M <: Module[_]](inline t: Any): Factory[M] = ${ makeImpl[M]('t) }
+    def makeImpl[T: Type](using Quotes)(t: Expr[Any]): Expr[Factory[T]] = {
+      import quotes.reflect.*
+
+      val shims = ShimService.reflect
+
+      val tpe = TypeRepr.of[T]
+
+      val cls = tpe.classSymbol.getOrElse(
+        report.errorAndAbort(s"Cross type ${tpe.show} must be trait", Position.ofMacroExpansion)
+      )
+
+      if (!cls.flags.is(Flags.Trait)) abortOldStyleClass(tpe)
+
+      val wrappedT: Expr[Seq[Any]] = t match
+        case '{ $t1: Seq[elems] } => t1
+        case '{ $t1: t1 } => '{ Seq.apply($t1) }
+
+      val elems0: Type[?] = t match {
+        case '{ $t1: Seq[elems] } => TypeRepr.of[elems].widen.asType
+        case '{ $t1: elems } => TypeRepr.of[elems].widen.asType
+      }
+      val elemTypes: (Expr[Seq[Seq[Any]]], Seq[(Type[?], Expr[?] => Expr[?])]) = {
+        def select[T: Type](n: Int): Expr[?] => Expr[T] = {
+          elems0 match {
+            case '[type elems1 <: NonEmptyTuple; `elems1`] =>
+              arg => arg match {
+                case '{ $arg: `elems1` } =>
+                  '{ $arg.apply(${Expr(n)}) }.asExprOf[T]
+              }
+          }
+        }
+        def asSeq(tpe: Type[?], n: Int): Seq[(Type[?], Expr[?] => Expr[?])] = tpe match {
+          case '[t *: ts] => (Type.of[t], select[t](n)) +: asSeq(Type.of[ts], n + 1)
+          case '[EmptyTuple] => Nil
+        }
+        elems0 match {
+          case '[type elems <: Tuple; `elems`] =>
+            val wrappedElems = wrappedT.asExprOf[Seq[elems]]
+            (
+              '{ $wrappedElems.map(_.productIterator.toList) },
+              asSeq(elems0, 0)
+            )
+          case '[t] =>
+            (
+              '{ $wrappedT.map(List(_)) },
+              List((Type.of[t], identity))
+            )
+        }
+      }
+
+      def exPair(n: Int): (Type[?], Expr[?] => Expr[?]) = {
+        elemTypes(1).lift(n).getOrElse(
+          report.errorAndAbort(
+            s"expected at least ${n + 1} elements, got ${elemTypes(1).size}",
+            Position.ofMacroExpansion
+          )
+        )
+      }
+
+      def exType(n: Int): TypeRepr = {
+        val (elemType, _) = exPair(n)
+        elemType match
+          case '[t] => TypeRepr.of[t]
+      }
+
+      def exTerm(n: Int): Expr[?] => Expr[?] = {
+        exPair(n)(1)
+      }
+
+      def mkSegmentsCall[T: Type](t: Expr[T]): Expr[List[String]] = {
+        val summonCall = Expr.summon[ToSegments[T]].getOrElse(
+          report.errorAndAbort(s"Could not summon ToSegments[${Type.show[T]}]", Position.ofMacroExpansion)
+        )
+        '{mill.define.Cross.ToSegments[T]($t)(using $summonCall) }
+      }
+
+      def mkSegmentsCallN(n: Int)(arg: Expr[?]): Expr[List[String]] = {
+        exTerm(n)(arg) match {
+          case '{ $v1: t1 } => mkSegmentsCall[t1](v1)
+        }
+      }
+
+      def newGetter(name: String, res: TypeRepr, flags: Flags = Flags.Override): Symbol => Symbol =
+        cls =>
+          Symbol.newMethod(
+            parent = cls,
+            name = name,
+            tpe = ByNameType(res),
+            flags = flags,
+            privateWithin = Symbol.noSymbol
+          )
+      def newField(name: String, res: TypeRepr, flags: Flags): Symbol => Symbol =
+        cls =>
+          Symbol.newVal(
+            parent = cls,
+            name = name,
+            tpe = res,
+            flags = flags,
+            privateWithin = Symbol.noSymbol
+          )
+
+      def newGetterTree(name: String, rhs: Expr[?] => Expr[?]): (Symbol, Expr[?]) => Statement = {
+        (cls, arg) =>
+          val sym = cls.declaredMethod(name)
+            .headOption
+            .getOrElse(report.errorAndAbort(s"could not find method $name in $cls", Position.ofMacroExpansion))
+          DefDef(sym, _ => Some(rhs(arg).asTerm))
+      }
+
+      def newValTree(name: String, rhs: Option[Term]): (Symbol, Expr[?]) => Statement = {
+        (cls, _) =>
+          val sym = {
+            val sym0 = cls.declaredField(name)
+            if sym0 != Symbol.noSymbol then sym0
+            else report.errorAndAbort(s"could not find field $name in $cls", Position.ofMacroExpansion)
+          }
+          ValDef(sym, rhs)
+      }
+
+      extension (sym: Symbol) {
+        def mkRef(debug: => String): Ref = {
+          if sym.isTerm then
+            Ref(sym)
+          else
+            report.errorAndAbort(s"could not ref ${debug}, it was not a term")
+        }
+      }
+
+      val newSyms = List.newBuilder[Symbol => Symbol]
+      val newTrees = collection.mutable.Buffer.empty[(Symbol, Expr[?]) => Statement]
+      val valuesTree: Expr[Seq[Seq[Any]]] = elemTypes(0)
+      val pathSegmentsTrees = List.newBuilder[Expr[?] => Expr[List[String]]]
+
+      def pushElemTrees(n: Int): Unit = {
+        val name = s"crossValue${if n > 0 then (n + 1).toString else ""}"
+        newSyms += newGetter(name, res = exType(n))
+        newTrees += newGetterTree(name, rhs = exTerm(n))
+        pathSegmentsTrees += mkSegmentsCallN(n)
+      }
+
+      newSyms += newField(
+        "local_ctx",
+        res = TypeRepr.of[mill.define.Ctx],
+        flags = Flags.PrivateLocal | Flags.ParamAccessor)
+
+      newTrees += newValTree("local_ctx", rhs = None)
+
+      if tpe <:< TypeRepr.of[Module[?]] then
+        pushElemTrees(0)
+      else
+        report.errorAndAbort(
+          s"Cross type ${tpe.show} must implement Cross.Module[T]",
+          Position.ofMacroExpansion
+        )
+
+      if tpe <:< TypeRepr.of[Module2[?, ?]] then
+        pushElemTrees(1)
+
+      if tpe <:< TypeRepr.of[Module3[?, ?, ?]] then
+        pushElemTrees(2)
+
+      if (tpe <:< TypeRepr.of[Module4[?, ?, ?, ?]])
+        pushElemTrees(3)
+
+      if (tpe <:< TypeRepr.of[Module5[?, ?, ?, ?, ?]])
+        pushElemTrees(4)
+
+      val pathSegmentsTree: Expr[?] => Expr[List[String]] =
+        pathSegmentsTrees.result().reduceLeft((a, b) => arg => '{ ${a(arg)} ++ ${b(arg)} })
+
+      def newCtor(cls: Symbol): (List[String], List[TypeRepr]) =
+        (List("local_ctx"), List(TypeRepr.of[mill.define.Ctx]))
+
+      def newClassDecls(cls: Symbol): List[Symbol] = {
+        newSyms.result().map(_(cls))
+      }
+
+      def clsFactory()(using Quotes): Symbol = {
+        shims.Symbol.newClass(
+          parent = cls,
+          name = s"${cls.name}_impl",
+          parents = List(TypeRepr.of[mill.define.Module.BaseClass], tpe),
+          ctor = newCtor,
+          decls = newClassDecls,
+          selfType = None
+        )
+      }
+
+      // We need to create a `class $concreteCls` here, rather than just
+      // creating an anonymous sub-type of $tpe, because our task resolution
+      // logic needs to use java reflection to identify sub-modules and java
+      // reflect can only properly identify nested `object`s inside Scala
+      // `object` and `class`es.
+      elems0 match {
+        case '[elems] =>
+          val wrappedElems = wrappedT.asExprOf[Seq[elems]]
+          val ref = '{
+            new mill.define.Cross.Factory[T](
+              makeList = $wrappedElems.map { (v2: elems) =>
+                ${
+                  val concreteCls = clsFactory()
+                  val concreteClsDef = shims.ClassDef(
+                    cls = concreteCls,
+                    parents = {
+                      val parentCtor =
+                        New(TypeTree.of[mill.define.Module.BaseClass]).select(
+                          TypeRepr.of[mill.define.Module.BaseClass].typeSymbol.primaryConstructor
+                        )
+                      val parentApp =
+                        parentCtor.appliedToNone.appliedTo(
+                          concreteCls.declaredField("local_ctx").mkRef(s"${concreteCls} field local_ctx")
+                        )
+                      List(parentApp, TypeTree.of[T])
+                    },
+                    body = newTrees.toList.map(_(concreteCls, 'v2))
+                  )
+                  val clsOf = Ref(defn.Predef_classOf).appliedToType(concreteCls.typeRef)
+                  def newCls(ctx0: Expr[mill.define.Ctx]): Expr[T] = {
+                    New(TypeTree.ref(concreteCls))
+                      .select(concreteCls.primaryConstructor)
+                      .appliedTo(ctx0.asTerm)
+                      .asExprOf[T]
+                  }
+                  Block(
+                    List(concreteClsDef),
+                    '{ (${clsOf.asExprOf[Class[?]]}, (ctx0: mill.define.Ctx) => ${newCls('ctx0)}) }.asTerm
+                  ).asExprOf[(Class[?], mill.define.Ctx => T)]
+                }
+              },
+              crossSegmentsList = $wrappedElems.map((segArg: elems) => ${pathSegmentsTree('segArg)}),
+              crossValuesListLists = $valuesTree,
+              crossValuesRaw = $wrappedT
+            )(using compiletime.summonInline[reflect.ClassTag[T]])
+          }
+          // report.errorAndAbort(s"made factory ${ref.show}")
+          ref
+      }
+    }
+
+    def abortOldStyleClass(using Quotes)(tpe: quotes.reflect.TypeRepr): Nothing = {
+      import quotes.reflect.*
+
       val primaryConstructorArgs =
-        tpe.typeSymbol.asClass.primaryConstructor.typeSignature.paramLists.head
+        tpe.classSymbol.get.primaryConstructor.paramSymss.head
 
       val oldArgStr = primaryConstructorArgs
-        .map { s => s"${s.name}: ${s.typeSignature}" }
+        .map { s => s"${s.name}: ${s.termRef.widen.show}" }
         .mkString(", ")
 
       def parenWrap(s: String) =
         if (primaryConstructorArgs.size == 1) s
         else s"($s)"
 
-      val newTypeStr = primaryConstructorArgs.map(_.typeSignature.toString).mkString(", ")
-      val newForwarderStr = primaryConstructorArgs.map(_.name.toString).mkString(", ")
+      val newTypeStr = primaryConstructorArgs.map(_.termRef.widen.show).mkString(", ")
+      val newForwarderStr = primaryConstructorArgs.map(_.name).mkString(", ")
 
-      c.abort(
-        c.enclosingPosition,
+      report.errorAndAbort(
         s"""
            |Cross type ${tpe.typeSymbol.name} must be trait, not a class. Please change:
            |
@@ -255,7 +419,8 @@ object Cross {
            |you may remove it. If you do not have this definition, you can
            |preserve the old behavior via `def millSourcePath = super.millSourcePath / crossValue`
            |
-           |""".stripMargin
+           |""".stripMargin,
+           Position.ofMacroExpansion
       )
     }
   }