add UnnecessarySortRewriteConfig

fix #19

Author: xuwei-k

build.sbt

@@ -101,6 +101,7 @@ lazy val rules = projectMatrix
   .jvmPlatform(rulesCrossVersions)
 
 lazy val inputOutputCommon = Def.settings(
+  libraryDependencies += "org.scala-lang.modules" %% "scala-collection-compat" % "2.6.0",
   libraryDependencies ++= {
     if (scalaBinaryVersion.value == "2.13") {
       Seq("io.circe" %% "circe-generic-extras" % "0.14.1")

input/src/main/scala/fix/UnnecessarySortRewriteTestAddImport1.scala

@@ -0,0 +1,16 @@
+package fix
+
+/*
+rule = UnnecessarySortRewrite
+UnnecessarySortRewrite.rewriteConfig = addCompatImport
+ */
+abstract class UnnecessarySortRewriteTestAddImport1 {
+  def seq: Seq[(Int, Int)]
+
+  def foo: Unit = {
+    seq.sortBy(_._1).head
+    seq.sortBy(_._2).last
+    seq.sortBy(_._1).headOption
+    seq.sortBy(_._2).lastOption
+  }
+}

input/src/main/scala/fix/UnnecessarySortRewriteTestAddImport2.scala

@@ -0,0 +1,14 @@
+package fix
+
+/*
+rule = UnnecessarySortRewrite
+UnnecessarySortRewrite.rewriteConfig = addCompatImport
+ */
+abstract class UnnecessarySortRewriteTestAddImport2 {
+  def seq: Seq[(Int, Int)]
+
+  def foo: Unit = {
+    seq.sortBy(_._1).head
+    seq.sortBy(_._2).last
+  }
+}

input/src/main/scala/fix/UnnecessarySortRewriteTestOnly212.scala

@@ -0,0 +1,16 @@
+package fix
+
+/*
+rule = UnnecessarySortRewrite
+UnnecessarySortRewrite.rewriteConfig = only212methods
+ */
+abstract class UnnecessarySortRewriteTestOnly212 {
+  def seq: Seq[(Int, Int)]
+
+  def foo: Unit = {
+    seq.sortBy(_._1).head
+    seq.sortBy(_._2).last
+    seq.sortBy(_._1).headOption
+    seq.sortBy(_._2).lastOption
+  }
+}

output/src/main/scala/fix/UnnecessarySortRewriteTestAddImport1.scala

@@ -0,0 +1,15 @@
+package fix
+
+
+
+import scala.collection.compat._
+abstract class UnnecessarySortRewriteTestAddImport1 {
+  def seq: Seq[(Int, Int)]
+
+  def foo: Unit = {
+    seq.minBy(_._1)
+    seq.maxBy(_._2)
+    seq.minByOption(_._1)
+    seq.maxByOption(_._2)
+  }
+}

output/src/main/scala/fix/UnnecessarySortRewriteTestAddImport2.scala

@@ -0,0 +1,11 @@
+package fix
+
+
+abstract class UnnecessarySortRewriteTestAddImport2 {
+  def seq: Seq[(Int, Int)]
+
+  def foo: Unit = {
+    seq.minBy(_._1)
+    seq.maxBy(_._2)
+  }
+}

output/src/main/scala/fix/UnnecessarySortRewriteTestOnly212.scala

@@ -0,0 +1,13 @@
+package fix
+
+
+abstract class UnnecessarySortRewriteTestOnly212 {
+  def seq: Seq[(Int, Int)]
+
+  def foo: Unit = {
+    seq.minBy(_._1)
+    seq.maxBy(_._2)
+    seq.sortBy(_._1).headOption
+    seq.sortBy(_._2).lastOption
+  }
+}

rules/src/main/scala/fix/UnnecessarySort.scala

@@ -9,12 +9,14 @@ import scala.meta.Term
 import scala.meta.Position
 
 object UnnecessarySort {
-  val map: Map[String, String] = Map(
-    "head" -> "minBy",
+  val scala213Methods: Map[String, String] = Map(
     "headOption" -> "minByOption",
-    "last" -> "maxBy",
     "lastOption" -> "maxByOption"
   )
+  val map: Map[String, String] = Map(
+    "head" -> "minBy",
+    "last" -> "maxBy",
+  ) ++ scala213Methods
 }
 
 class UnnecessarySort extends SyntacticRule("UnnecessarySort") {

rules/src/main/scala/fix/UnnecessarySortRewrite.scala

@@ -1,24 +1,95 @@
 package fix
 
+import metaconfig.ConfDecoder
+import metaconfig.ConfError
+import metaconfig.Configured
+import metaconfig.generic.Surface
 import scalafix.Patch
+import scalafix.v1.Configuration
+import scalafix.v1.Rule
 import scalafix.v1.SyntacticDocument
 import scalafix.v1.SyntacticRule
+import scala.meta.Pkg
+import scala.meta.Source
 import scala.meta.Term
 
-class UnnecessarySortRewrite extends SyntacticRule("UnnecessarySortRewrite") {
+case class UnnecessarySortRewriteConfig(rewriteConfig: UnnecessarySortRewriteConfig.RewriteConfig)
+
+object UnnecessarySortRewriteConfig {
+  sealed abstract class RewriteConfig(val value: String) extends Product with Serializable
+  object RewriteConfig {
+    val map: Map[String, RewriteConfig] = Seq(Only212Methods, AddCompatImport, Default).map(a => a.value -> a).toMap
+
+    implicit val decoder: ConfDecoder[RewriteConfig] =
+      implicitly[ConfDecoder[String]].flatMap { str =>
+        Configured.opt(map.get(str))(ConfError.message(s"invalid type ${str}"))
+      }
+  }
+  case object Only212Methods extends RewriteConfig("only212methods")
+  case object AddCompatImport extends RewriteConfig("addCompatImport")
+  case object Default extends RewriteConfig("default")
+
+  val default: UnnecessarySortRewriteConfig = UnnecessarySortRewriteConfig(rewriteConfig = Default)
+
+  implicit val surface: Surface[UnnecessarySortRewriteConfig] =
+    metaconfig.generic.deriveSurface[UnnecessarySortRewriteConfig]
+
+  implicit val decoder: ConfDecoder[UnnecessarySortRewriteConfig] =
+    metaconfig.generic.deriveDecoder(default)
+}
+
+class UnnecessarySortRewrite(config: UnnecessarySortRewriteConfig) extends SyntacticRule("UnnecessarySortRewrite") {
+
+  def this() = this(UnnecessarySortRewriteConfig.default)
+
+  override def withConfiguration(config: Configuration): Configured[Rule] = {
+    config.conf.getOrElse("UnnecessarySortRewrite")(this.config).map(newConfig => new UnnecessarySortRewrite(newConfig))
+  }
+
   override def fix(implicit doc: SyntacticDocument): Patch = {
-    doc.tree.collect {
-      case t @ Term.Select(
-            Term.Apply(Term.Select(x1, Term.Name("sortBy")), List(x2)),
-            Term.Name(methodName)
-          ) if UnnecessarySort.map.contains(methodName) =>
-        // TODO add `import scala.collection.compat._`
-        // if Scala 2.12 and minByOption, minByOption
-
-        Patch.replaceTree(
-          t,
-          Term.Apply(Term.Select(x1, Term.Name(UnnecessarySort.map(methodName))), List(x2)).toString,
-        )
+    doc.tree.collect { case src: Source =>
+      val result = src.collect {
+        case t @ Term.Select(
+              Term.Apply(Term.Select(x1, Term.Name("sortBy")), List(x2)),
+              Term.Name(methodName)
+            ) if UnnecessarySort.map.contains(methodName) =>
+          val patch1 = Patch.replaceTree(
+            t,
+            Term.Apply(Term.Select(x1, Term.Name(UnnecessarySort.map(methodName))), List(x2)).toString,
+          )
+
+          config.rewriteConfig match {
+            case UnnecessarySortRewriteConfig.Only212Methods =>
+              if (!UnnecessarySort.scala213Methods.contains(methodName)) {
+                Option((patch1, false))
+              } else {
+                None
+              }
+            case UnnecessarySortRewriteConfig.AddCompatImport =>
+              if (UnnecessarySort.scala213Methods.contains(methodName)) {
+                Option((patch1, true))
+              } else {
+                Option((patch1, false))
+              }
+            case UnnecessarySortRewriteConfig.Default =>
+              Option((patch1, false))
+          }
+      }.flatten
+
+      if (result.nonEmpty) {
+        val patch = result.map(_._1).asPatch
+        val pkg = src.collect { case p: Pkg => p.stats.head }.head
+        if (result.exists(_._2)) {
+          Seq(
+            Patch.addLeft(pkg, "\nimport scala.collection.compat._\n"),
+            patch,
+          ).asPatch
+        } else {
+          patch
+        }
+      } else {
+        Patch.empty
+      }
     }.asPatch
   }
 }