remove Eff.pure in for. not only pureEff

xuwei-k · Dec 12, 2023 · 412ac64 · 412ac64
1 parent 701edea
commit 412ac64
Show file tree

Hide file tree

Showing 3 changed files with 34 additions and 4 deletions.
diff --git a/input/src/main/scala/fix/RemovePureEffTest.scala b/input/src/main/scala/fix/RemovePureEffTest.scala
@@ -8,11 +8,13 @@ rule = RemovePureEff
  */
 class RemovePureEffTest {
 
-  def x[R]: Eff[R, (Int, Int, Eff[R, Int], Int)] = for {
+  def x[R]: Eff[R, (Int, Int, Eff[R, Int], Int, Int, String)] = for {
     a1 <- 1.pureEff[R]
     a2 <- 2.pureEff[R]
     a3 = 3.pureEff[R]
     a4 <- 4.pureEff[R]
-  } yield (a1, a2, a3, a4)
+    a5 <- Eff.pure(5)
+    a6 <- Eff.pure[R, String]("aaa")
+  } yield (a1, a2, a3, a4, a5, a6)
 
 }
diff --git a/output/src/main/scala/fix/RemovePureEffTest.scala b/output/src/main/scala/fix/RemovePureEffTest.scala
@@ -6,11 +6,13 @@ import org.atnos.eff.syntax.eff.*
 
 class RemovePureEffTest {
 
-  def x[R]: Eff[R, (Int, Int, Eff[R, Int], Int)] = for {
+  def x[R]: Eff[R, (Int, Int, Eff[R, Int], Int, Int, String)] = for {
     a1 <- 1.pureEff[R]
     a2 = 2
     a3 = 3.pureEff[R]
     a4 = 4
-  } yield (a1, a2, a3, a4)
+    a5 = 5
+    a6 = "aaa"
+  } yield (a1, a2, a3, a4, a5, a6)
 
 }
diff --git a/rules/src/main/scala/fix/RemovePureEff.scala b/rules/src/main/scala/fix/RemovePureEff.scala
@@ -19,6 +19,32 @@ class RemovePureEff extends SyntacticRule("RemovePureEff") {
                 ApplyType.Initial(Term.Select(rhs, Term.Name("pureEff")), Type.Name(_) :: Nil)
               ) =>
             Patch.replaceTree(x, Enumerator.Val(x.pat, rhs).toString)
+          case x @ Enumerator.Generator(
+                _,
+                Term.Apply.After_4_6_0(
+                  Term.ApplyType.After_4_6_0(
+                    Term.Select(
+                      Term.Name("Eff"),
+                      Term.Name("pure")
+                    ),
+                    Type.ArgClause(_ :: _ :: Nil)
+                  ),
+                  Term.ArgClause(arg :: Nil, None)
+                )
+              ) =>
+            Patch.replaceTree(x, Enumerator.Val(x.pat, arg).toString)
+
+          case x @ Enumerator.Generator(
+                _,
+                Term.Apply.After_4_6_0(
+                  Term.Select(
+                    Term.Name("Eff"),
+                    Term.Name("pure")
+                  ),
+                  Term.ArgClause(arg :: Nil, None)
+                )
+              ) =>
+            Patch.replaceTree(x, Enumerator.Val(x.pat, arg).toString)
         }
         .asPatch
     }.asPatch