Support Color picker for more Color declarations in Compose

Support open Color Picker for the constructions of
androidx.ui.graphics.Color class. After this CL only
Color(Float,Float,Float,Float,ColorSpaces) function is not supported.

Fixes: 176808581
Test: Covered and updated ComposeColorAnnotatorTest.kt
Change-Id: If0973984c94c9ab0a58be5fbc33a05a18038e489
diff --git a/compose-ide-plugin/src/com/android/tools/compose/ComposeColorAnnotator.kt b/compose-ide-plugin/src/com/android/tools/compose/ComposeColorAnnotator.kt
index 0cfb561..7e4a4b2 100644
--- a/compose-ide-plugin/src/com/android/tools/compose/ComposeColorAnnotator.kt
+++ b/compose-ide-plugin/src/com/android/tools/compose/ComposeColorAnnotator.kt
@@ -36,6 +36,7 @@
 import com.intellij.psi.PsiType
 import com.intellij.psi.impl.source.tree.LeafPsiElement
 import com.intellij.util.ui.ColorIcon
+import org.jetbrains.annotations.VisibleForTesting
 import org.jetbrains.kotlin.psi.KtCallElement
 import org.jetbrains.kotlin.psi.KtConstantExpression
 import org.jetbrains.uast.UCallExpression
@@ -129,7 +130,7 @@
  * picker.
  * Currently only updates the value of the Color declaration in the editor if it's using the [ComposeColorConstructor.INT] or
  * [ComposeColorConstructor.LONG].
