Repository: square/anvil
Issue: compiler-utils Annotation toAnnotationSpec error

Anvil: 2.4.0-M2

While writing a custom `CodeGenerator` for anvil I've encountered a bug in `compiler-utils` module.

```kotlin
@IntKey(Int.MAX_VALUE)
class A
```

When calling `com.squareup.anvil.compiler.internal.reference.AnnotationReference#toAnnotationSpec` on this AnnotationReference I get an error:
```
Couldn't resolve FqName Int.MAX_VALUE for Psi element: Int.MAX_VALUE

        at com.squareup.anvil.compiler.internal.PsiUtilsKt.requireFqName(PsiUtils.kt:266)
        at com.squareup.anvil.compiler.internal.reference.AnnotationArgumentReference$Psi$value$2.invoke$parsePsiElement(AnnotationArgumentReference.kt:231)
        at com.squareup.anvil.compiler.internal.reference.AnnotationArgumentReference$Psi$value$2.invoke(AnnotationArgumentReference.kt:236)
        at kotlin.UnsafeLazyImpl.getValue(Lazy.kt:81)
        at com.squareup.anvil.compiler.internal.reference.AnnotationArgumentReference$Psi.getValue(AnnotationArgumentReference.kt:121)
        at com.squareup.anvil.compiler.internal.reference.AnnotationArgumentReference.value(AnnotationArgumentReference.kt:58)
        at com.squareup.anvil.compiler.internal.reference.AnnotationReference.toAnnotationSpec(AnnotationReference.kt:88)

```

Changed files:
compiler-utils/src/main/java/com/squareup/anvil/compiler/internal/DescriptorUtils.kt
compiler/src/test/java/com/squareup/anvil/compiler/internal/reference/AnnotationReferenceTest.kt
compiler-utils/src/main/java/com/squareup/anvil/compiler/internal/reference/AnnotationArgumentReference.kt

diff --git a/compiler-utils/src/main/java/com/squareup/anvil/compiler/internal/DescriptorUtils.kt b/compiler-utils/src/main/java/com/squareup/anvil/compiler/internal/DescriptorUtils.kt
index b2964cc..c5a1003 100644
--- a/compiler-utils/src/main/java/com/squareup/anvil/compiler/internal/DescriptorUtils.kt
+++ b/compiler-utils/src/main/java/com/squareup/anvil/compiler/internal/DescriptorUtils.kt
@@ -6,7 +6,9 @@ import com.squareup.anvil.annotations.ExperimentalAnvilApi
 import com.squareup.anvil.compiler.api.AnvilCompilationException
 import org.jetbrains.kotlin.descriptors.ClassDescriptor
 import org.jetbrains.kotlin.descriptors.ModuleDescriptor
+import org.jetbrains.kotlin.descriptors.PropertyDescriptor
 import org.jetbrains.kotlin.descriptors.findClassAcrossModuleDependencies
+import org.jetbrains.kotlin.incremental.components.NoLookupLocation
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.resolve.constants.ConstantValue
@@ -58,3 +60,13 @@ public fun ConstantValue<*>.argumentType(module: ModuleDescriptor): KotlinType {
       "Couldn't resolve class across module dependencies for class ID: $classId"
     )
 }
+
+/**
+ * Assumes that the FqName is a top-level property, e.g. com.squareup.CONSTANT.
+ */
+@ExperimentalAnvilApi
+public fun FqName.getContributedPropertyOrNull(module: ModuleDescriptor): PropertyDescriptor? {
+  return module.getPackage(parent()).memberScope
+    .getContributedVariables(shortName(), NoLookupLocation.FROM_BACKEND)
+    .singleOrNull()
+}
diff --git a/compiler-utils/src/main/java/com/squareup/anvil/compiler/internal/reference/AnnotationArgumentReference.kt b/compiler-utils/src/main/java/com/squareup/anvil/compiler/internal/reference/AnnotationArgumentReference.kt
index 83f2740..9060b4b 100644
--- a/compiler-utils/src/main/java/com/squareup/anvil/compiler/internal/reference/AnnotationArgumentReference.kt
+++ b/compiler-utils/src/main/java/com/squareup/anvil/compiler/internal/reference/AnnotationArgumentReference.kt
@@ -5,11 +5,13 @@ import com.squareup.anvil.compiler.internal.argumentType
 import com.squareup.anvil.compiler.internal.classDescriptor
 import com.squareup.anvil.compiler.internal.descendant
 import com.squareup.anvil.compiler.internal.fqNameOrNull
+import com.squareup.anvil.compiler.internal.getContributedPropertyOrNull
 import com.squareup.anvil.compiler.internal.reference.AnnotationArgumentReference.Descriptor
 import com.squareup.anvil.compiler.internal.reference.AnnotationArgumentReference.Psi
 import com.squareup.anvil.compiler.internal.requireFqName
 import org.jetbrains.kotlin.com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.lexer.KtTokens
+import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.psi.KtClass
 import org.jetbrains.kotlin.psi.KtClassLiteralExpression
