Repository: square/anvil
Issue: Regression: generic typealias expands only one type argument

Anvil 2.4.1
Kotlin 1.7.0

Given:
    
    class RealType<T, U>
    
    typealias AliasType<X> = RealType<String, X>
    
    class User {
        @Inject
        lateinit var value: AliasType<Int>
    }

Anvil turns `AliasType<Int>` into `RealType<Int>` in the generated code, which causes a build error because `RealType` takes two type arguments, not one.

This is a regression from Anvil 2.4.0 and Kotlin 1.6.21 where this issue did not occur.

Changed files:
compiler-utils/src/main/java/com/squareup/anvil/compiler/internal/reference/TypeReference.kt
compiler-utils/src/main/java/com/squareup/anvil/compiler/internal/reference/AnvilModuleDescriptor.kt
compiler/src/test/java/com/squareup/anvil/compiler/dagger/MembersInjectorGeneratorTest.kt
compiler/src/main/java/com/squareup/anvil/compiler/codegen/reference/RealAnvilModuleDescriptor.kt
compiler/src/test/java/com/squareup/anvil/compiler/internal/reference/TypeReferenceTest.kt

diff --git a/compiler-utils/src/main/java/com/squareup/anvil/compiler/internal/reference/AnvilModuleDescriptor.kt b/compiler-utils/src/main/java/com/squareup/anvil/compiler/internal/reference/AnvilModuleDescriptor.kt
index d78e993..f263c9e 100644
--- a/compiler-utils/src/main/java/com/squareup/anvil/compiler/internal/reference/AnvilModuleDescriptor.kt
+++ b/compiler-utils/src/main/java/com/squareup/anvil/compiler/internal/reference/AnvilModuleDescriptor.kt
@@ -6,6 +6,7 @@ import com.squareup.anvil.compiler.internal.reference.ClassReference.Descriptor
 import com.squareup.anvil.compiler.internal.reference.ClassReference.Psi
 import org.jetbrains.kotlin.descriptors.ClassDescriptor
 import org.jetbrains.kotlin.descriptors.ModuleDescriptor
+import org.jetbrains.kotlin.descriptors.TypeAliasDescriptor
 import org.jetbrains.kotlin.incremental.components.LookupLocation
 import org.jetbrains.kotlin.incremental.components.NoLookupLocation.FROM_BACKEND
 import org.jetbrains.kotlin.name.ClassId
@@ -22,6 +23,10 @@ public interface AnvilModuleDescriptor : ModuleDescriptor {
     lookupLocation: LookupLocation = FROM_BACKEND
   ): ClassDescriptor?
 
+  public fun resolveTypeAliasFqNameOrNull(
+    fqName: FqName
+  ): TypeAliasDescriptor?
+
   public fun getClassAndInnerClassReferences(ktFile: KtFile): List<Psi>
 
   public fun getClassReference(clazz: KtClassOrObject): Psi
diff --git a/compiler-utils/src/main/java/com/squareup/anvil/compiler/internal/reference/TypeReference.kt b/compiler-utils/src/main/java/com/squareup/anvil/compiler/internal/reference/TypeReference.kt
index 2e447f2..98d448d 100644
--- a/compiler-utils/src/main/java/com/squareup/anvil/compiler/internal/reference/TypeReference.kt
+++ b/compiler-utils/src/main/java/com/squareup/anvil/compiler/internal/reference/TypeReference.kt
@@ -128,7 +128,7 @@ public sealed class TypeReference {
       }
 
     override val unwrappedTypes: List<TypeReference> by lazy(NONE) {
-      type.typeElement!!.children
+      val unwrappedTypes = type.typeElement!!.children
         .filterIsInstance<KtTypeArgumentList>()
         .singleOrNull()
         ?.children
@@ -145,6 +145,8 @@ public sealed class TypeReference {
               .toTypeReference(declaringClass)
           }
         } ?: emptyList()
+
+      findUnwrappedTypesWithTypeAlias(unwrappedTypes)
     }
 
     @Suppress("RedundantNullableReturnType")
@@ -306,6 +308,42 @@ public sealed class TypeReference {
 
       return (typeElement ?: fail()).requireTypeName()
     }