- * TODO(lukeegan): Implement this for each of the Color parameter combinations
+ * TODO(lukeegan): Implement for ComposeColorConstructor.FLOAT_X4_COLORSPACE Color parameter
  */
 data class ColorIconRenderer(val element: UCallExpression, val color: Color) : GutterIconRenderer() {
   private val ICON_SIZE = 8
@@ -140,8 +141,14 @@
 
   override fun getClickAction(): AnAction? {
     val constructorType = getConstructorType(element.valueArguments) ?: return null
-    if (!constructorType.canBeOverwritten()) {
-      return null
+    val project = element.sourcePsi?.project ?: return null
+
+    val setColorTask: (Color) -> Unit = getSetColorTask(constructorType) ?: return null
+
+    val pickerListener = ColorPickerListener { color, _ ->
+      ApplicationManager.getApplication().invokeLater(Runnable {
+        WriteCommandAction.runWriteCommandAction(project, "Change Color", null, Runnable { setColorTask.invoke(color) })
+      }, project.disposed)
     }
     return object : AnAction() {
       override fun actionPerformed(e: AnActionEvent) {
@@ -155,7 +162,7 @@
             .addColorValuePanel().withFocus()
             .addSeparator()
             .addCustomComponent(MaterialColorPaletteProvider)
-            .addColorPickerListener(ColorPickerListener { color, _ -> setColorToAttribute(color) })
+            .addColorPickerListener(pickerListener)
             .focusWhenDisplay(true)
             .setFocusCycleRoot(true)
             .build()
@@ -165,27 +172,53 @@
     }
   }
 
-  private fun ComposeColorConstructor.canBeOverwritten(): Boolean {
-    return ComposeColorConstructor.INT == this || ComposeColorConstructor.LONG == this
+  @VisibleForTesting
+  fun getSetColorTask(constructorType: ComposeColorConstructor): ((Color) -> Unit)? {
+    return when(constructorType) {
+      ComposeColorConstructor.INT, ComposeColorConstructor.LONG -> { color -> setColorToAttribute(color) }
+      ComposeColorConstructor.INT_X3 -> { color -> setColorToAttribute(color.red, color.green, color.blue) }
+      ComposeColorConstructor.INT_X4 -> { color -> setColorToAttribute(color.red, color.green, color.blue, color.alpha) }
+      ComposeColorConstructor.FLOAT_X3 -> { color -> setColorToAttribute(color.red / 255f, color.green / 255f, color.blue / 255f) }
+      ComposeColorConstructor.FLOAT_X4 -> { color ->
+        setColorToAttribute(color.red / 255f, color.green / 255f, color.blue / 255f, color.alpha / 255f)
+      }
+      ComposeColorConstructor.FLOAT_X4_COLORSPACE -> return null // TODO: support ComposeColorConstructor.FLOAT_X4_COLORSPACE in the future.
+    }
   }
 
-  fun setColorToAttribute(color: Color) {
-    val constructorType = getConstructorType(element.valueArguments) ?: return
-    if (!constructorType.canBeOverwritten()) {
-      return
-    }
-    val runnable =
-      Runnable {
-        val hexString = "0x${(Integer.toHexString(color.rgb)).toUpperCase(Locale.getDefault())}"
-        val firstArgument = element.valueArguments[0].sourcePsi as? KtConstantExpression ?: return@Runnable
-        if ((firstArgument as PsiElement).isValid) {
-          (firstArgument.node?.firstChildNode as? LeafPsiElement)?.replaceWithText(hexString)
-        }
-      }
-    val project = element.sourcePsi?.project ?: return
-    ApplicationManager.getApplication().invokeLater(Runnable {
-      WriteCommandAction.runWriteCommandAction(project, "Change Color", null, runnable)
-    }, project.disposed)
+  private fun setColorToAttribute(color: Color) {
+    val hexString = "0x${(Integer.toHexString(color.rgb)).toUpperCase(Locale.getDefault())}"
+    (element.valueArguments[0].sourcePsi as? KtConstantExpression)?.replaceWithTextToFirstChildNode(hexString) ?: return
+  }
+
+  private fun setColorToAttribute(red: Int, green: Int, blue: Int) {
+    val valueArguments = element.valueArguments
+    (valueArguments[0].sourcePsi as? KtConstantExpression)?.replaceWithTextToFirstChildNode("$red") ?: return
+    (valueArguments[1].sourcePsi as? KtConstantExpression)?.replaceWithTextToFirstChildNode("$green") ?: return
+    (valueArguments[2].sourcePsi as? KtConstantExpression)?.replaceWithTextToFirstChildNode("$blue") ?: return
+  }
+
+  private fun setColorToAttribute(red: Int, green: Int, blue: Int, alpha: Int) {
+    val valueArguments = element.valueArguments
+    (valueArguments[0].sourcePsi as? KtConstantExpression)?.replaceWithTextToFirstChildNode("$red") ?: return
+    (valueArguments[1].sourcePsi as? KtConstantExpression)?.replaceWithTextToFirstChildNode("$green") ?: return
+    (valueArguments[2].sourcePsi as? KtConstantExpression)?.replaceWithTextToFirstChildNode("$blue") ?: return
+    (valueArguments[3].sourcePsi as? KtConstantExpression)?.replaceWithTextToFirstChildNode("$alpha") ?: return
+  }
+
+  private fun setColorToAttribute(red: Float, green: Float, blue: Float) {
+    val valueArguments = element.valueArguments
+    (valueArguments[0].sourcePsi as? KtConstantExpression)?.replaceWithTextToFirstChildNode("${red}f") ?: return
+    (valueArguments[1].sourcePsi as? KtConstantExpression)?.replaceWithTextToFirstChildNode("${green}f") ?: return
+    (valueArguments[2].sourcePsi as? KtConstantExpression)?.replaceWithTextToFirstChildNode("${blue}f") ?: return
+  }
+
+  private fun setColorToAttribute(red: Float, green: Float, blue: Float, alpha: Float) {
+    val valueArguments = element.valueArguments
+    (valueArguments[0].sourcePsi as? KtConstantExpression)?.replaceWithTextToFirstChildNode("${red}f") ?: return
+    (valueArguments[1].sourcePsi as? KtConstantExpression)?.replaceWithTextToFirstChildNode("${green}f") ?: return
+    (valueArguments[2].sourcePsi as? KtConstantExpression)?.replaceWithTextToFirstChildNode("${blue}f") ?: return
+    (valueArguments[3].sourcePsi as? KtConstantExpression)?.replaceWithTextToFirstChildNode("${alpha}f") ?: return
   }
 }
 
@@ -225,6 +258,14 @@
   return constant.value
 }
 
+private fun KtConstantExpression.replaceWithTextToFirstChildNode(text: String): Boolean {
+  if (!this.isValid) {
+    return false
+  }
+  (this.node.firstChildNode as? LeafPsiElement)?.replaceWithText(text) ?: return false
+  return true
+}
+
 enum class ComposeColorConstructor {
   INT, LONG, INT_X3, INT_X4, FLOAT_X3, FLOAT_X4, FLOAT_X4_COLORSPACE
 }
diff --git a/compose-ide-plugin/testSrc/com/android/tools/compose/ComposeColorAnnotatorTest.kt b/compose-ide-plugin/testSrc/com/android/tools/compose/ComposeColorAnnotatorTest.kt
index 05bb3c1..d2c5b3c 100644
--- a/compose-ide-plugin/testSrc/com/android/tools/compose/ComposeColorAnnotatorTest.kt
+++ b/compose-ide-plugin/testSrc/com/android/tools/compose/ComposeColorAnnotatorTest.kt
@@ -26,6 +26,8 @@
 import com.google.common.truth.Truth.assertThat
 import com.intellij.codeInsight.daemon.impl.AnnotationHolderImpl
 import com.intellij.lang.annotation.AnnotationSession
+import com.intellij.openapi.application.ApplicationManager
+import com.intellij.openapi.command.WriteCommandAction
 import com.intellij.psi.PsiElement
 import com.intellij.psi.util.parentOfType
 import com.intellij.testFramework.EdtRule
@@ -58,10 +60,12 @@
     StudioFlags.COMPOSE_EDITOR_SUPPORT.override(true)
     (myFixture.module.getModuleSystem() as DefaultModuleSystem).usesCompose = true
     myFixture.addClass(
-      //language=kotlin
+      //language=java
       """
-      package androidx.compose.ui.graphics
-      class ColorSpace
+      package androidx.compose.ui.graphics;
+      class ColorSpace {
+        public static final ColorSpace TEST_SPACE = ColorSpace();
+      }
       """)
     myFixture.addFileToProject(
       "src/com/androidx/compose/ui/graphics/Color.kt",
@@ -111,7 +115,7 @@
       ),
       includeClickAction = true
     )
-    setNewColor("Co|lor(0xFF4A8A7B)", Color(0xFFAABBCC.toInt()))
+    setNewColor("Co|lor(0xFF4A8A7B)", Color(0xFFAABBCC.toInt()), ComposeColorConstructor.LONG)
     assertThat(myFixture.editor.document.text).isEqualTo(
       //language=kotlin
       """
@@ -152,7 +156,7 @@
         Color(87, 173, 40, 0)),
       includeClickAction = true
     )
-    setNewColor("Co|lor(0x4A8A7B)", Color(0xFFAABBCC.toInt()))
+    setNewColor("Co|lor(0x4A8A7B)", Color(0xFFAABBCC.toInt()), ComposeColorConstructor.INT)
     assertThat(myFixture.editor.document.text).isEqualTo(
       //language=kotlin
       """
@@ -168,19 +172,8 @@
       """.trimIndent())
   }
 
-  private fun setNewColor(window: String, newColor: Color) {
-    runInEdtAndWait {
-      val element = myFixture.moveCaret(window)
-      val annotationHolder = AnnotationHolderImpl(AnnotationSession(myFixture.file))
-      annotationHolder.runAnnotatorWithContext(element.parentOfType<KtCallExpression>()!! as PsiElement, ComposeColorAnnotator())
-      val iconRenderer = annotationHolder[0].gutterIconRenderer as ColorIconRenderer
-      iconRenderer.setColorToAttribute(newColor)
-      PlatformTestUtil.dispatchAllInvocationEventsInIdeEventQueue()
-    }
-  }
-
   @Test
-  fun testColorIntx3x4() {
+  fun testColorInt_X3() {
     val psiFile = myFixture.addFileToProject(
       "src/com/android/test/A.kt",
       //language=kotlin
@@ -188,36 +181,10 @@
       package com.android.test
       import androidx.compose.ui.graphics.Color
       class A {
-        val other = Color(0.194f, 0f, 0.41f)
+        val other = Color(0xC2, 0x00, 0x29)
         fun () {
-          val primary = Color(0.74f, 0.138f, 0.3f, 0.845f)
-          val primaryVariant = Color(red = 0.87f, green = 0.173f, blue = 0.4f)
-        }
-      }
-      """.trimIndent())
-    myFixture.configureFromExistingVirtualFile(psiFile.virtualFile)
-    checkGutterIconInfos(
-      listOf(
-        Color(49, 0, 105),
-        Color(189, 35, 77, 215),
-        Color(222, 44, 102)),
-      includeClickAction = false
-    )
-  }
-
-  @Test
-  fun testFloatIntx3x4() {
-    val psiFile = myFixture.addFileToProject(
-      "src/com/android/test/A.kt",
-      //language=kotlin
-      """
-      package com.android.test
-      import androidx.compose.ui.graphics.Color
-      class A {
-        val other = Color(194, 0, 41)
-        fun () {
-          val primary = Color(74, 138, 123, 145)
-          val primaryVariant = Color(red = 87, green = 173, blue = 40)
+          val primary = Color(0x4A, 0x8A, 0x7B)
+          val primaryVariant = Color(red = 0x57, green = 0xAD, blue = 0x28)
         }
       }
       """.trimIndent())
@@ -225,12 +192,191 @@
     checkGutterIconInfos(
       listOf(
         Color(194, 0, 41),
-        Color(74, 138, 123, 145),
+        Color(74, 138, 123),
         Color(87, 173, 40)),
+      includeClickAction = true
+    )
+    setNewColor("Co|lor(0x4A, 0x8A, 0x7B)", Color(0xFFAABBCC.toInt()), ComposeColorConstructor.INT_X3)
+    assertThat(myFixture.editor.document.text).isEqualTo(
+      //language=kotlin
+      """
+      package com.android.test
+      import androidx.compose.ui.graphics.Color
+      class A {
+        val other = Color(0xC2, 0x00, 0x29)
+        fun () {
+          val primary = Color(170, 187, 204)
+          val primaryVariant = Color(red = 0x57, green = 0xAD, blue = 0x28)
+        }
+      }
+      """.trimIndent())
+  }
+
+  @Test
+  fun testColorInt_X4() {
+    val psiFile = myFixture.addFileToProject(
+      "src/com/android/test/A.kt",
+      //language=kotlin
+      """
+      package com.android.test
+      import androidx.compose.ui.graphics.Color
+      class A {
+        val other = Color(0xC2, 0x00, 0x29, 0xFF)
+        fun () {
+          val primary = Color(0x4A, 0x8A, 0x7B, 0xFF)
+          val primaryVariant = Color(red = 0x57, green = 0xAD, blue = 0x28, alpha = 0xFF)
+        }
+      }
+      """.trimIndent())
+    myFixture.configureFromExistingVirtualFile(psiFile.virtualFile)
+    checkGutterIconInfos(
+      listOf(
+        Color(194, 0, 41),
+        Color(74, 138, 123),
+        Color(87, 173, 40)),
+      includeClickAction = true
+    )
+    setNewColor("Co|lor(0x4A, 0x8A, 0x7B, 0xFF)", Color(0xFFAABBCC.toInt()), ComposeColorConstructor.INT_X4)
+    assertThat(myFixture.editor.document.text).isEqualTo(
+      //language=kotlin
+      """
+      package com.android.test
+      import androidx.compose.ui.graphics.Color
+      class A {
+        val other = Color(0xC2, 0x00, 0x29, 0xFF)
+        fun () {
+          val primary = Color(170, 187, 204, 255)
+          val primaryVariant = Color(red = 0x57, green = 0xAD, blue = 0x28, alpha = 0xFF)
+        }
+      }
+      """.trimIndent())
+  }
+
+
+  @Test
+  fun testColorFloat_X3() {
+    val psiFile = myFixture.addFileToProject(
+      "src/com/android/test/A.kt",
+      //language=kotlin
+      """
+      package com.android.test
+      import androidx.compose.ui.graphics.Color
+      class A {
+        val other = Color(0.14f, 0.0f, 0.16f)
+        fun () {
+          val primary = Color(0.3f, 0.54f, 0.48f)
+          val primaryVariant = Color(red = 0.34f, green = 0.68f, blue = 0.15f)
+        }
+      }
+      """.trimIndent())
+    myFixture.configureFromExistingVirtualFile(psiFile.virtualFile)
+    checkGutterIconInfos(
+      listOf(
+        Color(36, 0, 41),
+        Color(77, 138, 122),
+        Color(87, 173, 38)),
+      includeClickAction = true
+    )
+    setNewColor("Co|lor(0.3f, 0.54f, 0.48f)", Color(0xFFAABBCC.toInt()), ComposeColorConstructor.FLOAT_X3)
+    assertThat(myFixture.editor.document.text).isEqualTo(
+      //language=kotlin
+      """
+      package com.android.test
+      import androidx.compose.ui.graphics.Color
+      class A {
+        val other = Color(0.14f, 0.0f, 0.16f)
+        fun () {
+          val primary = Color(0.6666667f, 0.73333335f, 0.8f)
+          val primaryVariant = Color(red = 0.34f, green = 0.68f, blue = 0.15f)
+        }
+      }
+      """.trimIndent())
+  }
+
+  @Test
+  fun testColorFloat_X4() {
+    val psiFile = myFixture.addFileToProject(
+      "src/com/android/test/A.kt",
+      //language=kotlin
+      """
+      package com.android.test
+      import androidx.compose.ui.graphics.Color
+      class A {
+        val other = Color(0.194f, 0f, 0.41f, 0.5f)
+        fun () {
+          val primary = Color(0.74f, 0.138f, 0.3f, 0.845f)
+          val primaryVariant = Color(red = 0.87f, green = 0.173f, blue = 0.4f, alpha = 0.25f)
+        }
+      }
+      """.trimIndent())
+    myFixture.configureFromExistingVirtualFile(psiFile.virtualFile)
+    checkGutterIconInfos(
+      listOf(
+        Color(49, 0, 105, 128),
+        Color(189, 35, 77, 215),
+        Color(222, 44, 102, 64)),
+      includeClickAction = true
+    )
+    setNewColor("Co|lor(0.74f, 0.138f, 0.3f, 0.845f)", Color(0xFFAABBCC.toInt()), ComposeColorConstructor.FLOAT_X4)
+    assertThat(myFixture.editor.document.text).isEqualTo(
+      //language=kotlin
+      """
+      package com.android.test
+      import androidx.compose.ui.graphics.Color
+      class A {
+        val other = Color(0.194f, 0f, 0.41f, 0.5f)
+        fun () {
+          val primary = Color(0.6666667f, 0.73333335f, 0.8f, 1.0f)
+          val primaryVariant = Color(red = 0.87f, green = 0.173f, blue = 0.4f, alpha = 0.25f)
+        }
+      }
+      """.trimIndent())
+  }
+
+  @Test
+  fun testColorFloat_X4_ColorSpace() {
+    val psiFile = myFixture.addFileToProject(
+      "src/com/android/test/A.kt",
+      //language=kotlin
+      """
+      package com.android.test
+      import androidx.compose.ui.graphics.Color
+      import androidx.compose.ui.graphics.ColorSpace
+      class A {
+        val other = Color(0.194f, 0f, 0.41f, 0.5f, ColorSpace.TEST_SPACE)
+        fun () {
+          val primary = Color(0.74f, 0.138f, 0.3f, 0.845f, ColorSpace.TEST_SPACE)
+          val primaryVariant = Color(red = 0.87f, green = 0.173f, blue = 0.4f, alpha = 0.25f, colorSpace = ColorSpace.TEST_SPACE)
+        }
+      }
+      """.trimIndent())
+    myFixture.configureFromExistingVirtualFile(psiFile.virtualFile)
+    checkGutterIconInfos(
+      listOf(
+        Color(49, 0, 105, 128),
+        Color(189, 35, 77, 215),
+        Color(222, 44, 102, 64)),
       includeClickAction = false
     )
   }
 
+  private fun setNewColor(window: String, newColor: Color, constructorType: ComposeColorConstructor) {
+    runInEdtAndWait {
+      val element = myFixture.moveCaret(window)
+      val annotationHolder = AnnotationHolderImpl(AnnotationSession(myFixture.file))
+      annotationHolder.runAnnotatorWithContext(element.parentOfType<KtCallExpression>()!! as PsiElement, ComposeColorAnnotator())
+      val iconRenderer = annotationHolder[0].gutterIconRenderer as ColorIconRenderer
+      val project = myFixture.project
+
+      val setColorTask = iconRenderer.getSetColorTask(constructorType) ?: return@runInEdtAndWait
+      ApplicationManager.getApplication().invokeLater({
+        WriteCommandAction.runWriteCommandAction(project, "Change Color", null, { setColorTask.invoke(newColor) })
+      }, project.disposed)
+
+      PlatformTestUtil.dispatchAllInvocationEventsInIdeEventQueue()
+    }
+  }
+
   private fun checkGutterIconInfos(expectedColorIcons: List<Color>, includeClickAction: Boolean) {
     val iconList = myFixture.doHighlighting().filter { it.gutterIconRenderer is ColorIconRenderer }.sortedBy { it.startOffset }
     assertThat(iconList).hasSize(3)
@@ -292,6 +438,5 @@
     val icons = myFixture.findAllGutters()
     val colorGutterIconRenderer = icons.firstIsInstance<AndroidAnnotatorUtil.ColorRenderer>()
     assertThat((colorGutterIconRenderer.icon as MultipleColorIcon).colors).containsExactlyElementsIn(arrayOf(Color(63, 81, 181)))
-
   }
-}
\ No newline at end of file
+}