Scala白盒宏：递归移除Case Class字段Option包装的实现问题

问题描述 #

我要实现一个Scala注解宏，功能是从给定的case class生成新的case class，递归移除所有字段的Option类型包装——如果字段本身是case class，也要移除其内部字段的Option。目前已经完成了非嵌套case class的处理，但递归处理时无法获取case class类型字段的ClassTag，不清楚如何在语法树阶段完成类型检查。现有代码如下：

import scala.annotation.StaticAnnotation
import scala.collection.mutable
import scala.reflect.macros.blackbox.Context

import scala.language.experimental.macros
import scala.annotation.compileTimeOnly

class RemoveOptionFromFields extends StaticAnnotation {
  def macroTransform(annottees: Any*): Any = macro RemoveOptionFromFields.impl
}

object RemoveOptionFromFields {
  def impl(c: Context)(annottees: c.Expr[Any]*): c.Expr[Any] = {
    import c.universe._
 
    def modifiedClass(classDecl: ClassDef, compDeclOpt: Option[ModuleDef]) = {
      val result = classDecl match {
        case q"case class $className(..$fields) extends ..$parents { ..$body }" =>
          val fieldsWithoutOption = fields.map {
            case ValDef(mods, name, tpt, rhs) =>
              tpt.children match {
                case List(first, second) if first.toString() == "Option" =>
                  // Check if `second` is a case class?
                  // Get it's fields if so
                  val innerType = tpt.children(1)
                  ValDef(mods, name, innerType, rhs)
                case _ => 
                  ValDef(mods, name, tpt, rhs)
              }
          }

          val withOptionRemovedFromFieldsClassDecl = q"case class WithOptionRemovedFromFields(..$fieldsWithoutOption)"

          val newCompanionDecl = compDeclOpt.fold(
            q"""
            object ${className.toTermName} {
              $withOptionRemovedFromFieldsClassDecl
            }
            """
          ) {
            compDecl =>
              val q"object $obj extends ..$bases { ..$body }" = compDecl
              q"""
              object $obj extends ..$bases {
                ..$body

                $withOptionRemovedFromFieldsClassDecl
              }
              """
          }

          q"""
            $classDecl
            $newCompanionDecl
          """
      }
      c.Expr[Any](result)
    }
 
    annottees.map(_.tree) match {
      case (classDecl: ClassDef) :: Nil => modifiedClass(classDecl, None)
      case (classDecl: ClassDef) :: (compDecl: ModuleDef) :: Nil => modifiedClass(classDecl, Some(compDecl))
      case _ => c.abort(c.enclosingPosition, "This annotation only supports classes")
    }
  }
}

解决方案 #

要实现递归处理嵌套case class，核心是在宏的语法树阶段通过类型符号解析判断目标类型是否为case class，并递归生成对应的无Option版本。以下是修改后的完整代码：

import scala.annotation.StaticAnnotation
import scala.reflect.macros.blackbox.Context

import scala.language.experimental.macros
import scala.annotation.compileTimeOnly

class RemoveOptionFromFields extends StaticAnnotation {
  def macroTransform(annottees: Any*): Any = macro RemoveOptionFromFields.impl
}

object RemoveOptionFromFields {
  def impl(c: Context)(annottees: c.Expr[Any]*): c.Expr[Any] = {
    import c.universe._

    // 存储已生成的case class定义，避免重复生成
    val generatedClasses = mutable.Set[Type]()

    // 判断类型是否为case class
    def isCaseClass(tpe: Type): Boolean = {
      val sym = tpe.typeSymbol
      sym.isClass && sym.asClass.isCaseClass
    }

    // 递归移除类型的Option包装，若为case class则生成对应的无Option版本
    def stripOption(tpt: Tree): Tree = tpt match {
      case AppliedTypeTree(Ident(TypeName("Option")), List(innerTpt)) =>
        val strippedInner = stripOption(innerTpt)
        strippedInner
      case _ =>
        val tpe = c.typecheck(tpt, mode = c.TYPEmode).tpe
        if (isCaseClass(tpe)) {
          generateCaseClassWithoutOption(tpe)
          Ident(TypeName(s"${tpe.typeSymbol.name}WithoutOption"))
        } else {
          tpt
        }
    }

    // 针对给定的case class类型，生成其无Option字段的版本
    def generateCaseClassWithoutOption(tpe: Type): Tree = {
      if (generatedClasses.contains(tpe)) return EmptyTree
      generatedClasses.add(tpe)

      val className = tpe.typeSymbol.name.toTypeName
      val fields = tpe.decls.collect {
        case m: MethodSymbol if m.isCaseAccessor =>
          val name = m.name.toTermName
          val originalTpt = m.returnType
          val strippedTpt = stripOption(TypeTree(originalTpt))
          ValDef(Modifiers(Flag.PARAM), name, strippedTpt, EmptyTree)
      }.toList

      q"case class ${TypeName(s"${className}WithoutOption")}(..$fields)"
    }

    def modifiedClass(classDecl: ClassDef, compDeclOpt: Option[ModuleDef]) = {
      val result = classDecl match {
        case q"case class $className(..$fields) extends ..$parents { ..$body }" =>
          // 处理当前类的所有字段，递归移除Option
          val fieldsWithoutOption = fields.map {
            case ValDef(mods, name, tpt, rhs) =>
              val strippedTpt = stripOption(tpt)
              ValDef(mods, name, strippedTpt, rhs)
          }

          // 生成当前类的无Option版本
          val currentClassWithoutOption = q"case class ${TypeName(s"${className}WithoutOption")}(..$fieldsWithoutOption)"

          // 收集所有递归生成的嵌套case class定义
          val nestedClasses = generatedClasses.filter(_ != c.typecheck(q"$className", mode = c.TYPEmode).tpe)
            .map(generateCaseClassWithoutOption)

          val newCompanionDecl = compDeclOpt.fold(
            q"""
            object ${className.toTermName} {
              $currentClassWithoutOption
              ..$nestedClasses
            }
            """
          ) {
            compDecl =>
              val q"object $obj extends ..$bases { ..$body }" = compDecl
              q"""
              object $obj extends ..$bases {
                ..$body
                $currentClassWithoutOption
                ..$nestedClasses
              }
              """
          }

          q"""
            $classDecl
            $newCompanionDecl
          """
      }
      c.Expr[Any](result)
    }

    annottees.map(_.tree) match {
      case (classDecl: ClassDef) :: Nil => modifiedClass(classDecl, None)
      case (classDecl: ClassDef) :: (compDecl: ModuleDef) :: Nil => modifiedClass(classDecl, Some(compDecl))
      case _ => c.abort(c.enclosingPosition, "This annotation only supports classes")
    }
  }
}

关键逻辑说明 #

• isCaseClass：通过类型符号判断是否为case class，比字符串匹配更可靠
• stripOption：递归处理类型树，遇到Option[T]就提取T并继续递归；遇到case class则触发生成其无Option版本
• generateCaseClassWithoutOption：解析case class的所有访问器方法（即字段），递归处理每个字段的类型，生成新的case class定义，并用generatedClasses避免重复生成
• 嵌套类收集：将递归生成的所有嵌套case class定义注入到原伴生对象中，确保可以直接使用

内容的提问来源于stack exchange，提问作者Aiono