+
+    private fun findUnwrappedTypesWithTypeAlias(
+      unwrappedTypes: List<TypeReference>
+    ): List<TypeReference> {
+      // For Psi we must inspect the type and check if further arguments have been specified in
+      // a type alias, e.g.
+      //
+      //   typealias MyType<X> = Pair<X, String>
+      //
+      //   val myType: MyType<Int>
+      //
+      // Without inspecting the type, Psi would only tell us that there's one argument `Int`. Type
+      // aliases are pretty much inlined, so we need to treat them specially. The idea is to check
+      // if the type is a type alias and if the type arguments differ. If no, then nothing special
+      // needs to happen. If yes, then we need to merge both type argument lists.
+
+      val children = type.typeElement?.children ?: return unwrappedTypes
+      if (children.size != 2 || children[1] !is KtTypeArgumentList) return unwrappedTypes
+
+      val typeAlias = children[0].fqNameOrNull(module)
+        ?.let { module.asAnvilModuleDescriptor().resolveTypeAliasFqNameOrNull(it) }
+        ?: return unwrappedTypes
+
+      val typeAliasArguments = typeAlias.defaultType.arguments
+      val delegatedTypeArguments = typeAlias.expandedType.arguments
+      if (typeAliasArguments.size == delegatedTypeArguments.size) return unwrappedTypes
+
+      return delegatedTypeArguments.map { typeProjection ->
+        val index = typeAliasArguments.indexOf(typeProjection)
+        if (index >= 0) {
+          unwrappedTypes[index]
+        } else {
+          typeProjection.type.toTypeReference(declaringClass)
+        }
+      }
+    }
   }
 
   public class Descriptor internal constructor(
diff --git a/compiler/src/main/java/com/squareup/anvil/compiler/codegen/reference/RealAnvilModuleDescriptor.kt b/compiler/src/main/java/com/squareup/anvil/compiler/codegen/reference/RealAnvilModuleDescriptor.kt
index 171c4ab..12046df 100644
--- a/compiler/src/main/java/com/squareup/anvil/compiler/codegen/reference/RealAnvilModuleDescriptor.kt
+++ b/compiler/src/main/java/com/squareup/anvil/compiler/codegen/reference/RealAnvilModuleDescriptor.kt
@@ -12,6 +12,7 @@ import com.squareup.anvil.compiler.internal.reference.ClassReference.Psi
 import com.squareup.anvil.compiler.internal.requireFqName
 import org.jetbrains.kotlin.descriptors.ClassDescriptor
 import org.jetbrains.kotlin.descriptors.ModuleDescriptor
+import org.jetbrains.kotlin.descriptors.TypeAliasDescriptor
 import org.jetbrains.kotlin.descriptors.findTypeAliasAcrossModuleDependencies
 import org.jetbrains.kotlin.descriptors.resolveClassByFqName
 import org.jetbrains.kotlin.incremental.components.LookupLocation
@@ -80,10 +81,14 @@ class RealAnvilModuleDescriptor private constructor(
     return resolveDescriptorCache.getOrPut(fqName) {
       // In the case of a typealias, we need to look up the original reference instead.
       resolveClassByFqName(fqName, lookupLocation)
-        ?: findTypeAliasAcrossModuleDependencies(fqName.classIdBestGuess())?.classDescriptor
+        ?: resolveTypeAliasFqNameOrNull(fqName)?.classDescriptor
     }
   }
 
+  override fun resolveTypeAliasFqNameOrNull(fqName: FqName): TypeAliasDescriptor? {
+    return findTypeAliasAcrossModuleDependencies(fqName.classIdBestGuess())
+  }
+
   override fun getClassReference(clazz: KtClassOrObject): Psi {
     return classReferenceCache.getOrPut(clazz.toClassReferenceCacheKey()) {
       Psi(clazz, clazz.toClassId(), this)
diff --git a/compiler/src/test/java/com/squareup/anvil/compiler/dagger/MembersInjectorGeneratorTest.kt b/compiler/src/test/java/com/squareup/anvil/compiler/dagger/MembersInjectorGeneratorTest.kt
index 613c8f3..f90bcec 100644
--- a/compiler/src/test/java/com/squareup/anvil/compiler/dagger/MembersInjectorGeneratorTest.kt
+++ b/compiler/src/test/java/com/squareup/anvil/compiler/dagger/MembersInjectorGeneratorTest.kt
@@ -2170,6 +2170,57 @@ public final class InjectClass_MembersInjector<T, U, V> implements MembersInject
     }
   }
 
+  @Test
+  fun `a member injector is generated for a class with a typealias injected parameter and different type arguments`() {
+
+    compile(
+      """
+      package com.squareup.test
+
+      import javax.inject.Inject
+
+      typealias AliasType<X> = Map<String, X>
+
+      class InjectClass {
+        @Inject lateinit var value: AliasType<Int>
+     
+        override fun equals(other: Any?): Boolean {
+          if (this === other) return true
+          if (javaClass != other?.javaClass) return false
+     
+          other as InjectClass
+     
+          if (value != other.value) return false
+     
+          return true
+        }
+        
+        override fun hashCode() = value.hashCode()
+      }
+      """
+    ) {
+      val injectClassMembersInjector = injectClass.membersInjector()
+
+      val constructor = injectClassMembersInjector.declaredConstructors.single()
+      assertThat(constructor.parameterTypes.toList())
+        .containsExactly(Provider::class.java)
+
+      val membersInjectorInstance = constructor
+        .newInstance(Provider { mapOf("a" to 1) }) as MembersInjector<Any>
+
+      val injectInstanceConstructor = injectClass.createInstance()
+      membersInjectorInstance.injectMembers(injectInstanceConstructor)
+
+      val injectInstanceStatic = injectClass.createInstance()
+
+      injectClassMembersInjector.staticInjectMethod("value")
+        .invoke(null, injectInstanceStatic, mapOf("a" to 1))
+
+      assertThat(injectInstanceConstructor).isEqualTo(injectInstanceStatic)
+      assertThat(injectInstanceConstructor).isNotSameInstanceAs(injectInstanceStatic)
+    }
+  }
+
   @Test
   fun `a member injector is generated for a class with a typealias superclass`() {
 
diff --git a/compiler/src/test/java/com/squareup/anvil/compiler/internal/reference/TypeReferenceTest.kt b/compiler/src/test/java/com/squareup/anvil/compiler/internal/reference/TypeReferenceTest.kt
index d171b55..13fc8b9 100644
--- a/compiler/src/test/java/com/squareup/anvil/compiler/internal/reference/TypeReferenceTest.kt
+++ b/compiler/src/test/java/com/squareup/anvil/compiler/internal/reference/TypeReferenceTest.kt
@@ -214,4 +214,52 @@ class TypeReferenceTest {
       assertThat(exitCode).isEqualTo(OK)
     }
   }
+
+  @Test fun `unwrapped types contain types from type aliases`() {
+    compile(
+      """
+      package com.squareup.test
+      
+      typealias AliasType<X> = Map<String, X>
+
+      typealias MyPair<X, Y> = Triple<X, String, Y>
+      typealias MySingle1<Z> = MyPair<Z, Long>
+      typealias MySingle2<Z> = Triple<Z, Z, Z>
+
+      class SomeClass {
+        lateinit var map: AliasType<Int>
+        lateinit var single1: MySingle1<Int>
+        lateinit var single2: MySingle2<Int>
+      }
+      """,
+      allWarningsAsErrors = false,
+      codeGenerators = listOf(
+        simpleCodeGenerator { psiRef ->
+          when (psiRef.shortName) {
+            "SomeClass" -> {
+              listOf(psiRef, psiRef.toDescriptorReference()).forEach { ref ->
+                assertThat(
+                  ref.properties.single { it.name == "map" }.type().unwrappedTypes
+                    .map { it.asClassReference().shortName }
+                ).containsExactly("String", "Int").inOrder()
+                assertThat(
+                  ref.properties.single { it.name == "single1" }.type().unwrappedTypes
+                    .map { it.asClassReference().shortName }
+                ).containsExactly("Int", "String", "Long").inOrder()
+                assertThat(
+                  ref.properties.single { it.name == "single2" }.type().unwrappedTypes
+                    .map { it.asClassReference().shortName }
+                ).containsExactly("Int", "Int", "Int").inOrder()
+              }
+            }
+            else -> throw NotImplementedError(psiRef.shortName)
+          }
+
+          null
+        }
+      )
+    ) {
+      assertThat(exitCode).isEqualTo(OK)
+    }
+  }
 }