[Compose-CL] Autocomplete IDs in constraint array

When constraining an anchor. The first element in the array should be an
ID, different from current ID.

Refactored Pattern helper functions into PatternUtils.kt

Bug: 207030860
Test: completeConstraintIdsInArray
Change-Id: I4e729c21b3964401911dea07eafd635d82f7c163

6 files changed, 153 insertions, 29 deletions

diff --git a/compose-ide-plugin/src/com/android/tools/compose/code/completion/constraintlayout/Constants.kt b/compose-ide-plugin/src/com/android/tools/compose/code/completion/constraintlayout/Constants.kt
index f87e768a4fe..6b17dd6b5c3 100644
--- a/compose-ide-plugin/src/com/android/tools/compose/code/completion/constraintlayout/Constants.kt
+++ b/compose-ide-plugin/src/com/android/tools/compose/code/completion/constraintlayout/Constants.kt
@@ -18,6 +18,7 @@ package com.android.tools.compose.code.completion.constraintlayout
 internal object KeyWords {
   const val ConstraintSets = "ConstraintSets"
   const val Extends = "Extends"
+  const val ParentId = "parent"
 }
 
 /**
diff --git a/compose-ide-plugin/src/com/android/tools/compose/code/completion/constraintlayout/ConstraintLayoutJsonCompletionContributor.kt b/compose-ide-plugin/src/com/android/tools/compose/code/completion/constraintlayout/ConstraintLayoutJsonCompletionContributor.kt
index 0473e878500..a60523e5c2a 100644
--- a/compose-ide-plugin/src/com/android/tools/compose/code/completion/constraintlayout/ConstraintLayoutJsonCompletionContributor.kt
+++ b/compose-ide-plugin/src/com/android/tools/compose/code/completion/constraintlayout/ConstraintLayoutJsonCompletionContributor.kt
@@ -15,6 +15,7 @@
  */
 package com.android.tools.compose.code.completion.constraintlayout
 
+import com.android.tools.compose.code.completion.constraintlayout.provider.ConstraintIdsProvider
 import com.android.tools.compose.code.completion.constraintlayout.provider.ConstraintSetFieldsProvider
 import com.android.tools.compose.code.completion.constraintlayout.provider.ConstraintSetNamesProvider
 import com.android.tools.compose.code.completion.constraintlayout.provider.ConstraintsProvider
@@ -24,16 +25,10 @@ import com.intellij.codeInsight.completion.CompletionContributor
 import com.intellij.codeInsight.completion.CompletionParameters
 import com.intellij.codeInsight.completion.CompletionResultSet
 import com.intellij.codeInsight.completion.CompletionType
-import com.intellij.json.JsonElementTypes
 import com.intellij.json.JsonLanguage
-import com.intellij.json.psi.JsonProperty
-import com.intellij.json.psi.JsonReferenceExpression
 import com.intellij.json.psi.JsonStringLiteral
-import com.intellij.patterns.PlatformPatterns
-import com.intellij.patterns.PsiElementPattern
-import com.intellij.psi.PsiElement
 
-private const val BASE_DEPTH_FOR_LITERAL_IN_PROPERTY = 2
+internal const val BASE_DEPTH_FOR_LITERAL_IN_PROPERTY = 2
 
 /** Depth for a literal of a property of the list of ConstraintSets. With respect to the ConstraintSets root element. */
 private const val CONSTRAINT_SET_LIST_PROPERTY_DEPTH = BASE_DEPTH_FOR_LITERAL_IN_PROPERTY + BASE_DEPTH_FOR_LITERAL_IN_PROPERTY
@@ -42,7 +37,7 @@ private const val CONSTRAINT_SET_LIST_PROPERTY_DEPTH = BASE_DEPTH_FOR_LITERAL_IN
 private const val CONSTRAINT_SET_PROPERTY_DEPTH = CONSTRAINT_SET_LIST_PROPERTY_DEPTH + BASE_DEPTH_FOR_LITERAL_IN_PROPERTY
 
 /** Depth for a literal of a property of a Constraints block. With respect to the ConstraintSets root element. */
