fix return

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

@@ -42,7 +42,7 @@ object elapsedMacro {
     import c.universe._
 
     private lazy val start: c.universe.TermName = TermName("$elapsedBegin")
-    private lazy val valDef = TermName("$elapsed")
+    private lazy val valDef: c.universe.TermName = TermName("$elapsed")
 
     private def getLogLevel(logLevel: Tree): LogLevel = {
       if (logLevel.children.exists(t => t.toString().contains(PACKAGE))) {
@@ -108,7 +108,19 @@ object elapsedMacro {
     }
 
     // There may be a half-way exit, rather than the one whose last expression is exit.
-    private def returnEarly(defDef: DefDef, trees: List[Tree]): List[Tree] = {
+    // Unreliable function!!!
+    private def returnEarly(defDef: DefDef, trees: Tree*): List[Tree] = {
+      val ifElseMatch = (f: If) => {
+        if (f.elsep.nonEmpty) {
+          if (f.elsep.children.nonEmpty && f.elsep.children.size > 1) {
+            If(f.cond, f.thenp, q"..${returnEarly(defDef, f.elsep.children: _*)}")
+          } else {
+            If(f.cond, f.thenp, q"..${returnEarly(defDef, f.elsep)}")
+          }
+        } else {
+          f
+        }
+      }
       if (trees.isEmpty) return Nil
       trees.map {
         case r: Return =>
@@ -117,16 +129,24 @@ object elapsedMacro {
              $r
             """
         case f: If => //support if return
-          if (!f.thenp.isEmpty) {
-            If(f.cond, q"..${returnEarly(defDef, List(f.thenp))}", f.elsep)
-          } else if (!f.elsep.isEmpty) {
-            If(f.cond, f.thenp, q"..${returnEarly(defDef, List(f.elsep))}")
+          c.info(c.enclosingPosition, s"returnEarly: thenp: ${f.thenp}, children: ${f.thenp.children}, cond: ${f.cond}", force = true)
+          c.info(c.enclosingPosition, s"returnEarly: elsep: ${f.elsep}, children: ${f.elsep.children}, cond: ${f.cond}", force = true)
+          if (f.thenp.nonEmpty) {
+            if (f.thenp.children.nonEmpty && f.thenp.children.size > 1) {
+              val ifTree = If(f.cond, q"..${returnEarly(defDef, f.thenp.children: _*)}", f.elsep)
+              ifElseMatch(ifTree)
+            } else {
+              val ifTree = If(f.cond, q"..${returnEarly(defDef, f.thenp)}", f.elsep)
+              ifElseMatch(ifTree)
+            }
           } else {
-            f
+            ifElseMatch(f)
           }
-        case t => t
-        // TODO support for/while/switch
-      }
+        case t =>
+          // TODO support for/while/switch
+          c.info(c.enclosingPosition, s"returnEarly: not support expr: $t", force = true)
+          t
+      }.toList
     }
 
     private def getNewMethod(defDef: DefDef): DefDef = {
@@ -135,7 +155,7 @@ object elapsedMacro {
         val last = defDef.rhs.children.last
         q"""
           $getStartExpr
-          ..${returnEarly(defDef, heads)}
+          ..${returnEarly(defDef, heads: _*)}
           ..${getPrintlnLog(defDef.name)}
           $last
         """

src/test/scala/io/github/dreamylost/ElapsedTest.scala

@@ -34,7 +34,7 @@ import scala.concurrent.Future
  */
 class ElapsedTest extends AnyFlatSpec with Matchers {
 
-  "elapsed1" should "not calculate anything, the return type is not specified" in {
+  "elapsed1" should "failed, not calculate anything, the return type is not specified" in {
     //Duration and TimeUnit must Full class name
     """
       |    class A {
@@ -53,10 +53,92 @@ class ElapsedTest extends AnyFlatSpec with Matchers {
       |    }
       |
       |    class D {
-      |      @elapsed(limit = scala.concurrent.duration.Duration(1, java.util.concurrent.TimeUnit.SECONDS), logLevel = io.github.dreamylost.LogLevel.ERROR)
-      |      def j = ???
-      |    }|
+      |      @elapsed(limit = scala.concurrent.duration.Duration(1, java.util.concurrent.TimeUnit.SECONDS), logLevel = io.github.dreamylost.LogLevel.INFO)
+      |      def i:String = ???
+      |   }
       |""".stripMargin shouldNot compile
+
+    class A {
+
+      @elapsed(limit = scala.concurrent.duration.Duration(1, java.util.concurrent.TimeUnit.SECONDS), logLevel = io.github.dreamylost.LogLevel.INFO)
+      def j: Int = {
+        var i = 1
+        if (i == 1) {
+          val h = 0
+          var l = 0
+          if (j == 0) {
+            val h = 0
+            var l = 0
+            return 1
+          } else {
+            val j = 0
+            return 0
+          }
+        } else {
+          i
+        }
+        i
+      }
+
+      @elapsed(limit = scala.concurrent.duration.Duration(1, java.util.concurrent.TimeUnit.SECONDS), logLevel = io.github.dreamylost.LogLevel.INFO)
+      def k: Unit = {
+        var i = 1
+        if (i == 1) {
+          val i = 0
+          val k = 0
+          if (i == 0) {
+            1
+          } else { 2 }
+        } else {
+          val u = 0
+          0
+        }
+
+        var k = 1
+        if (k == 1) {
+          val i = 0
+          val k = 0
+          if (i == 0) {
+            return ()
+          } else { return () }
+        } else {
+          val u = 0
+          return u
+        }
+        1
+      }
+
+      @elapsed(limit = scala.concurrent.duration.Duration(1, java.util.concurrent.TimeUnit.SECONDS), logLevel = io.github.dreamylost.LogLevel.INFO)
+      def l: Int = {
+        val i = 0
+        for (i <- Seq(1)) {
+          if (i == 1) {
+            return 1 //not support
+          }
+        }
+        0
+      }
+
+      @elapsed(limit = scala.concurrent.duration.Duration(1, java.util.concurrent.TimeUnit.SECONDS), logLevel = io.github.dreamylost.LogLevel.INFO)
+      def m: Int = {
+        var i = 1
+        if (i == 1) {
+        } else {
+          val u = 0
+          return 0
+        }
+
+        if (i == 1) {
+          return 1
+        } else {
+          val u = 0
+          return 0
+        }
+
+        1
+      }
+
+    }
   }
 
   "elapsed2" should "ok, get the returnType of the method " in {
@@ -124,8 +206,19 @@ class ElapsedTest extends AnyFlatSpec with Matchers {
       |
       |    val a = new A
       |    a.helloScala1
+      |
+      |     class B {
+      |      @elapsed(limit = scala.concurrent.duration.Duration(1, java.util.concurrent.TimeUnit.SECONDS), logLevel = io.github.dreamylost.LogLevel.INFO)
+      |      def helloScala11: String = {
+      |        val s = "hello"
+      |        if (s == "hello") {
+      |          return "world" + "wooo"
+      |        }
+      |        val x = "world"
+      |        return s
+      |      }
+      |    }
       |""".stripMargin should compile
-
   }
 
   "elapsed5" should "ok, return Future" in {