[IJ plugin] Improve compat->operationBased migration

intellij-plugin/src/main/kotlin/com/apollographql/ijplugin/refactoring/migration/compattooperationbased/CompatToOperationBasedCodegenMigrationProcessor.kt

@@ -17,7 +17,7 @@ class CompatToOperationBasedCodegenMigrationProcessor(project: Project) : Apollo
 
   override val migrationItems = listOf(
       UpdateCodegenInBuildKts,
-      RemoveFragmentsField,
       ReworkInlineFragmentFields,
+      RemoveFragmentsField,
   )
 }

intellij-plugin/src/main/kotlin/com/apollographql/ijplugin/refactoring/migration/compattooperationbased/item/RemoveFragmentsField.kt

@@ -1,25 +1,24 @@
 package com.apollographql.ijplugin.refactoring.migration.compattooperationbased.item
 
-import com.apollographql.ijplugin.refactoring.findInheritorsOfClass
 import com.apollographql.ijplugin.refactoring.findReferences
 import com.apollographql.ijplugin.refactoring.migration.item.MigrationItem
 import com.apollographql.ijplugin.refactoring.migration.item.MigrationItemUsageInfo
 import com.apollographql.ijplugin.refactoring.migration.item.toMigrationItemUsageInfo
 import com.intellij.openapi.project.Project
 import com.intellij.psi.PsiMigration
 import com.intellij.psi.search.GlobalSearchScope
-import org.jetbrains.kotlin.asJava.classes.KtLightClassBase
+import com.intellij.psi.util.childrenOfType
+import org.jetbrains.kotlin.psi.KtCallExpression
 import org.jetbrains.kotlin.psi.KtClass