-private const val CONSTRAINT_BLOCK_PROPERTY_DEPTH = CONSTRAINT_SET_PROPERTY_DEPTH + BASE_DEPTH_FOR_LITERAL_IN_PROPERTY
+internal const val CONSTRAINT_BLOCK_PROPERTY_DEPTH = CONSTRAINT_SET_PROPERTY_DEPTH + BASE_DEPTH_FOR_LITERAL_IN_PROPERTY
 
 /**
  * [CompletionContributor] for the JSON5 format supported in ConstraintLayout-Compose (and MotionLayout).
@@ -67,10 +62,19 @@ class ConstraintLayoutJsonCompletionContributor : CompletionContributor() {
     extend(
       CompletionType.BASIC,
       // Complete ConstraintSet names in Extends keyword
-      jsonPropertyStringValue()
+      jsonStringValue()
         .withPropertyParentAtLevel(2, KeyWords.Extends),
       ConstraintSetNamesProvider
     )
+    extend(
+      CompletionType.BASIC,
+      // Complete IDs in the constraint array (first position)
+      jsonStringValue()
+        // First element in the array, ie: there is no PsiElement preceding the desired one at this level
+        .withParent(psiElement<JsonStringLiteral>().atIndexOfJsonArray(0))
+        .insideConstraintArray(),
+      ConstraintIdsProvider
+    )
   }
 
   override fun fillCompletionVariants(parameters: CompletionParameters, result: CompletionResultSet) {
@@ -82,22 +86,4 @@ class ConstraintLayoutJsonCompletionContributor : CompletionContributor() {
     }
     super.fillCompletionVariants(parameters, result)
   }
-}
-
-// region ConstraintLayout Pattern Helpers
-private fun jsonPropertyName() = PlatformPatterns.psiElement(JsonElementTypes.IDENTIFIER)
-
-private fun jsonPropertyStringValue() =
-  PlatformPatterns.psiElement(JsonElementTypes.SINGLE_QUOTED_STRING).withParent<JsonStringLiteral>()
-
-private fun PsiElementPattern<*, *>.withConstraintSetsParentAtLevel(level: Int) = withPropertyParentAtLevel(level, KeyWords.ConstraintSets)
-// endregion
-
-// region Kotlin Syntax Helpers
-private inline fun <reified T : PsiElement> psiElement() = PlatformPatterns.psiElement(T::class.java)
-
-private inline fun <reified T : PsiElement> PsiElementPattern<*, *>.withParent() = this.withParent(T::class.java)
-
-private fun PsiElementPattern<*, *>.withPropertyParentAtLevel(level: Int, name: String) =
-  this.withSuperParent(level, psiElement<JsonProperty>().withChild(psiElement<JsonReferenceExpression>().withText(name)).save(name))
-// endregion
\ No newline at end of file
+}
\ No newline at end of file
diff --git a/compose-ide-plugin/src/com/android/tools/compose/code/completion/constraintlayout/PatternUtils.kt b/compose-ide-plugin/src/com/android/tools/compose/code/completion/constraintlayout/PatternUtils.kt
new file mode 100644
index 00000000000..84f14bb15d9
--- /dev/null
+++ b/compose-ide-plugin/src/com/android/tools/compose/code/completion/constraintlayout/PatternUtils.kt
@@ -0,0 +1,78 @@
+/*
+ * Copyright (C) 2022 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.android.tools.compose.code.completion.constraintlayout
+
+import com.intellij.json.JsonElementTypes
+import com.intellij.json.psi.JsonArray
+import com.intellij.json.psi.JsonProperty
+import com.intellij.json.psi.JsonReferenceExpression
+import com.intellij.json.psi.JsonStringLiteral
+import com.intellij.json.psi.JsonValue
+import com.intellij.patterns.PatternCondition
+import com.intellij.patterns.PlatformPatterns
+import com.intellij.patterns.PsiElementPattern
+import com.intellij.patterns.StandardPatterns
+import com.intellij.psi.PsiElement
+import com.intellij.util.ProcessingContext
+
+// region ConstraintLayout Pattern Helpers
+internal fun jsonPropertyName() = PlatformPatterns.psiElement(JsonElementTypes.IDENTIFIER)
+
+internal fun jsonStringValue() =
+  PlatformPatterns.psiElement(JsonElementTypes.SINGLE_QUOTED_STRING).withParent<JsonStringLiteral>()
+
+internal fun PsiElementPattern<*, *>.withConstraintSetsParentAtLevel(level: Int) = withPropertyParentAtLevel(level, KeyWords.ConstraintSets)
+
+internal fun PsiElementPattern<*, *>.insideConstraintArray() =
+  withSuperParent(2, psiElement<JsonArray>())
+    .withSuperParent(
+      BASE_DEPTH_FOR_LITERAL_IN_PROPERTY + 1, // JsonArray adds one level
+      psiElement<JsonProperty>().withChild(
+        // The parent property name may only be a StandardAnchor
+        psiElement<JsonReferenceExpression>().withText(StandardPatterns.string().oneOf(StandardAnchor.values().map { it.keyWord }))
+      )
+    )
+    .withConstraintSetsParentAtLevel(CONSTRAINT_BLOCK_PROPERTY_DEPTH + 1) // JsonArray adds one level
+// endregion
+
+// region Kotlin Syntax Helpers
+internal inline fun <reified T : PsiElement> psiElement(): PsiElementPattern<T, PsiElementPattern.Capture<T>> =
+  PlatformPatterns.psiElement(T::class.java)
+
+internal inline fun <reified T : PsiElement> PsiElementPattern<*, *>.withParent() = this.withParent(T::class.java)
+
+/**
+ * Pattern such that when traversing up the tree from the current element, the element at [level] is a [JsonProperty]. Which name matches
+ * one of the given [names].
+ */
+internal fun PsiElementPattern<*, *>.withPropertyParentAtLevel(level: Int, vararg names: String) =
+  this.withSuperParent(level, psiElement<JsonProperty>().withChild(
+    psiElement<JsonReferenceExpression>().withText(StandardPatterns.string().oneOf(*names)))
+  )
+
+/**
+ * Verifies that the current element is at the given [index] of the elements contained by its [JsonArray] parent.
+ */
+internal fun <T : JsonValue> PsiElementPattern<T, PsiElementPattern.Capture<T>>.atIndexOfJsonArray(index: Int) =
+  with(object : PatternCondition<T>("atIndexOfJsonArray") {
+    override fun accepts(element: T, context: ProcessingContext?): Boolean {
+      val parent = element.context as? JsonArray ?: return false
+      val children = parent.valueList
+      val indexOfSelf = children.indexOf(element)
+      return index == indexOfSelf
+    }
+  })
+// endregion
\ No newline at end of file
diff --git a/compose-ide-plugin/src/com/android/tools/compose/code/completion/constraintlayout/provider/ConstraintSetCompletionProviders.kt b/compose-ide-plugin/src/com/android/tools/compose/code/completion/constraintlayout/provider/ConstraintSetCompletionProviders.kt
index 5c2cd09b578..0df772d2400 100644
--- a/compose-ide-plugin/src/com/android/tools/compose/code/completion/constraintlayout/provider/ConstraintSetCompletionProviders.kt
+++ b/compose-ide-plugin/src/com/android/tools/compose/code/completion/constraintlayout/provider/ConstraintSetCompletionProviders.kt
@@ -155,6 +155,30 @@ internal object ConstraintsProvider : BaseConstraintSetsCompletionProvider() {
   }
 }
 