@@ -21,6 +23,7 @@ import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.psi.KtNameReferenceExpression
 import org.jetbrains.kotlin.psi.KtNamedFunction
 import org.jetbrains.kotlin.psi.KtObjectDeclaration
+import org.jetbrains.kotlin.psi.KtPrefixExpression
 import org.jetbrains.kotlin.psi.KtProperty
 import org.jetbrains.kotlin.psi.KtStringTemplateEntry
 import org.jetbrains.kotlin.psi.KtStringTemplateExpression
@@ -199,20 +202,68 @@ public sealed class AnnotationArgumentReference {
         } ?: fail()
       }
 
+      fun resolvePrimitiveConstant(fqName: FqName): Any? {
+        // This could be a constant from a primitive type, e.g. Int.MAX_VALUE
+        val classFqName = fqName.parent()
+        val constantName = fqName.shortName().asString()
+
+        // If this constant is coming from a companion object, then we'll find it this way.
+        classFqName.toClassReferenceOrNull(module)
+          ?.let { if (it.isObject()) listOf(it) else it.companionObjects() }
+          ?.flatMap { it.properties }
+          ?.singleOrNull { it.name == constantName }
+          ?.let { property ->
+            return when (property) {
+              is PropertyReference.Descriptor ->
+                property.property.compileTimeInitializer?.value
+              is PropertyReference.Psi ->
+                property.property.initializer?.let { parsePrimitiveType(it.text) }
+            }
+          }
+
+        // Is the constant a top-level constant?
+        return fqName.getContributedPropertyOrNull(module)
+          ?.compileTimeInitializer
+          ?.value
+      }
+
+      fun resolveConstant(fqName: FqName): Any? {
+        return if (fqName.toClassReferenceOrNull(module) != null) {
+          // That's an enum constant.
+          fqName
+        } else {
+          // That's hopefully a constant for a primitive type.
+          resolvePrimitiveConstant(fqName)
+        }
+      }
+
+      fun resolveConstant(psiElement: PsiElement): Any? {
+        val fqName = psiElement.fqNameOrNull(module)
+          ?: (psiElement as? KtDotQualifiedExpression)
+            // That's necessary for constants like kotlin.Int.MAX_VALUE, where the real FqName is
+            // actually Kotlin.Int.Companion.MAX_VALUE.
+            ?.firstChild
+            ?.fqNameOrNull(module)
+            ?.descendant(psiElement.children.last().text)
+          ?: return null
+
+        return resolveConstant(fqName)
+      }
+
       fun parsePsiElement(psiElement: PsiElement): Any {
         return when (psiElement) {
           is KtClassLiteralExpression -> psiElement.requireFqName(module).toClassReference(module)
           is KtCollectionLiteralExpression ->
             psiElement.children.map { parsePsiElement(it) }.convertToArrayIfNeeded()
-          is KtConstantExpression -> parsePrimitiveType(psiElement.text)
+          is KtConstantExpression, is KtPrefixExpression -> parsePrimitiveType(psiElement.text)
           is KtStringTemplateExpression ->
             psiElement
               .getChildOfType<KtStringTemplateEntry>()
               ?.text
               ?: fail()
           is KtNameReferenceExpression -> {
-            // Those are likely enum values.
-            psiElement.fqNameOrNull(module)?.let { return it }
+            // Those are likely enum values or another primitive constant.
+            resolveConstant(psiElement)?.let { return it }
 
             // Check if it's a constant. Note that we don't return the constant FqName directly,
             // but return the actual value from the constant instead. That's also what the
@@ -228,7 +279,12 @@ public sealed class AnnotationArgumentReference {
             return psiElement.containingKtFile.packageFqName
               .descendant(psiElement.getReferencedName())
           }
-          is KtDotQualifiedExpression -> psiElement.requireFqName(module)
+          is KtDotQualifiedExpression -> {
+            // Maybe it's a fully qualified name, so check this too.
+            return resolveConstant(psiElement)
+              ?: resolveConstant(FqName(psiElement.text))
+              ?: fail()
+          }
           else -> fail()
         }
       }
diff --git a/compiler/src/test/java/com/squareup/anvil/compiler/internal/reference/AnnotationReferenceTest.kt b/compiler/src/test/java/com/squareup/anvil/compiler/internal/reference/AnnotationReferenceTest.kt
index 6f3822f..d0f58e0 100644
--- a/compiler/src/test/java/com/squareup/anvil/compiler/internal/reference/AnnotationReferenceTest.kt
+++ b/compiler/src/test/java/com/squareup/anvil/compiler/internal/reference/AnnotationReferenceTest.kt
@@ -137,10 +137,18 @@ class AnnotationReferenceTest {
   }
 
   @Test fun `an int annotation argument can be parsed`() {
+    @Suppress("RemoveRedundantQualifierName")
     compile(
       """
       package com.squareup.test
  
+      import com.squareup.test2.CONSTANT_2
+      import com.squareup.test2.Abc
+      import com.squareup.test2.Abc.Companion.CONSTANT_4
+      import com.squareup.test2.SomeObject
+      import com.squareup.test2.SomeObject.CONSTANT_3
+      import kotlin.Int.Companion.MAX_VALUE
+ 
       annotation class BindingKey(val value: Int)
       
       @BindingKey(1)
@@ -150,20 +158,117 @@ class AnnotationReferenceTest {
 
       @BindingKey(CONSTANT)
       interface SomeClass2
+      
+      @BindingKey(-5)
+      interface SomeClass3
+      
+      @BindingKey(MAX_VALUE)
+      interface SomeClass4
+      
+      @BindingKey(Int.MAX_VALUE)
+      interface SomeClass5
+      
+      @BindingKey(kotlin.Int.MAX_VALUE)
+      interface SomeClass6
+      
+      @BindingKey(CONSTANT_2)
+      interface SomeClass7
+      
+      @BindingKey(com.squareup.test2.CONSTANT_2)
+      interface SomeClass8
+      
+      @BindingKey(CONSTANT_3)
+      interface SomeClass9
+      
+      @BindingKey(SomeObject.CONSTANT_3)
+      interface SomeClass10
+      
+      @BindingKey(com.squareup.test2.SomeObject.CONSTANT_3)
+      interface SomeClass11
+      
+      @BindingKey(CONSTANT_4)
+      interface SomeClass12
+      
+      @BindingKey(Abc.CONSTANT_4)
+      interface SomeClass13
+      
+      @BindingKey(com.squareup.test2.Abc.CONSTANT_4)
+      interface SomeClass14
+      """,
+      """
+      package com.squareup.test2
+        
+      const val CONSTANT_2 = 2
+      
+      object SomeObject {
+        const val CONSTANT_3 = 3
+      }
+      
+      class Abc {
+        companion object {
+          const val CONSTANT_4 = 4
+        }
+      }
       """,
       allWarningsAsErrors = false,
       codeGenerators = listOf(
         simpleCodeGenerator { psiRef ->
-          if (psiRef.shortName in listOf("BindingKey"))
+          if (psiRef.shortName in listOf("BindingKey", "SomeObject", "Abc", "Companion"))
             return@simpleCodeGenerator null
 
           listOf(psiRef, psiRef.toDescriptorReference()).forEach { ref ->
-            val argument = ref.annotations.single().arguments.single()
-
-            assertThat(argument.value<Int>()).isEqualTo(1)
-
-            assertThat(ref.annotations.single().toAnnotationSpec().toString())
-              .isEqualTo("@com.squareup.test.BindingKey(value = 1)")
+            when (psiRef.shortName) {
+              "SomeClass1" -> {
+                val argument = ref.annotations.single().arguments.single()
+                assertThat(argument.value<Int>()).isEqualTo(1)
+
+                assertThat(ref.annotations.single().toAnnotationSpec().toString())
+                  .isEqualTo("@com.squareup.test.BindingKey(value = 1)")
+              }
+              "SomeClass2" -> {
+                val argument = ref.annotations.single().arguments.single()
+                assertThat(argument.value<Int>()).isEqualTo(1)
+
+                assertThat(ref.annotations.single().toAnnotationSpec().toString())
+                  .isEqualTo("@com.squareup.test.BindingKey(value = 1)")
+              }
+              "SomeClass3" -> {
+                val argument = ref.annotations.single().arguments.single()
+                assertThat(argument.value<Int>()).isEqualTo(-5)
+
+                assertThat(ref.annotations.single().toAnnotationSpec().toString())
+                  .isEqualTo("@com.squareup.test.BindingKey(value = -5)")
+              }
+              "SomeClass4", "SomeClass5", "SomeClass6" -> {
+                val argument = ref.annotations.single().arguments.single()
+                assertThat(argument.value<Int>()).isEqualTo(Int.MAX_VALUE)
+
+                assertThat(ref.annotations.single().toAnnotationSpec().toString())
+                  .isEqualTo("@com.squareup.test.BindingKey(value = 2147483647)")
+              }
+              "SomeClass7", "SomeClass8" -> {
+                val argument = ref.annotations.single().arguments.single()
+                assertThat(argument.value<Int>()).isEqualTo(2)
+
+                assertThat(ref.annotations.single().toAnnotationSpec().toString())
+                  .isEqualTo("@com.squareup.test.BindingKey(value = 2)")
+              }
+              "SomeClass9", "SomeClass10", "SomeClass11" -> {
+                val argument = ref.annotations.single().arguments.single()
+                assertThat(argument.value<Int>()).isEqualTo(3)
+
+                assertThat(ref.annotations.single().toAnnotationSpec().toString())
+                  .isEqualTo("@com.squareup.test.BindingKey(value = 3)")
+              }
+              "SomeClass12", "SomeClass13", "SomeClass14" -> {
+                val argument = ref.annotations.single().arguments.single()
+                assertThat(argument.value<Int>()).isEqualTo(4)
+
+                assertThat(ref.annotations.single().toAnnotationSpec().toString())
+                  .isEqualTo("@com.squareup.test.BindingKey(value = 4)")
+              }
+              else -> throw NotImplementedError(psiRef.shortName)
+            }
           }
 
           null