Repository: square/anvil
Issue: Regression in 2.3.7 with generic assisted factories

we have a base interface for our factories that looks like this

```kotlin
interface AssistedWorkerFactory<T : ListenableWorker> {
  fun create(context: Context, params: WorkerParameters): T
}
```

with 2.3.7, this now results in this compilation exception

```
e: com.squareup.anvil.compiler.api.AnvilCompilationException: Back-end (JVM) Internal error: Invalid return type: slack.workmanager.workerfactory.AssistedWorkerFactory.create. An assisted factory's abstract method must return a type with an @AssistedInject-annotated constructor.
File being compiled: (17,3) in /mnt/buildkite-agent/builds/mbr-buildkite-agent-android-iad-ju4p-1/slack/android-pull-request/services/work-manager/src/main/kotlin/slack/workmanager/workerfactory/AssistedWorkerFactory.kt
 
	at com.squareup.anvil.compiler.codegen.dagger.AssistedFactoryGenerator$AssistedFactoryFunction$Companion$toAssistedFactoryFunction$returnTypeLazy$2.invoke(AssistedFactoryGenerator.kt:313)
```

I think this is an issue with it not properly reading the generic type. As an example, this is what we are doing

```kotlin
class FlushPendingExposuresWork @AssistedInject constructor(
  @Assisted appContext: Context,
  @Assisted params: WorkerParameters,
  // ...
) : CoroutineWorker(appContext, params) {

  override suspend fun doWork(): Result {
    // ...
  }

  @AssistedFactory
  interface Factory : AssistedWorkerFactory<FlushPendingExposuresWork>
}
```

Changed files:
compiler/src/main/java/com/squareup/anvil/compiler/codegen/dagger/AssistedFactoryGenerator.kt
compiler-utils/src/main/java/com/squareup/anvil/compiler/internal/PsiUtils.kt
compiler-utils/src/main/java/com/squareup/anvil/compiler/internal/AnvilModuleDescriptor.kt
compiler/src/test/java/com/squareup/anvil/compiler/dagger/AssistedFactoryGeneratorTest.kt
compiler-utils/src/main/java/com/squareup/anvil/compiler/internal/CompilerUtils.kt

diff --git a/compiler-utils/src/main/java/com/squareup/anvil/compiler/internal/AnvilModuleDescriptor.kt b/compiler-utils/src/main/java/com/squareup/anvil/compiler/internal/AnvilModuleDescriptor.kt
index 15a29af..d2c4d1b 100644
--- a/compiler-utils/src/main/java/com/squareup/anvil/compiler/internal/AnvilModuleDescriptor.kt
+++ b/compiler-utils/src/main/java/com/squareup/anvil/compiler/internal/AnvilModuleDescriptor.kt
@@ -1,6 +1,7 @@
 package com.squareup.anvil.compiler.internal
 
 import com.squareup.anvil.annotations.ExperimentalAnvilApi
+import com.squareup.anvil.compiler.api.AnvilCompilationException
 import org.jetbrains.kotlin.descriptors.ModuleDescriptor
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.FqName
@@ -29,6 +30,12 @@ public fun ModuleDescriptor.getKtClassOrObjectOrNull(
   fqName: FqName
 ): KtClassOrObject? = asAnvilModuleDescriptor().getKtClassOrObjectOrNull(fqName)
 