+/**
+ * Provides IDs when autocompleting a constraint array.
+ *
+ * The ID may be either 'parent' or any of the declared IDs in all ConstraintSets, except the ID of the constraints block from which this
+ * provider was invoked.
+ */
+internal object ConstraintIdsProvider : BaseConstraintSetsCompletionProvider() {
+  override fun addCompletions(
+    constraintSetsPropertyModel: ConstraintSetsPropertyModel,
+    parameters: CompletionParameters,
+    result: CompletionResultSet
+  ) {
+    val possibleIds = constraintSetsPropertyModel.constraintSets.flatMap { it.declaredIds }.toCollection(HashSet())
+    // Parent ID should always be present
+    possibleIds.add(KeyWords.ParentId)
+    // Remove the current ID
+    getJsonPropertyParent(parameters)?.name?.let(possibleIds::remove)
+
+    possibleIds.forEach { id ->
+      result.addLookupElement(id)
+    }
+  }
+}
+
 private fun CompletionResultSet.addLookupElement(name: String, tailText: String? = null, format: InsertionFormat? = null) {
   var lookupBuilder = if (format == null) {
     LookupElementBuilder.create(name)
diff --git a/compose-ide-plugin/src/com/android/tools/compose/code/completion/constraintlayout/provider/model/ConstraintSetModel.kt b/compose-ide-plugin/src/com/android/tools/compose/code/completion/constraintlayout/provider/model/ConstraintSetModel.kt
index 7b8ee497657..26f45edb9f5 100644
--- a/compose-ide-plugin/src/com/android/tools/compose/code/completion/constraintlayout/provider/model/ConstraintSetModel.kt
+++ b/compose-ide-plugin/src/com/android/tools/compose/code/completion/constraintlayout/provider/model/ConstraintSetModel.kt
@@ -26,6 +26,11 @@ import com.intellij.json.psi.JsonStringLiteral
  */
 internal class ConstraintSetModel(jsonProperty: JsonProperty) : BaseJsonPropertyModel(jsonProperty) {
   /**
+   * List of properties that have a constraint block assigned to it.
+   */
+  private val propertiesWithConstraints = innerProperties.filter { it.name != KeyWords.Extends }
+
+  /**
    * Name of the ConstraintSet.
    */
   val name: String? = elementPointer.element?.name
@@ -36,12 +41,17 @@ internal class ConstraintSetModel(jsonProperty: JsonProperty) : BaseJsonProperty
   val extendsFrom: String? = (findProperty(KeyWords.Extends)?.value as? JsonStringLiteral)?.value
 
   /**
+   * List of IDs declared in this ConstraintSet.
+   */
+  val declaredIds = propertiesWithConstraints.map { it.name }
+
+  /**
    * The constraints (by widget ID) explicitly declared in this ConstraintSet.
    *
    * Note that it does not resolve constraints inherited from [extendsFrom].
    */
   val constraintsById: Map<String, ConstraintsModel> =
-    innerProperties.associate { property ->
+    propertiesWithConstraints.associate { property ->
       property.name to ConstraintsModel(property)
     }
 
diff --git a/compose-ide-plugin/testSrc/com/android/tools/compose/code/completion/constraintlayout/ConstraintLayoutJsonCompletionContributorTest.kt b/compose-ide-plugin/testSrc/com/android/tools/compose/code/completion/constraintlayout/ConstraintLayoutJsonCompletionContributorTest.kt
index 6a64949ff36..e20d6e7a845 100644
--- a/compose-ide-plugin/testSrc/com/android/tools/compose/code/completion/constraintlayout/ConstraintLayoutJsonCompletionContributorTest.kt
+++ b/compose-ide-plugin/testSrc/com/android/tools/compose/code/completion/constraintlayout/ConstraintLayoutJsonCompletionContributorTest.kt
@@ -124,6 +124,31 @@ internal class ConstraintLayoutJsonCompletionContributorTest {
   }
 
   @Test
+  fun completeConstraintIdsInArray() {
+    @Language("JSON5")
+    val content =
+      """
+      {
+        ConstraintSets: {
+          start: {
+            id1: {
+              start: ['$caret', 'start', 0]
+            },
+            id2: {},
+            id3: {}
+          }
+        }
+      }
+    """.trimIndent()
+    myFixture.configureByText("myscene.json", content)
+    myFixture.completeBasic()
+    val lookupElements = myFixture.lookupElementStrings!!
+
+    assertThat(lookupElements).hasSize(3)
+    assertThat(lookupElements).containsExactly("id2", "id3", "parent")
+  }
+
+  @Test
   fun constraintAnchorHandlerResult() {
     @Language("JSON5")
     val content =