From 8b04aadc3c5df4a86d68be1ce02cbe1f5c6f3ced Mon Sep 17 00:00:00 2001
From: Matt Dziuban <mrdziuban@gmail.com>
Date: Fri, 13 Sep 2024 09:34:46 -0400
Subject: [PATCH] Add `NoUnnecessaryForComprehension` rule.

---
 .../fix/NoUnnecessaryForComprehension.scala   | 43 +++++++++++++
 .../fix/NoUnnecessaryForComprehension.scala   | 12 ++++
 .../META-INF/services/scalafix.v1.Rule        |  1 +
 .../fix/NoUnnecessaryForComprehension.scala   | 62 +++++++++++++++++++
 4 files changed, 118 insertions(+)
 create mode 100644 scalafix/input/src/main/scala-3/fix/NoUnnecessaryForComprehension.scala
 create mode 100644 scalafix/output/src/main/scala-3/fix/NoUnnecessaryForComprehension.scala
 create mode 100644 scalafix/rules/src/main/scala/fix/NoUnnecessaryForComprehension.scala

diff --git a/scalafix/input/src/main/scala-3/fix/NoUnnecessaryForComprehension.scala b/scalafix/input/src/main/scala-3/fix/NoUnnecessaryForComprehension.scala
new file mode 100644
index 0000000..f3fc703
--- /dev/null
+++ b/scalafix/input/src/main/scala-3/fix/NoUnnecessaryForComprehension.scala
@@ -0,0 +1,43 @@
+/*
+rule = NoUnnecessaryForComprehension
+NoUnnecessaryForComprehension.color = false
+*/
+package fix
+
+object NoUnnecessaryForComprehension {
+  val x = for {/* assert: NoUnnecessaryForComprehension
+          ^
+
+A for comprehension with only one statement can be simplified
+
+Cases where the yield returns the result of the statement can just be the statement itself:
+
+Before:
+
+  for {
+    x <- someStatementHere
+  } yield x
+
+After:
+
+  someStatementHere
+
+Cases where the yield performs an additional computation can be rewritten with map:
+
+Before:
+
+  for {
+    x <- someStatementHere
+  } yield doSomethingElse(x)
+
+After:
+
+  someStatementHere.map(doSomethingElse) */
+    i <- Option(1)
+  } yield i
+
+  val y = for {
+    i <- Option(1)
+    j <- Option(2)
+  } yield i + j
+}
diff --git a/scalafix/output/src/main/scala-3/fix/NoUnnecessaryForComprehension.scala b/scalafix/output/src/main/scala-3/fix/NoUnnecessaryForComprehension.scala
new file mode 100644
index 0000000..2502a9b
--- /dev/null
+++ b/scalafix/output/src/main/scala-3/fix/NoUnnecessaryForComprehension.scala
@@ -0,0 +1,12 @@
+package fix
+
+object NoUnnecessaryForComprehension {
+  val x = for {
+    i <- Option(1)
+  } yield i
+
+  val y = for {
+    i <- Option(1)
+    j <- Option(2)
+  } yield i + j
+}
diff --git a/scalafix/rules/src/main/resources/META-INF/services/scalafix.v1.Rule b/scalafix/rules/src/main/resources/META-INF/services/scalafix.v1.Rule
index 83bf89b..ade4a93 100644
--- a/scalafix/rules/src/main/resources/META-INF/services/scalafix.v1.Rule
+++ b/scalafix/rules/src/main/resources/META-INF/services/scalafix.v1.Rule
@@ -1,3 +1,4 @@
 fix.NoUnnecessaryCase
+fix.NoUnnecessaryForComprehension
 fix.StrictSubclassAccess
 fix.NoWithForExtends
diff --git a/scalafix/rules/src/main/scala/fix/NoUnnecessaryForComprehension.scala b/scalafix/rules/src/main/scala/fix/NoUnnecessaryForComprehension.scala
new file mode 100644
index 0000000..de00e6f
--- /dev/null
+++ b/scalafix/rules/src/main/scala/fix/NoUnnecessaryForComprehension.scala
@@ -0,0 +1,62 @@
+package fix
+
+import metaconfig.{ConfDecoder, Configured}
+import metaconfig.generic.{deriveDecoder, deriveSurface, Surface}
+import scalafix.v1._
+import scala.io.AnsiColor
+import scala.meta._
+
+case class NoUnnecessaryForComprehensionConfig(color: Boolean)
+object NoUnnecessaryForComprehensionConfig {
+  val default = NoUnnecessaryForComprehensionConfig(true)
+
+  implicit val surface: Surface[NoUnnecessaryForComprehensionConfig] = deriveSurface
+  implicit val decoder: ConfDecoder[NoUnnecessaryForComprehensionConfig] = deriveDecoder(default)
+}
+
+case class UnnecessaryForComprehensionLint(position: Position, color: Boolean) extends Diagnostic {
+  private def withColor(c: String): String => String = s => if (color) c ++ s ++ AnsiColor.RESET else s
+  private val magenta = withColor(AnsiColor.MAGENTA)
+  private val red = withColor(AnsiColor.RED)
+
+  override def message = s"""|
+    |A ${magenta("for")} comprehension with only one statement can be simplified
+    |
+    |Cases where the ${magenta("yield")} returns the result of the statement can just be the statement itself:
+    |
+    |Before:
+    |
+    |  ${magenta("for")} {
+    |    ${red("x")} <- someStatementHere
+    |  } ${magenta("yield")} x
+    |
+    |After:
+    |
+    |  someStatementHere
+    |
+    |Cases where the ${magenta("yield")} performs an additional computation can be rewritten with ${magenta("map")}:
+    |
+    |Before:
+    |
+    |  ${magenta("for")} {
+    |    ${red("x")} <- someStatementHere
+    |  } ${magenta("yield")} doSomethingElse(x)
+    |
+    |After:
+    |
+    |  someStatementHere.map(doSomethingElse)
+    |""".stripMargin
+}
+
+class NoUnnecessaryForComprehension(config: NoUnnecessaryForComprehensionConfig)
+extends SyntacticRule("NoUnnecessaryForComprehension") {
+  def this() = this(NoUnnecessaryForComprehensionConfig.default)
+
+  override def withConfiguration(config: Configuration): Configured[Rule] =
+    config.conf.getOrElse("NoUnnecessaryForComprehension")(this.config).map(new NoUnnecessaryForComprehension(_))
+
+  override def fix(implicit doc: SyntacticDocument): Patch =
+    doc.tree.collect {
+      case t @ Term.ForYield(List(_), _) => Patch.lint(UnnecessaryForComprehensionLint(t.pos, config.color))
+    }.asPatch
+}