-import org.jetbrains.kotlin.psi.KtDotQualifiedExpression
+import org.jetbrains.kotlin.psi.KtQualifiedExpression
+import org.jetbrains.kotlin.psi.KtValueArgument
+import org.jetbrains.kotlin.psi.KtValueArgumentList
+import org.jetbrains.kotlin.psi.KtValueArgumentName
 import org.jetbrains.kotlin.psi.psiUtil.findPropertyByName
 
 object RemoveFragmentsField : MigrationItem() {
   override fun findUsages(project: Project, migration: PsiMigration, searchScope: GlobalSearchScope): List<MigrationItemUsageInfo> {
-    val operationInheritors = findInheritorsOfClass(project, "com.apollographql.apollo3.api.Operation").filterIsInstance<KtLightClassBase>()
-    val fragmentDataInheritors = findInheritorsOfClass(project, "com.apollographql.apollo3.api.Fragment.Data").filterIsInstance<KtLightClassBase>()
-    val allModels = (operationInheritors + fragmentDataInheritors).flatMap {
-      it.kotlinOrigin?.body?.declarations.orEmpty().filterIsInstance<KtClass>()
-    }
+    val allModels: List<KtClass> = findAllModels(project)
     val fragmentsProperties = allModels.mapNotNull { model ->
       model.findPropertyByName("fragments")
     }
@@ -28,15 +27,25 @@ object RemoveFragmentsField : MigrationItem() {
     }
     return references
         .mapNotNull {
-          val parent = it.parent as? KtDotQualifiedExpression ?: return@mapNotNull null
-          when {
-            // fragments.x
-            parent.receiverExpression.text == "fragments" -> {
-              parent.toMigrationItemUsageInfo(true)
+          when (val parent = it.parent) {
+            is KtQualifiedExpression -> {
+              when {
+                // fragments.x
+                parent.receiverExpression.text == "fragments" -> {
+                  parent.toMigrationItemUsageInfo(true)
+                }
+                // x.fragments
+                parent.selectorExpression?.text == "fragments" -> {
+                  parent.toMigrationItemUsageInfo(false)
+                }
+
+                else -> null
+              }
             }
-            // x.fragments
-            parent.selectorExpression?.text == "fragments" -> {
-              parent.toMigrationItemUsageInfo(false)
+
+            is KtValueArgumentName -> {
+              // fragments = ...
+              (parent.parent as? KtValueArgument)?.toMigrationItemUsageInfo()
             }
 
             else -> null
@@ -45,13 +54,31 @@ object RemoveFragmentsField : MigrationItem() {
   }
 
   override fun performRefactoring(project: Project, migration: PsiMigration, usage: MigrationItemUsageInfo) {
-    val element = usage.element as KtDotQualifiedExpression
-    if (usage.attachedData()) {
-      // fragments.x -> x
-      element.replace(element.selectorExpression!!)
-    } else {
-      // x.fragments -> x
-      element.replace(element.receiverExpression)
+    when (val element = usage.element) {
+      is KtQualifiedExpression -> {
+        if (usage.attachedData()) {
+          // fragments.x -> x
+          element.replace(element.selectorExpression!!)
+        } else {
+          // x.fragments -> x
+          element.replace(element.receiverExpression)
+        }
+      }
+
+      // fragments = Xxx.Fragments(yyy = ..., xxx = ...) -> yyy = ..., xxx = ...
+      is KtValueArgument -> {
+        // Xxx.Fragments(yyy = ..., xxx = ...)
+        val callExpression = element.getArgumentExpression()?.childrenOfType<KtCallExpression>()?.firstOrNull()
+        // yyy = ..., xxx = ...
+        val enclosedArguments = callExpression?.valueArgumentList ?: return
+        val parentArgumentList = element.parent as KtValueArgumentList
+        for (enclosedArgument in enclosedArguments.arguments) {
+          parentArgumentList.addArgument(enclosedArgument)
+        }
+        parentArgumentList.removeArgument(element)
+      }
+
+      else -> {}
     }
   }
 }

intellij-plugin/src/main/kotlin/com/apollographql/ijplugin/refactoring/migration/compattooperationbased/item/ReworkInlineFragmentFields.kt

@@ -1,7 +1,6 @@
 package com.apollographql.ijplugin.refactoring.migration.compattooperationbased.item
 
 import com.apollographql.ijplugin.refactoring.findClassReferences
-import com.apollographql.ijplugin.refactoring.findInheritorsOfClass
 import com.apollographql.ijplugin.refactoring.findReferences
 import com.apollographql.ijplugin.refactoring.migration.item.MigrationItem
 import com.apollographql.ijplugin.refactoring.migration.item.MigrationItemUsageInfo
@@ -10,7 +9,6 @@ import com.intellij.openapi.project.Project
 import com.intellij.psi.PsiMigration
 import com.intellij.psi.PsiReference
 import com.intellij.psi.search.GlobalSearchScope
-import org.jetbrains.kotlin.asJava.classes.KtLightClassBase
 import org.jetbrains.kotlin.idea.base.utils.fqname.fqName
 import org.jetbrains.kotlin.nj2k.postProcessing.type
 import org.jetbrains.kotlin.psi.KtClass
@@ -21,12 +19,7 @@ import org.jetbrains.kotlin.psi.KtPsiFactory
 object ReworkInlineFragmentFields : MigrationItem() {
   override fun findUsages(project: Project, migration: PsiMigration, searchScope: GlobalSearchScope): List<MigrationItemUsageInfo> {
     val usageInfo = mutableListOf<MigrationItemUsageInfo>()
-
-    val operationInheritors = findInheritorsOfClass(project, "com.apollographql.apollo3.api.Operation").filterIsInstance<KtLightClassBase>()
-    val fragmentDataInheritors = findInheritorsOfClass(project, "com.apollographql.apollo3.api.Fragment.Data").filterIsInstance<KtLightClassBase>()
-    val allModels = (operationInheritors + fragmentDataInheritors).flatMap {
-      it.kotlinOrigin?.body?.declarations.orEmpty().filterIsInstance<KtClass>()
-    }
+    val allModels: List<KtClass> = findAllModels(project)
     val inlineFragmentProperties = allModels.flatMap { model ->
       model.inlineFragmentProperties()
     }

intellij-plugin/src/main/kotlin/com/apollographql/ijplugin/refactoring/migration/compattooperationbased/item/UpdateCodegenInBuildKts.kt

@@ -4,33 +4,39 @@ import com.apollographql.ijplugin.refactoring.migration.item.DeletesElements
 import com.apollographql.ijplugin.refactoring.migration.item.MigrationItem
 import com.apollographql.ijplugin.refactoring.migration.item.MigrationItemUsageInfo
 import com.apollographql.ijplugin.refactoring.migration.item.toMigrationItemUsageInfo
+import com.apollographql.ijplugin.util.cast
 import com.apollographql.ijplugin.util.findPsiFilesByName
+import com.apollographql.ijplugin.util.getMethodName
+import com.apollographql.ijplugin.util.unquoted
 import com.intellij.openapi.project.Project
-import com.intellij.psi.PsiFile
 import com.intellij.psi.PsiMigration
 import com.intellij.psi.search.GlobalSearchScope
-import com.intellij.psi.util.parentOfType
+import org.jetbrains.kotlin.psi.KtCallExpression
 import org.jetbrains.kotlin.psi.KtDotQualifiedExpression
 import org.jetbrains.kotlin.psi.KtFile
-import org.jetbrains.kotlin.psi.KtLiteralStringTemplateEntry
 import org.jetbrains.kotlin.psi.KtNameReferenceExpression
 import org.jetbrains.kotlin.psi.KtTreeVisitorVoid
 
 object UpdateCodegenInBuildKts : MigrationItem(), DeletesElements {
   override fun findUsages(project: Project, migration: PsiMigration, searchScope: GlobalSearchScope): List<MigrationItemUsageInfo> {
-    val buildGradleKtsFiles: List<PsiFile> = project.findPsiFilesByName("build.gradle.kts", searchScope)
     val usages = mutableListOf<MigrationItemUsageInfo>()
+    val buildGradleKtsFiles: List<KtFile> = project.findPsiFilesByName("build.gradle.kts", searchScope).filterIsInstance<KtFile>()
     for (file in buildGradleKtsFiles) {
-      if (file !is KtFile) continue
       file.accept(object : KtTreeVisitorVoid() {
-        override fun visitLiteralStringTemplateEntry(entry: KtLiteralStringTemplateEntry) {
-          super.visitLiteralStringTemplateEntry(entry)
-          // codegenModels.set("compat")
-          if (entry.text == "compat") {
-            val dotQualifiedExpression = entry.parentOfType<KtDotQualifiedExpression>() ?: return
-            if ((dotQualifiedExpression.receiverExpression as? KtNameReferenceExpression)?.getReferencedName() == "codegenModels") {
-              usages.add(dotQualifiedExpression.toMigrationItemUsageInfo())
-            }
+        override fun visitCallExpression(expression: KtCallExpression) {
+          super.visitCallExpression(expression)
+          if (expression.getMethodName() == "apollo") {
+            expression.accept(object : KtTreeVisitorVoid() {
+              override fun visitDotQualifiedExpression(expression: KtDotQualifiedExpression) {
+                super.visitDotQualifiedExpression(expression)
+                if (expression.receiverExpression.cast<KtNameReferenceExpression>()?.getReferencedName() == "codegenModels") {
+                  val argumentText = expression.selectorExpression.cast<KtCallExpression>()?.valueArguments?.firstOrNull()?.text ?: return
+                  if (argumentText.unquoted() == "compat" || argumentText.contains("MODELS_COMPAT")) {
+                    usages.add(expression.toMigrationItemUsageInfo())
+                  }
+                }
+              }
+            })
           }
         }
       })

intellij-plugin/src/main/kotlin/com/apollographql/ijplugin/refactoring/migration/compattooperationbased/item/findAllModels.kt

@@ -0,0 +1,15 @@
+package com.apollographql.ijplugin.refactoring.migration.compattooperationbased.item
+
+import com.apollographql.ijplugin.refactoring.findInheritorsOfClass
+import com.intellij.openapi.project.Project
+import org.jetbrains.kotlin.asJava.classes.KtLightClassBase
+import org.jetbrains.kotlin.psi.KtClass
+
+fun findAllModels(project: Project): List<KtClass> {
+  val operationInheritors = findInheritorsOfClass(project, "com.apollographql.apollo3.api.Operation").filterIsInstance<KtLightClassBase>()
+  val fragmentDataInheritors = findInheritorsOfClass(project, "com.apollographql.apollo3.api.Fragment.Data").filterIsInstance<KtLightClassBase>()
+  val allModels: List<KtClass> = (operationInheritors + fragmentDataInheritors).flatMap {
+    it.kotlinOrigin?.body?.declarations.orEmpty().filterIsInstance<KtClass>()
+  } + fragmentDataInheritors.map { it.kotlinOrigin }.filterIsInstance<KtClass>()
+  return allModels
+}

intellij-plugin/src/test/testData/migration/compat-to-operationbased/removeFragmentsField.kt

@@ -55,4 +55,20 @@ suspend fun main() {
 
   val id = data.launches.launches[0]!!.fragments.launchFields.id
   val id2 = data.launches.launches[0]!!.apply { fragments.launchFields.id }
+  val id3 = data.launches.launches[0]?.fragments?.launchFields.id
+
+  //@formatter:off
+  val launch = MyQuery.Launch(
+      __typename = "Launch",
+      fragments = Launch.Fragments(
+          launchFields = LaunchFields(
+              id = "id",
+              site = "site",
+              mission = LaunchFields.Mission(
+                  name = "name",
+                  missionPatch = "missionPatch",
+              ),
+          ),
+      ),
+  )
 }

intellij-plugin/src/test/testData/migration/compat-to-operationbased/removeFragmentsField_after.kt

@@ -55,4 +55,18 @@ suspend fun main() {
 
   val id = data.launches.launches[0]!!.launchFields.id
   val id2 = data.launches.launches[0]!!.apply { launchFields.id }
+  val id3 = data.launches.launches[0]?.launchFields.id
+
+  //@formatter:off
+  val launch = MyQuery.Launch(
+      __typename = "Launch",
+          launchFields = LaunchFields(
+              id = "id",
+              site = "site",
+              mission = LaunchFields.Mission(
+                  name = "name",
+                  missionPatch = "missionPatch",
+              ),
+          ),
+  )
 }

intellij-plugin/src/test/testData/migration/compat-to-operationbased/reworkInlineFragmentFields.kt

@@ -92,4 +92,22 @@ suspend fun main() {
   val asGranny: AsGranny = data.fruitList.list[0].asApple!!.asGranny!!
   val isGranny = asGranny.isGranny
   val id2 = data.fruitList.list[0].asApple!!.asGranny!!.id
+
+  val newAsGranny: FruitListQuery.AsApple = FruitListQuery.AsApple(
+      __typename = "Apple",
+      id = "id",
+      color = "color",
+      asGolden = FruitListQuery.AsGolden(
+          __typename = "Golden",
+          id = "id",
+          color = "color",
+          isGolden = true,
+      ),
+      asGranny = FruitListQuery.AsGranny(
+          __typename = "Granny",
+          id = "id",
+          color = "color",
+          isGranny = true,
+      ),
+  )
 }

intellij-plugin/src/test/testData/migration/compat-to-operationbased/reworkInlineFragmentFields_after.kt

@@ -92,4 +92,22 @@ suspend fun main() {
   val asGranny: OnGranny = data.fruitList.list[0].onApple!!.onGranny!!
   val isGranny = asGranny.isGranny
   val id2 = data.fruitList.list[0].onApple!!.onGranny!!.id
+
+  val newAsGranny: FruitListQuery.OnApple = FruitListQuery.OnApple(
+      __typename = "Apple",
+      id = "id",
+      color = "color",
+      onGolden = FruitListQuery.OnGolden(
+          __typename = "Golden",
+          id = "id",
+          color = "color",
+          isGolden = true,
+      ),
+      onGranny = FruitListQuery.OnGranny(
+          __typename = "Granny",
+          id = "id",
+          color = "color",
+          isGranny = true,
+      ),
+  )
 }

intellij-plugin/src/test/testData/migration/compat-to-operationbased/updateCodegenInBuildKts.gradle.kts

@@ -2,5 +2,7 @@ apollo {
   service("service") {
     packageName.set("com.example.rocketreserver")
     codegenModels.set("compat")
+    codegenModels.set(com.apollographql.apollo3.compiler.MODELS_COMPAT)
+    codegenModels.set(MODELS_COMPAT)
   }
 }