refactor

src/main/scala/io/github/dreamylost/macros/constructorMacro.scala

@@ -102,7 +102,11 @@ object constructorMacro {
       val resTree = appendClassBody(
         classDecl,
         classInfo => Seq(getThisMethodWithCurrying(classInfo.classParamss, classInfo.body)))
-      c.Expr(resTree)
+      c.Expr(
+        q"""
+          ${compDeclOpt.fold(EmptyTree)(x => x)}
+          $resTree
+         """)
     }
 
     override def impl(annottees: c.universe.Expr[Any]*): c.universe.Expr[Any] = {
@@ -110,9 +114,9 @@ object constructorMacro {
       if (isCaseClass(annotateeClass)) {
         c.abort(c.enclosingPosition, ErrorMessage.ONLY_CLASS)
       }
-      val res = returnWithCompanionObject(collectCustomExpr(annottees)(createCustomExpr).tree, annottees)
-      printTree(force = extractArgumentsDetail._1, res)
-      c.Expr(res)
+      val res = collectCustomExpr(annottees)(createCustomExpr)
+      printTree(force = extractArgumentsDetail._1, res.tree)
+      res
     }
   }
 

src/main/scala/io/github/dreamylost/macros/equalsAndHashCodeMacro.scala

@@ -49,9 +49,8 @@ object equalsAndHashCodeMacro {
         c.abort(c.enclosingPosition, ErrorMessage.ONLY_CLASS)
       }
       val resTree = collectCustomExpr(annottees)(createCustomExpr)
-      val res = returnWithCompanionObject(resTree.tree, annottees)
-      printTree(force = extractArgumentsDetail._1, res)
-      c.Expr(res)
+      printTree(force = extractArgumentsDetail._1, resTree.tree)
+      resTree
     }
 
     /**
@@ -120,11 +119,15 @@ object equalsAndHashCodeMacro {
       }
       val classDefinition = mapToClassDeclInfo(classDecl)
       val res = appendClassBody(classDecl, classInfo =>
-        getEqualsMethod(classDefinition.className, map(classInfo),
-          classDefinition.superClasses, classDefinition.body) ++
+        getEqualsMethod(classDefinition.className, map(classInfo), classDefinition.superClasses, classDefinition.body) ++
           List(getHashcodeMethod(map(classInfo), classDefinition.superClasses))
       )
-      c.Expr(res)
+
+      c.Expr(
+        q"""
+          ${compDeclOpt.fold(EmptyTree)(x => x)}
+          $res
+         """)
     }
   }
 

src/main/scala/io/github/dreamylost/macros/jacksonEnumMacro.scala

@@ -81,10 +81,9 @@ object jacksonEnumMacro {
     }
 
     override def impl(annottees: c.universe.Expr[Any]*): c.universe.Expr[Any] = {
-      val newClassExpr = collectCustomExpr(annottees)(createCustomExpr)
-      val res = returnWithCompanionObject(newClassExpr.tree, annottees)
-      printTree(force = extractArgumentsDetail._1, res)
-      c.Expr(res)
+      val res = collectCustomExpr(annottees)(createCustomExpr)
+      printTree(force = extractArgumentsDetail._1, res.tree)
+      res
     }
 
     override def createCustomExpr(classDecl: c.universe.ClassDef, compDeclOpt: Option[c.universe.ModuleDef]): Any = {
@@ -95,11 +94,17 @@ object jacksonEnumMacro {
       val q"$mods class $tpname[..$tparams] $ctorMods(...$paramss) extends ..$bases { ..$body }" = classDecl
       val newFieldss = paramss.asInstanceOf[List[List[Tree]]].map(_.map(replaceAnnotation))
       val newClass = q"$mods class $tpname[..$tparams] $ctorMods(...$newFieldss) extends ..$bases { ..$body }"
-      c.Expr(
+      val res =
         q"""
            ..$typeReferClasses
              
            $newClass 
+         """
+
+      c.Expr(
+        q"""
+          ${compDeclOpt.fold(EmptyTree)(x => x)}
+          $res
          """)
     }
   }

src/main/scala/io/github/dreamylost/macros/toStringMacro.scala

@@ -74,9 +74,14 @@ object toStringMacro {
       // Check the type of the class, which can only be defined on the ordinary class
       val annotateeClass: ClassDef = checkAndGetClassDef(annottees)
       val resTree = appendClassBody(annotateeClass, _ => Seq(toStringTemplateImpl(argument, annotateeClass)))
-      val res = returnWithCompanionObject(resTree, annottees)
-      printTree(argument.verbose, res)
-      c.Expr[Any](res)
+      val compDeclOpt = getCompanionObject(annottees)
+      val res = c.Expr(
+        q"""
+          ${compDeclOpt.fold(EmptyTree)(x => x)}
+          $resTree
+         """)
+      printTree(force = extractArgumentsDetail._1, res.tree)
+      res
     }
 
     private def printField(argument: Argument, lastParam: Option[String], field: Tree): Tree = {