+@ExperimentalAnvilApi
+public fun ModuleDescriptor.requireKtClassOrObject(
+  fqName: FqName
+): KtClassOrObject = getKtClassOrObjectOrNull(fqName)
+  ?: throw AnvilCompilationException("Couldn't resolve KtClassOrObject for ${fqName.asString()}")
+
 @ExperimentalAnvilApi
 public fun ModuleDescriptor.canResolveFqName(
   fqName: FqName
diff --git a/compiler-utils/src/main/java/com/squareup/anvil/compiler/internal/CompilerUtils.kt b/compiler-utils/src/main/java/com/squareup/anvil/compiler/internal/CompilerUtils.kt
index ee22f0a..9dc2f74 100644
--- a/compiler-utils/src/main/java/com/squareup/anvil/compiler/internal/CompilerUtils.kt
+++ b/compiler-utils/src/main/java/com/squareup/anvil/compiler/internal/CompilerUtils.kt
@@ -21,6 +21,9 @@ import org.jetbrains.kotlin.resolve.constants.KClassValue
 import org.jetbrains.kotlin.resolve.constants.KClassValue.Value.NormalClass
 import org.jetbrains.kotlin.resolve.descriptorUtil.annotationClass
 import org.jetbrains.kotlin.resolve.descriptorUtil.fqNameSafe
+import org.jetbrains.kotlin.resolve.descriptorUtil.getAllSuperClassifiers
+import org.jetbrains.kotlin.resolve.descriptorUtil.getSuperClassNotAny
+import org.jetbrains.kotlin.resolve.descriptorUtil.getSuperInterfaces
 import org.jetbrains.kotlin.resolve.descriptorUtil.module
 import org.jetbrains.kotlin.types.ErrorType
 import org.jetbrains.kotlin.types.KotlinType
@@ -62,6 +65,16 @@ public fun ClassDescriptor.annotation(
   "Couldn't find $annotationFqName with scope $scope for $fqNameSafe."
 }
 
+/**
+ * Returns only the super class (excluding [Any]) and implemented interfaces declared directly by
+ * this class. This is different from [getAllSuperClassifiers] in that the latter returns the entire
+ * hierarchy.
+ */
+public fun ClassDescriptor.directSuperClassAndInterfaces(): List<ClassDescriptor> {
+  return listOfNotNull(getSuperClassNotAny())
+    .plus(getSuperInterfaces())
+}
+
 @ExperimentalAnvilApi
 public fun ConstantValue<*>.toType(
   module: ModuleDescriptor,
@@ -76,8 +89,9 @@ public fun ConstantValue<*>.toType(
 public fun KotlinType.argumentType(): KotlinType = arguments.first().type
 
 @ExperimentalAnvilApi
-public fun KotlinType.classDescriptorForType(): ClassDescriptor =
-  DescriptorUtils.getClassDescriptorForType(this)
+public fun KotlinType.classDescriptorForType(): ClassDescriptor {
+  return DescriptorUtils.getClassDescriptorForType(this)
+}
 
 @ExperimentalAnvilApi
 public fun AnnotationDescriptor.getAnnotationValue(key: String): ConstantValue<*>? =
diff --git a/compiler-utils/src/main/java/com/squareup/anvil/compiler/internal/PsiUtils.kt b/compiler-utils/src/main/java/com/squareup/anvil/compiler/internal/PsiUtils.kt
index 8e43e5b..5f34208 100644
--- a/compiler-utils/src/main/java/com/squareup/anvil/compiler/internal/PsiUtils.kt
+++ b/compiler-utils/src/main/java/com/squareup/anvil/compiler/internal/PsiUtils.kt
@@ -4,6 +4,8 @@ package com.squareup.anvil.compiler.internal
 
 import com.squareup.anvil.annotations.ExperimentalAnvilApi
 import com.squareup.anvil.compiler.api.AnvilCompilationException
+import com.squareup.anvil.compiler.internal.ClassReference.Descriptor
+import com.squareup.anvil.compiler.internal.ClassReference.Psi
 import com.squareup.kotlinpoet.TypeVariableName
 import org.jetbrains.kotlin.com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.descriptors.ClassDescriptor
@@ -31,12 +33,16 @@ import org.jetbrains.kotlin.psi.KtProperty
 import org.jetbrains.kotlin.psi.KtPureElement
 import org.jetbrains.kotlin.psi.KtSuperTypeListEntry
 import org.jetbrains.kotlin.psi.KtTypeArgumentList
+import org.jetbrains.kotlin.psi.KtTypeParameter
 import org.jetbrains.kotlin.psi.KtTypeReference
 import org.jetbrains.kotlin.psi.KtUserType
 import org.jetbrains.kotlin.psi.KtValueArgument
 import org.jetbrains.kotlin.psi.KtValueArgumentName
+import org.jetbrains.kotlin.psi.psiUtil.containingClass
+import org.jetbrains.kotlin.psi.psiUtil.getChildOfType
 import org.jetbrains.kotlin.psi.psiUtil.parents
 import org.jetbrains.kotlin.psi.psiUtil.parentsWithSelf
+import org.jetbrains.kotlin.resolve.descriptorUtil.fqNameSafe
 import kotlin.reflect.KClass
 
 private val kotlinAnnotations = listOf(jvmSuppressWildcardsFqName, publishedApiFqName)
@@ -220,34 +226,79 @@ public fun KtClassOrObject.allPsiSuperClasses(
 }
 
 /**
- * This will return all super types which are parsed as PSI, which means anything defined in this
- * module. This will include generated code, assuming it has already been generated. It **will not**
- * include any classes from other modules or external dependencies.
+ * Used to create a common type between [KtClassOrObject] class references and [ClassDescriptor]
+ * references, to streamline parsing.
  *
- * The first elements in the returned list represents the direct superclass to the receiver. The
- * last elements represent the types which are furthest up-stream.
+ * @see allSuperTypeClassReferences
+ */
+@ExperimentalAnvilApi
+public sealed interface ClassReference {
+
+  public val fqName: FqName
+
+  public data class Psi(
+    val clazz: KtClassOrObject,
+    override val fqName: FqName
+  ) : ClassReference
+
+  public data class Descriptor(
+    val clazz: ClassDescriptor,
+    override val fqName: FqName
+  ) : ClassReference
+}
+
+/**
+ * This will return all super types as [ClassReference], whether they're parsed as [KtClassOrObject]
+ * or [ClassDescriptor]. This will include generated code, assuming it has already been generated.
+ * The returned sequence will be distinct by FqName, and Psi types are preferred over Descriptors.
  *
- * If you have a class hierarchy which extends into dependencies, and you need to reference those
- * super classes, you'll need to take the last class in this list and resolve its `ClassDescriptor`,
- * then resolve its super-classes that way. Note that this will only work if the last class
- * in the list is not generated, so that it's parsed and part of the plain ModuleDescriptor.
+ * The first elements in the returned sequence represent the direct superclass to the receiver. The
+ * last elements represent the types which are furthest up-stream.
  *
- * **Note** This function returns all classes and interfaces. See [allPsiSuperClasses] for a version
- * which only returns classes.
+ * @param includeSelf If true, the receiver class is the first element of the sequence
  */
 @ExperimentalAnvilApi
-public fun KtClassOrObject.allPsiSuperTypes(
-  module: ModuleDescriptor
-): Sequence<KtClassOrObject> {
+public fun KtClassOrObject.allSuperTypeClassReferences(
+  module: ModuleDescriptor,
+  includeSelf: Boolean = false
+): Sequence<ClassReference> {
 
-  if (this is KtEnumEntry) return emptySequence()
+  fun KtClassOrObject.superTypeEntryNames() = superTypeListEntries
+    .map { it.requireFqName(module) }
 
-  val supers = superTypeListEntries
-    .asSequence()
-    .mapNotNull { it.typeReference?.fqNameOrNull(module) }
-    .mapNotNull { module.getKtClassOrObjectOrNull(it) }
+  fun FqName.typeRef(module: ModuleDescriptor): ClassReference {
+    return classDescriptorOrNull(module)
+      ?.let { Descriptor(it, this) }
+      ?: Psi(module.requireKtClassOrObject(this), this)
+  }
+
+  fun List<FqName>.typeRefs(module: ModuleDescriptor): Sequence<ClassReference> {
+    return asSequence().map { it.typeRef(module) }
+  }
 
-  return supers + supers.flatMap { it.allPsiSuperTypes(module) }
+  val initial = if (includeSelf) {
+    listOf(requireFqName()) + superTypeEntryNames()
+  } else {
+    superTypeEntryNames()
+  }.typeRefs(module)
+
+  return generateSequence(initial) { superTypes ->
+    superTypes
+      .mapNotNull { classRef ->
+        when (classRef) {
+          is Psi -> classRef.clazz.superTypeEntryNames()
+            .takeIf { it.isNotEmpty() }
+            ?.typeRefs(module)
+          is Descriptor -> classRef.clazz.directSuperClassAndInterfaces()
+            .takeIf { it.isNotEmpty() }
+            ?.asSequence()
+            ?.map { it.fqNameSafe.typeRef(module) }
+        }
+      }
+      .flatten()
+      .takeIf { it.firstOrNull() != null }
+  }.flatten()
+    .distinctBy { it.fqName }
 }
 
 @ExperimentalAnvilApi
@@ -535,6 +586,81 @@ public fun KtTypeReference.isFunctionType(): Boolean = typeElement is KtFunction
 @ExperimentalAnvilApi
 public fun KtClassOrObject.isGenericClass(): Boolean = typeParameterList != null
 
+@ExperimentalAnvilApi
+public fun KtTypeParameter.requireIdentifier(): String = nameIdentifier?.text
+  ?: throw AnvilCompilationException(
+    "unable to determine the name of this type parameter.", element = this
+  )
+
+/**
+ * Safely resolves a PSI [KtTypeReference], when that type reference may be a generic expressed by a
+ * type variable name. This is done by inspecting the class hierarchy to find where the generic type
+ * is declared, then resolving *that* reference.
+ *
+ * For instance, given:
+ *
+ * ```
+ * interface Factory<T> {
+ *   fun create(): T
+ * }
+ *
+ * interface ServiceFactory : Factory<Service>
+ * ```
+ *
+ * The KtTypeReference `T` will fail to resolve, since it isn't a type. This function will instead
+ * look to the `ServiceFactory` interface, then look at the supertype declaration in order to
+ * determine the type.
+ *
+ * @receiver The class which actually references the type. In the above example, this would be
+ *   `ServiceFactory`.
+ */
+@ExperimentalAnvilApi
+public fun KtClassOrObject.resolveTypeReference(
+  module: ModuleDescriptor,
+  typeReference: KtTypeReference
+): KtTypeReference? {
+
+  // if the element isn't a type variable name like `T`, it can be resolved through imports.
+  typeReference.typeElement?.fqNameOrNull(module)?.let { return typeReference }
+
+  val declaringClass = typeReference.typeElement?.containingClass()
+    ?: throw AnvilCompilationException(
+      "Unable to find a containing class.",
+      element = typeReference
+    )
+  val declaringClassFqName = declaringClass.requireFqName()
+
+  // If the class/interface declaring the generic is the receiver class,
+  // then the generic hasn't been set to a concrete type and can't be resolved.
+  if (requireFqName() == declaringClassFqName) {
+    return null
+  }
+
+  // Used to determine which parameter to look at in a KtTypeArgumentList
+  val indexToResolve = declaringClass.typeParameters
+    .indexOfFirst { it.requireIdentifier() == typeReference.typeElement?.text }
+
+  // Find where the supertype is actually declared by matching the FqName of the SuperTypeListEntry
+  // to the type which declares the generic we're trying to resolve.
+  // After finding the SuperTypeListEntry, just take the TypeReference from its type argument list
+  val resolvedTypeReference = allSuperTypeClassReferences(module, includeSelf = true)
+    .filterIsInstance<Psi>()
+    .firstNotNullOfOrNull { typeRef ->
+      typeRef.clazz
+        .superTypeListEntries
+        .firstOrNull { it.requireFqName(module) == declaringClassFqName }
+    }
+    ?.typeReference
+    ?.typeElement
+    ?.getChildOfType<KtTypeArgumentList>()
+    ?.arguments
+    ?.get(indexToResolve)
+    ?.typeReference
+
+  return resolvedTypeReference?.takeIf { it.fqNameOrNull(module) != null }
+    ?: resolvedTypeReference?.let { resolveTypeReference(module, it) }
+}
+
 @ExperimentalAnvilApi
 public fun KtCallableDeclaration.requireTypeReference(module: ModuleDescriptor): KtTypeReference {
   typeReference?.let { return it }
diff --git a/compiler/src/main/java/com/squareup/anvil/compiler/codegen/dagger/AssistedFactoryGenerator.kt b/compiler/src/main/java/com/squareup/anvil/compiler/codegen/dagger/AssistedFactoryGenerator.kt
index 4c68d76..79cd82d 100644
--- a/compiler/src/main/java/com/squareup/anvil/compiler/codegen/dagger/AssistedFactoryGenerator.kt
+++ b/compiler/src/main/java/com/squareup/anvil/compiler/codegen/dagger/AssistedFactoryGenerator.kt
@@ -13,12 +13,14 @@ import com.squareup.anvil.compiler.codegen.PrivateCodeGenerator
 import com.squareup.anvil.compiler.codegen.dagger.AssistedFactoryGenerator.AssistedFactoryFunction.Companion.toAssistedFactoryFunction
 import com.squareup.anvil.compiler.codegen.dagger.AssistedFactoryGenerator.AssistedParameterKey.Companion.toKeysList
 import com.squareup.anvil.compiler.codegen.injectConstructor
-import com.squareup.anvil.compiler.internal.allPsiSuperTypes
+import com.squareup.anvil.compiler.internal.ClassReference
+import com.squareup.anvil.compiler.internal.ClassReference.Descriptor
+import com.squareup.anvil.compiler.internal.ClassReference.Psi
+import com.squareup.anvil.compiler.internal.allSuperTypeClassReferences
 import com.squareup.anvil.compiler.internal.asClassName
 import com.squareup.anvil.compiler.internal.asTypeName
 import com.squareup.anvil.compiler.internal.buildFile
 import com.squareup.anvil.compiler.internal.classDescriptorForType
-import com.squareup.anvil.compiler.internal.classDescriptorOrNull
 import com.squareup.anvil.compiler.internal.classesAndInnerClass
 import com.squareup.anvil.compiler.internal.findAnnotation
 import com.squareup.anvil.compiler.internal.findAnnotationArgument
@@ -30,6 +32,7 @@ import com.squareup.anvil.compiler.internal.hasAnnotation
 import com.squareup.anvil.compiler.internal.isInterface
 import com.squareup.anvil.compiler.internal.requireFqName
 import com.squareup.anvil.compiler.internal.requireTypeName
+import com.squareup.anvil.compiler.internal.resolveTypeReference
 import com.squareup.anvil.compiler.internal.safePackageString
 import com.squareup.anvil.compiler.internal.typeVariableNames
 import com.squareup.kotlinpoet.ClassName
@@ -101,7 +104,8 @@ internal class AssistedFactoryGenerator : PrivateCodeGenerator() {
     val returnTypeFqName = function.returnTypeLazy.value
 
     // The return type of the function must have an @AssistedInject constructor.
-    val constructor = module.getKtClassOrObjectOrNull(returnTypeFqName)
+    val injectedClass = module.getKtClassOrObjectOrNull(returnTypeFqName)
+    val constructor = injectedClass
       ?.injectConstructor(assistedInjectFqName, module)
       ?: throw AnvilCompilationException(
         "Invalid return type: $returnTypeFqName. An assisted factory's abstract " +
@@ -125,7 +129,7 @@ internal class AssistedFactoryGenerator : PrivateCodeGenerator() {
 
     // Compute for each parameter its key.
     val functionParameterKeys = function.parameterKeys.value
-    val assistedParameterKeys = assistedParameters.toKeysList(module)
+    val assistedParameterKeys = assistedParameters.toKeysList(module, injectedClass)
 
     // The factory function may not have two or more parameters with the same key.
     val duplicateKeys = functionParameterKeys
@@ -281,6 +285,7 @@ internal class AssistedFactoryGenerator : PrivateCodeGenerator() {
       fun FunctionDescriptor.toAssistedFactoryFunction(
         factoryClass: KtClassOrObject
       ): AssistedFactoryFunction {
+
         val returnTypeLazy = lazy(NONE) {
           returnType
             ?.classDescriptorForType()
@@ -305,11 +310,16 @@ internal class AssistedFactoryGenerator : PrivateCodeGenerator() {
       }
 
       fun KtNamedFunction.toAssistedFactoryFunction(
-        module: ModuleDescriptor
+        module: ModuleDescriptor,
+        assistedFactoryClass: KtClassOrObject
       ): AssistedFactoryFunction {
         val returnTypeLazy = lazy(NONE) {
           typeReference
-            ?.fqNameOrNull(module)
+            ?.let { typeReference ->
+              assistedFactoryClass.resolveTypeReference(module, typeReference)
+                ?.fqNameOrNull(module)
+            }
+            ?: typeReference?.fqNameOrNull(module)
             ?: throw AnvilCompilationException(
               "Invalid return type: ${requireFqName()}. An assisted factory's abstract " +
                 "method must return a type with an @AssistedInject-annotated constructor.",
@@ -320,7 +330,12 @@ internal class AssistedFactoryGenerator : PrivateCodeGenerator() {
           valueParameters.map { param ->
 
             val typeName = param.typeReference
-              ?.requireTypeName(module)
+              ?.let { typeReference ->
+
+                assistedFactoryClass.resolveTypeReference(module, typeReference)
+                  ?.requireTypeName(module)
+                  ?: typeReference.requireTypeName(module)
+              }
               ?: throw AnvilCompilationException(
                 "Couldn't get type for parameter.",
                 element = param
@@ -329,10 +344,11 @@ internal class AssistedFactoryGenerator : PrivateCodeGenerator() {
             param.nameAsSafeName.asString() to typeName
           }
         }
+
         return AssistedFactoryFunction(
           name = name!!,
           returnTypeLazy = returnTypeLazy,
-          parameterKeys = lazy(NONE) { valueParameters.toKeysList(module) },
+          parameterKeys = lazy(NONE) { valueParameters.toKeysList(module, assistedFactoryClass) },
           parameterPairs = parameterPairs
         )
       }
@@ -343,53 +359,48 @@ internal class AssistedFactoryGenerator : PrivateCodeGenerator() {
     module: ModuleDescriptor
   ): AssistedFactoryFunction {
 
-    val psiSuperTypes = allPsiSuperTypes(module).toList()
-
-    // Go as far as we can with Psi, then at the end, get the FqName of the last Psi element.
-    // In case there's a deep hierarchy, we need to try to resolve super types using descriptors
-    val superFqName = psiSuperTypes.lastOrNull()
-      ?.fqName
-      ?: requireFqName()
-
-    // Try to resolve super-types using descriptors, but don't use "require" in case the code is
-    // generated.  It's important to try resolving the most "super" FqName that we know about,
-    // because the annotated class could be generated.
-    val functionsFromDescriptor = superFqName.classDescriptorOrNull(module)
-      ?.unsubstitutedMemberScope
-      ?.getContributedDescriptors(DescriptorKindFilter.FUNCTIONS)
-      ?.asSequence()
-      ?.filterIsInstance<FunctionDescriptor>()
-      ?.filter { it.modality == ABSTRACT }
-      ?.filter { it.visibility == PUBLIC || it.visibility == PROTECTED }
-      ?.map { it.toAssistedFactoryFunction(this) }
-      ?.toList()
-      .orEmpty()
+    fun ClassReference.assistedFunctions(): List<AssistedFactoryFunction> {
+      return when (this) {
+        is Psi -> {
+          val ktFunctions = clazz.takeIf { it.isInterface() }
+            ?.functions(false)
+            ?: clazz.functions(false)
+              .filter { it.hasModifier(KtTokens.ABSTRACT_KEYWORD) }
+              .filter { it.isPublic || it.isProtected() }
+
+          ktFunctions
+            .map { it.toAssistedFactoryFunction(module, this@requireSingleAbstractFunction) }
+        }
+        is Descriptor -> {
+          clazz.unsubstitutedMemberScope
+            .getContributedDescriptors(DescriptorKindFilter.FUNCTIONS)
+            .asSequence()
+            .filterIsInstance<FunctionDescriptor>()
+            .filter { it.modality == ABSTRACT }
+            .filter { it.visibility == PUBLIC || it.visibility == PROTECTED }
+            .map { it.toAssistedFactoryFunction(this@requireSingleAbstractFunction) }
+            .toList()
+        }
+      }
+    }
 
     // `clazz` must be first in the list because of `distinctBy { ... }`, which keeps the first
     // matched element.  If the function's inherited, it can be overridden as well.  Prioritizing
     // the version from the file we're parsing ensures the correct variance of the referenced types.
-    val functions = listOf(this)
-      .plus(psiSuperTypes)
-      .flatMap { type ->
-
-        type.takeIf { it.isInterface() }
-          ?.functions(false)
-          ?: type.functions(false)
-            .filter { it.hasModifier(KtTokens.ABSTRACT_KEYWORD) }
-            .filter { it.isPublic || it.isProtected() }
-      }
-      .map { it.toAssistedFactoryFunction(module) }
-      .plus(functionsFromDescriptor)
+    val assistedFunctions = allSuperTypeClassReferences(module, includeSelf = true)
+      .distinctBy { it.fqName }
+      .flatMap { it.assistedFunctions() }
       .distinctBy { it.name }
+      .toList()
 
     // Check for exact number of functions.
-    return when (functions.size) {
+    return when (assistedFunctions.size) {
       0 -> throw AnvilCompilationException(
         "The @AssistedFactory-annotated type is missing an abstract, non-default method " +
           "whose return type matches the assisted injection type.",
         element = this
       )
-      1 -> functions[0]
+      1 -> assistedFunctions[0]
       else -> throw AnvilCompilationException(
         "The @AssistedFactory-annotated type should contain a single abstract, non-default " +
           "method but found multiple.",
@@ -428,27 +439,34 @@ internal class AssistedFactoryGenerator : PrivateCodeGenerator() {
           )
         }
 
-      @JvmName("ughGenerics")
-      fun Collection<KtParameter>.toKeysList(module: ModuleDescriptor): List<AssistedParameterKey> =
-        map { parameter ->
+      fun Collection<KtParameter>.toKeysList(
+        module: ModuleDescriptor,
+        classForParameter: KtClassOrObject
+      ): List<AssistedParameterKey> = map { parameter ->
 
-          val typeName = parameter.typeReference?.requireTypeName(module)
-            ?: throw AnvilCompilationException(
-              "Couldn't get type for parameter.",
-              element = parameter
-            )
-
-          val identifier = parameter.findAnnotation(assistedFqName, module)
-            ?.findAnnotationArgument<KtStringTemplateExpression>("value", 0)
-            ?.getChildOfType<KtStringTemplateEntry>()
-            ?.text
-            ?: ""
+        val typeName = parameter.typeReference
+          ?.let { typeReference ->
 
-          AssistedParameterKey(
-            typeName = typeName,
-            identifier = identifier
+            classForParameter.resolveTypeReference(module, typeReference)
+              ?: typeReference
+          }
+          ?.requireTypeName(module)
+          ?: throw AnvilCompilationException(
+            "Couldn't get type for parameter.",
+            element = parameter
           )
-        }
+
+        val identifier = parameter.findAnnotation(assistedFqName, module)
+          ?.findAnnotationArgument<KtStringTemplateExpression>("value", 0)
+          ?.getChildOfType<KtStringTemplateEntry>()
+          ?.text
+          ?: ""
+
+        AssistedParameterKey(
+          typeName = typeName,
+          identifier = identifier
+        )
+      }
     }
   }
 }
diff --git a/compiler/src/test/java/com/squareup/anvil/compiler/dagger/AssistedFactoryGeneratorTest.kt b/compiler/src/test/java/com/squareup/anvil/compiler/dagger/AssistedFactoryGeneratorTest.kt
index 0ae5d86..55729a1 100644
--- a/compiler/src/test/java/com/squareup/anvil/compiler/dagger/AssistedFactoryGeneratorTest.kt
+++ b/compiler/src/test/java/com/squareup/anvil/compiler/dagger/AssistedFactoryGeneratorTest.kt
@@ -237,6 +237,82 @@ public final class AssistedServiceFactory_Impl implements AssistedServiceFactory
     }
   }
 
+  @Test fun `the factory function may be provided by a generic super type`() {
+    compile(
+      """
+      package com.squareup.test
+      
+      import dagger.assisted.Assisted
+      import dagger.assisted.AssistedFactory
+      import dagger.assisted.AssistedInject
+      
+      data class AssistedService @AssistedInject constructor(
+        val int: Int,
+        @Assisted val string: String
+      )
+      
+      interface Base<R: CharSequence, T> {
+        fun create(r :R): T
+      }
+      
+      interface Mid<T> : Base<String, T>
+      
+      @AssistedFactory
+      interface AssistedServiceFactory : Mid<AssistedService>
+      """
+    ) {
+      val factoryImplClass = assistedServiceFactory.implClass()
+      val generatedFactoryInstance = assistedService.factoryClass().createInstance(Provider { 5 })
+      val factoryImplInstance = factoryImplClass.createInstance(generatedFactoryInstance)
+
+      val staticMethods = factoryImplClass.declaredMethods.filter { it.isStatic }
+      assertThat(staticMethods).hasSize(1)
+
+      val factoryProvider = staticMethods.single { it.name == "create" }
+        .invoke(null, generatedFactoryInstance) as Provider<*>
+      assertThat(factoryProvider.get()::class.java).isEqualTo(factoryImplClass)
+
+      val assistedServiceInstance = factoryImplClass.declaredMethods
+        .filterNot { it.isStatic }
+        .last { it.name == "create" }
+        .invoke(factoryImplInstance, "Hello")
+
+      assertThat(assistedServiceInstance).isEqualTo(assistedService.createInstance(5, "Hello"))
+    }
+  }
+
+  @Test fun `the factory function may be provided by a generic super type from another module`() {
+    compile(
+      """
+      package com.squareup.test
+      
+      import dagger.assisted.Assisted
+      import dagger.assisted.AssistedFactory
+      import dagger.assisted.AssistedInject
+      import kotlin.properties.ReadOnlyProperty
+      import kotlin.reflect.KProperty
+      
+      class AssistedService @AssistedInject constructor(
+          @Assisted val thisRef: String,
+          @Assisted val property: KProperty<*>
+      )
+      
+      @AssistedFactory
+      interface AssistedServiceFactory : ReadOnlyProperty<String, AssistedService> 
+      """
+    ) {
+      val factoryImplClass = assistedServiceFactory.implClass()
+      val generatedFactoryInstance = assistedService.factoryClass().createInstance()
+
+      val staticMethods = factoryImplClass.declaredMethods.filter { it.isStatic }
+      assertThat(staticMethods).hasSize(1)
+
+      val factoryProvider = staticMethods.single { it.name == "create" }
+        .invoke(null, generatedFactoryInstance) as Provider<*>
+      assertThat(factoryProvider.get()::class.java).isEqualTo(factoryImplClass)
+    }
+  }
+
   @Test fun `the implementation function name matches the factory name`() {
     compile(
       """