From d97d1665f025036ef55760b9f91986727b7c6b72 Mon Sep 17 00:00:00 2001
From: xuwei-k <6b656e6a69@gmail.com>
Date: Sat, 7 Dec 2024 09:55:32 +0900
Subject: [PATCH] fix CollectHeadOption

---
 input/src/main/scala/fix/CollectHeadOptionTest.scala  | 1 +
 output/src/main/scala/fix/CollectHeadOptionTest.scala | 1 +
 rules/src/main/scala/fix/CollectHeadOption.scala      | 7 ++++++-
 3 files changed, 8 insertions(+), 1 deletion(-)

diff --git a/input/src/main/scala/fix/CollectHeadOptionTest.scala b/input/src/main/scala/fix/CollectHeadOptionTest.scala
index fc3ba16..1ac2f91 100644
--- a/input/src/main/scala/fix/CollectHeadOptionTest.scala
+++ b/input/src/main/scala/fix/CollectHeadOptionTest.scala
@@ -5,4 +5,5 @@ package fix
 
 object CollectHeadOptionTest {
   def x1: Option[String] = List(1, 2, 3).collect { case n if n % 2 == 0 => n.toString }.headOption
+  def x2(f: PartialFunction[Int, String]): Option[String] = List(1, 2, 3).collect(f).headOption
 }
diff --git a/output/src/main/scala/fix/CollectHeadOptionTest.scala b/output/src/main/scala/fix/CollectHeadOptionTest.scala
index 52b6f62..e816244 100644
--- a/output/src/main/scala/fix/CollectHeadOptionTest.scala
+++ b/output/src/main/scala/fix/CollectHeadOptionTest.scala
@@ -2,4 +2,5 @@ package fix
 
 object CollectHeadOptionTest {
   def x1: Option[String] = List(1, 2, 3).collectFirst{ case n if n % 2 == 0 => n.toString }
+  def x2(f: PartialFunction[Int, String]): Option[String] = List(1, 2, 3).collectFirst(f)
 }
diff --git a/rules/src/main/scala/fix/CollectHeadOption.scala b/rules/src/main/scala/fix/CollectHeadOption.scala
index 21e08a6..40d5613 100644
--- a/rules/src/main/scala/fix/CollectHeadOption.scala
+++ b/rules/src/main/scala/fix/CollectHeadOption.scala
@@ -17,7 +17,12 @@ class CollectHeadOption extends SyntacticRule("CollectHeadOption") {
             ),
             Term.Name("headOption")
           ) =>
-        Patch.replaceTree(t, s"${obj}.collectFirst${func}")
+        func match {
+          case _: Term.PartialFunction =>
+            Patch.replaceTree(t, s"${obj}.collectFirst${func}")
+          case _ =>
+            Patch.replaceTree(t, s"${obj}.collectFirst(${func})")
+        }
     }
   }.asPatch
 }