Consolidate behaviour of receiver parameter between KAPT and KSP

This change adds a synthetic receiver parameter to KSP method elements that are extension functions. Similar to suspend functions the synthesised element makes it so that KAPT and KSP parameters are the same. This also fixes an issue in Javac / KAPT where the receiver parameter name was returning the next parameter name from the Kotlin metadata because metadata does not contain the receiver and there was an off by one error.

Bug: 160322705
Bug: 207676234
Test: XExecutableElementTest
Change-Id: I1b71df1ed6ae7dd14d012e0c12d9b28043a7847e
(cherry picked from commit 821535e5213c68f19bd31a0e7d7befe28196d08e)
Merged-In:I1b71df1ed6ae7dd14d012e0c12d9b28043a7847e
diff --git a/room/room-compiler-processing/src/main/java/androidx/room/compiler/processing/XMethodElement.kt b/room/room-compiler-processing/src/main/java/androidx/room/compiler/processing/XMethodElement.kt
index c85e818..600520d 100644
--- a/room/room-compiler-processing/src/main/java/androidx/room/compiler/processing/XMethodElement.kt
+++ b/room/room-compiler-processing/src/main/java/androidx/room/compiler/processing/XMethodElement.kt
@@ -98,11 +98,16 @@
     /**
      * Returns true if this is a suspend function.
      *
-     * @see XMethodType.getSuspendFunctionReturnType
+     * @see XSuspendMethodType
      */
     fun isSuspendFunction(): Boolean
 
     /**
+     * Returns true if this is an extension function.
+     */
+    fun isExtensionFunction(): Boolean
+
+    /**
      * Returns true if this method can be overridden without checking its enclosing [XElement].
      */
     fun isOverrideableIgnoringContainer(): Boolean {
diff --git a/room/room-compiler-processing/src/main/java/androidx/room/compiler/processing/XSuspendMethodType.kt b/room/room-compiler-processing/src/main/java/androidx/room/compiler/processing/XSuspendMethodType.kt
index 997efac..028c724 100644
--- a/room/room-compiler-processing/src/main/java/androidx/room/compiler/processing/XSuspendMethodType.kt
+++ b/room/room-compiler-processing/src/main/java/androidx/room/compiler/processing/XSuspendMethodType.kt
@@ -18,7 +18,7 @@
 
 interface XSuspendMethodType : XMethodType {
     /**
-     * IfReturns the real return type as seen by Kotlin.
+     * Returns the real suspend function return type as seen by Kotlin.
      */
     fun getSuspendFunctionReturnType(): XType
 }
\ No newline at end of file
diff --git a/room/room-compiler-processing/src/main/java/androidx/room/compiler/processing/javac/JavacConstructorElement.kt b/room/room-compiler-processing/src/main/java/androidx/room/compiler/processing/javac/JavacConstructorElement.kt
index 1eae5ff..5bc60e7 100644
--- a/room/room-compiler-processing/src/main/java/androidx/room/compiler/processing/javac/JavacConstructorElement.kt
+++ b/room/room-compiler-processing/src/main/java/androidx/room/compiler/processing/javac/JavacConstructorElement.kt
@@ -45,6 +45,19 @@
         element.requireEnclosingType(env)
     }
 
+    override val parameters: List<JavacMethodParameter> by lazy {
+        element.parameters.mapIndexed { index, variable ->
+            JavacMethodParameter(
+                env = env,
+                enclosingMethodElement = this,
+                containing = containing,
+                element = variable,
+                kotlinMetadataFactory = { kotlinMetadata?.parameters?.getOrNull(index) },
+                argIndex = index
+            )
+        }
+    }
+
     override val executableType: XConstructorType by lazy {
         val asMemberOf = env.typeUtils.asMemberOf(containing.type.typeMirror, element)
         JavacConstructorType(
diff --git a/room/room-compiler-processing/src/main/java/androidx/room/compiler/processing/javac/JavacExecutableElement.kt b/room/room-compiler-processing/src/main/java/androidx/room/compiler/processing/javac/JavacExecutableElement.kt
index c578727c..929fb02 100644
--- a/room/room-compiler-processing/src/main/java/androidx/room/compiler/processing/javac/JavacExecutableElement.kt
+++ b/room/room-compiler-processing/src/main/java/androidx/room/compiler/processing/javac/JavacExecutableElement.kt
@@ -39,18 +39,7 @@
         element.descriptor()
     }
 
-    override val parameters: List<JavacMethodParameter> by lazy {
-        element.parameters.mapIndexed { index, variable ->
-            JavacMethodParameter(
-                env = env,
-                enclosingMethodElement = this,
-                containing = containing,
-                element = variable,
-                kotlinMetadataFactory = { kotlinMetadata?.parameters?.getOrNull(index) },
-                argIndex = index
-            )
-        }
-    }
+    abstract override val parameters: List<JavacMethodParameter>
 
     override val equalityItems: Array<out Any?> by lazy {
         arrayOf(element, containing)
diff --git a/room/room-compiler-processing/src/main/java/androidx/room/compiler/processing/javac/JavacMethodElement.kt b/room/room-compiler-processing/src/main/java/androidx/room/compiler/processing/javac/JavacMethodElement.kt
index 66093c0..6337406 100644
--- a/room/room-compiler-processing/src/main/java/androidx/room/compiler/processing/javac/JavacMethodElement.kt
+++ b/room/room-compiler-processing/src/main/java/androidx/room/compiler/processing/javac/JavacMethodElement.kt
@@ -56,6 +56,22 @@
         element.requireEnclosingType(env)
     }
 
+    override val parameters: List<JavacMethodParameter> by lazy {
+        element.parameters.mapIndexed { index, variable ->
+            JavacMethodParameter(
+                env = env,
+                enclosingMethodElement = this,
+                containing = containing,
+                element = variable,
+                kotlinMetadataFactory = {
+                    val metadataParamIndex = if (isExtensionFunction()) index - 1 else index
+                    kotlinMetadata?.parameters?.getOrNull(metadataParamIndex)
+                },
+                argIndex = index
+            )
+        }
+    }
+
     override val kotlinMetadata: KmFunction? by lazy {
         (enclosingElement as? JavacTypeElement)?.kotlinMetadata?.getFunctionMetadata(element)
     }
@@ -104,6 +120,8 @@
 
     override fun isSuspendFunction() = kotlinMetadata?.isSuspend() == true
 
+    override fun isExtensionFunction() = kotlinMetadata?.isExtension() == true
+
     override fun overrides(other: XMethodElement, owner: XTypeElement): Boolean {
         check(other is JavacMethodElement)
         check(owner is JavacTypeElement)
diff --git a/room/room-compiler-processing/src/main/java/androidx/room/compiler/processing/javac/kotlin/KotlinClassMetadataUtils.kt b/room/room-compiler-processing/src/main/java/androidx/room/compiler/processing/javac/kotlin/KotlinClassMetadataUtils.kt
index d597e66..d2d1588 100644
--- a/room/room-compiler-processing/src/main/java/androidx/room/compiler/processing/javac/kotlin/KotlinClassMetadataUtils.kt
+++ b/room/room-compiler-processing/src/main/java/androidx/room/compiler/processing/javac/kotlin/KotlinClassMetadataUtils.kt
@@ -20,6 +20,7 @@
 import kotlinx.metadata.ClassName
 import kotlinx.metadata.Flag
 import kotlinx.metadata.Flags
+import kotlinx.metadata.KmAnnotation
 import kotlinx.metadata.KmClassVisitor
 import kotlinx.metadata.KmConstructorExtensionVisitor
 import kotlinx.metadata.KmConstructorVisitor
@@ -28,6 +29,7 @@
 import kotlinx.metadata.KmFunctionVisitor
 import kotlinx.metadata.KmPropertyExtensionVisitor
 import kotlinx.metadata.KmPropertyVisitor
+import kotlinx.metadata.KmTypeExtensionVisitor
 import kotlinx.metadata.KmTypeParameterVisitor
 import kotlinx.metadata.KmTypeVisitor
 import kotlinx.metadata.KmValueParameterVisitor
@@ -37,6 +39,7 @@
 import kotlinx.metadata.jvm.JvmFunctionExtensionVisitor
 import kotlinx.metadata.jvm.JvmMethodSignature
 import kotlinx.metadata.jvm.JvmPropertyExtensionVisitor
+import kotlinx.metadata.jvm.JvmTypeExtensionVisitor
 import kotlinx.metadata.jvm.KotlinClassMetadata
 import java.util.Locale
 
@@ -60,9 +63,11 @@
     val descriptor: String,
     private val flags: Flags,
     override val parameters: List<KmValueParameter>,
-    val returnType: KmType
+    val returnType: KmType,
+    val receiverType: KmType?
 ) : KmExecutable {
     fun isSuspend() = Flag.Function.IS_SUSPEND(flags)
+    fun isExtension() = receiverType != null
 }
 
 /**
@@ -91,10 +96,11 @@
 internal data class KmType(
     val flags: Flags,
     val typeArguments: List<KmType>,
-    val extendsBound: KmType?
+    val extendsBound: KmType?,
+    val isExtensionType: Boolean
 ) {
     fun isNullable() = Flag.Type.IS_NULLABLE(flags)
-    fun erasure(): KmType = KmType(flags, emptyList(), extendsBound?.erasure())
+    fun erasure(): KmType = KmType(flags, emptyList(), extendsBound?.erasure(), isExtensionType)
 }
 
 private data class KmTypeParameter(
@@ -105,7 +111,8 @@
     fun asKmType() = KmType(
         flags = flags,
         typeArguments = emptyList(),
-        extendsBound = extendsBound
+        extendsBound = extendsBound,
+        isExtensionType = false
     )
 }
 
@@ -136,6 +143,7 @@
             lateinit var methodSignature: JvmMethodSignature
             val parameters = mutableListOf<KmValueParameter>()
             lateinit var returnType: KmType
+            var receiverType: KmType? = null
 
             override fun visitValueParameter(
                 flags: Flags,
@@ -146,6 +154,12 @@
                 }
             }
 
+            override fun visitReceiverParameterType(flags: Flags): KmTypeVisitor? {
+                return TypeReader(flags) {
+                    receiverType = it
+                }
+            }
+
             override fun visitExtensions(type: KmExtensionType): KmFunctionExtensionVisitor {
                 if (type != JvmFunctionExtensionVisitor.TYPE) {
                     error("Unsupported extension type: $type")
@@ -171,7 +185,8 @@
                         descriptor = methodSignature.asString(),
                         flags = flags,
                         parameters = parameters,
-                        returnType = returnType
+                        returnType = returnType,
+                        receiverType = receiverType
                     )
                 )
             }
@@ -289,7 +304,8 @@
                                 descriptor = setterSignature.asString(),
                                 flags = 0,
                                 parameters = listOf(param),
-                                returnType = KM_VOID_TYPE
+                                returnType = KM_VOID_TYPE,
+                                receiverType = null
                             )
                         },
                         getter = getter?.let { getterSignature ->
@@ -299,7 +315,8 @@
                                 descriptor = getterSignature.asString(),
                                 flags = flags,
                                 parameters = emptyList(),
-                                returnType = returnType
+                                returnType = returnType,
+                                receiverType = null
                             )
                         }
                     )
@@ -353,6 +370,7 @@
 ) : KmTypeVisitor() {
     private val typeArguments = mutableListOf<KmType>()
     private var extendsBound: KmType? = null
+    private var isExtensionType = false
     override fun visitArgument(flags: Flags, variance: KmVariance): KmTypeVisitor {
         return TypeReader(flags) {
             typeArguments.add(it)
@@ -368,12 +386,24 @@
         }
     }
 
+    override fun visitExtensions(type: KmExtensionType): KmTypeExtensionVisitor? {
+        if (type != JvmTypeExtensionVisitor.TYPE) return null
+        return object : JvmTypeExtensionVisitor() {
+            override fun visitAnnotation(annotation: KmAnnotation) {
+                if (annotation.className == "kotlin/ExtensionFunctionType") {
+                    isExtensionType = true
+                }
+            }
+        }
+    }
+
     override fun visitEnd() {
         output(
             KmType(
                 flags = flags,
                 typeArguments = typeArguments,
-                extendsBound = extendsBound
+                extendsBound = extendsBound,
+                isExtensionType = isExtensionType
             )
         )
     }
@@ -443,7 +473,8 @@
                     typeArguments = typeParameters.map {
                         it.asKmType()
                     },
-                    extendsBound = null
+                    extendsBound = null,
+                    isExtensionType = false
                 ),
                 superType = superType
             )
@@ -515,5 +546,6 @@
 private val KM_VOID_TYPE = KmType(
     flags = 0,
     typeArguments = emptyList(),
-    extendsBound = null
+    extendsBound = null,
+    isExtensionType = false
 )
\ No newline at end of file
diff --git a/room/room-compiler-processing/src/main/java/androidx/room/compiler/processing/ksp/KspConstructorElement.kt b/room/room-compiler-processing/src/main/java/androidx/room/compiler/processing/ksp/KspConstructorElement.kt
index e8ff39e..f3d3fa0 100644
--- a/room/room-compiler-processing/src/main/java/androidx/room/compiler/processing/ksp/KspConstructorElement.kt
+++ b/room/room-compiler-processing/src/main/java/androidx/room/compiler/processing/ksp/KspConstructorElement.kt
@@ -18,6 +18,7 @@
 
 import androidx.room.compiler.processing.XConstructorElement
 import androidx.room.compiler.processing.XConstructorType
+import androidx.room.compiler.processing.XExecutableParameterElement
 import androidx.room.compiler.processing.XType
 import com.google.devtools.ksp.symbol.KSFunctionDeclaration
 
@@ -36,6 +37,17 @@
             ?: error("Constructor parent must be a type element $this")
     }
 
+    override val parameters: List<XExecutableParameterElement> by lazy {
+        declaration.parameters.mapIndexed { index, param ->
+            KspExecutableParameterElement(
+                env = env,
+                enclosingMethodElement = this,
+                parameter = param,
+                parameterIndex = index
+            )
+        }
+    }
+
     override val executableType: XConstructorType by lazy {
         KspConstructorType(
             env = env,
diff --git a/room/room-compiler-processing/src/main/java/androidx/room/compiler/processing/ksp/KspExecutableElement.kt b/room/room-compiler-processing/src/main/java/androidx/room/compiler/processing/ksp/KspExecutableElement.kt
index 2361bb7..50f23d0 100644
--- a/room/room-compiler-processing/src/main/java/androidx/room/compiler/processing/ksp/KspExecutableElement.kt
+++ b/room/room-compiler-processing/src/main/java/androidx/room/compiler/processing/ksp/KspExecutableElement.kt
@@ -18,7 +18,6 @@
 
 import androidx.room.compiler.processing.XAnnotated
 import androidx.room.compiler.processing.XExecutableElement
-import androidx.room.compiler.processing.XExecutableParameterElement
 import androidx.room.compiler.processing.XHasModifiers
 import androidx.room.compiler.processing.XType
 import androidx.room.compiler.processing.ksp.KspAnnotated.UseSiteFilter.Companion.NO_USE_SITE
@@ -52,17 +51,6 @@
         declaration.requireEnclosingMemberContainer(env)
     }
 
-    override val parameters: List<XExecutableParameterElement> by lazy {
-        declaration.parameters.mapIndexed { index, param ->
-            KspExecutableParameterElement(
-                env = env,
-                enclosingMethodElement = this,
-                parameter = param,
-                parameterIndex = index
-            )
-        }
-    }
-
     @OptIn(KspExperimental::class)
     override val thrownTypes: List<XType> by lazy {
         env.resolver.getJvmCheckedException(declaration).map {
diff --git a/room/room-compiler-processing/src/main/java/androidx/room/compiler/processing/ksp/KspMethodElement.kt b/room/room-compiler-processing/src/main/java/androidx/room/compiler/processing/ksp/KspMethodElement.kt
index dfa6955..2db3a0b 100644
--- a/room/room-compiler-processing/src/main/java/androidx/room/compiler/processing/ksp/KspMethodElement.kt
+++ b/room/room-compiler-processing/src/main/java/androidx/room/compiler/processing/ksp/KspMethodElement.kt
@@ -22,6 +22,7 @@
 import androidx.room.compiler.processing.XType
 import androidx.room.compiler.processing.XTypeElement
 import androidx.room.compiler.processing.ksp.synthetic.KspSyntheticContinuationParameterElement
+import androidx.room.compiler.processing.ksp.synthetic.KspSyntheticReceiverParameterElement
 import com.google.devtools.ksp.KspExperimental
 import com.google.devtools.ksp.closestClassDeclaration
 import com.google.devtools.ksp.symbol.ClassKind
@@ -52,6 +53,32 @@
         jvmName.getOrNull() ?: declaration.simpleName.asString()
     }
 
+    override val parameters: List<XExecutableParameterElement> by lazy {
+        buildList<XExecutableParameterElement> {
+            val extensionReceiver = declaration.extensionReceiver
+            if (extensionReceiver != null) {
+                // Synthesize the receiver parameter to be consistent with KAPT
+                add(
+                    KspSyntheticReceiverParameterElement(
+                        env = env,
+                        enclosingMethodElement = this@KspMethodElement,
+                        receiverType = extensionReceiver,
+                    )
+                )
+            }
+            addAll(
+                declaration.parameters.mapIndexed { index, param ->
+                    KspExecutableParameterElement(
+                        env = env,
+                        enclosingMethodElement = this@KspMethodElement,
+                        parameter = param,
+                        parameterIndex = index
+                    )
+                }
+            )
+        }
+    }
+
     override val executableType: XMethodType by lazy {
         KspMethodType.create(
             env = env,
@@ -101,6 +128,8 @@
             !isPrivate()
     }
 
+    override fun isExtensionFunction() = declaration.extensionReceiver != null
+
     override fun overrides(other: XMethodElement, owner: XTypeElement): Boolean {
         return env.resolver.overrides(this, other)
     }
diff --git a/room/room-compiler-processing/src/main/java/androidx/room/compiler/processing/ksp/synthetic/KspSyntheticPropertyMethodElement.kt b/room/room-compiler-processing/src/main/java/androidx/room/compiler/processing/ksp/synthetic/KspSyntheticPropertyMethodElement.kt
index 6f11ebc..71675b6 100644
--- a/room/room-compiler-processing/src/main/java/androidx/room/compiler/processing/ksp/synthetic/KspSyntheticPropertyMethodElement.kt
+++ b/room/room-compiler-processing/src/main/java/androidx/room/compiler/processing/ksp/synthetic/KspSyntheticPropertyMethodElement.kt
@@ -82,6 +82,8 @@
 
     final override fun isSuspendFunction() = false
 
+    final override fun isExtensionFunction() = false
+
     final override val enclosingElement: XMemberContainer
         get() = this.field.enclosingElement
 
diff --git a/room/room-compiler-processing/src/main/java/androidx/room/compiler/processing/ksp/synthetic/KspSyntheticReceiverParameterElement.kt b/room/room-compiler-processing/src/main/java/androidx/room/compiler/processing/ksp/synthetic/KspSyntheticReceiverParameterElement.kt
new file mode 100644
index 0000000..3748728
--- /dev/null
+++ b/room/room-compiler-processing/src/main/java/androidx/room/compiler/processing/ksp/synthetic/KspSyntheticReceiverParameterElement.kt
@@ -0,0 +1,102 @@
+/*
+ * Copyright 2021 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 androidx.room.compiler.processing.ksp.synthetic
+
+import androidx.room.compiler.processing.XAnnotated
+import androidx.room.compiler.processing.XEquality
+import androidx.room.compiler.processing.XExecutableParameterElement
+import androidx.room.compiler.processing.XType
+import androidx.room.compiler.processing.ksp.KspAnnotated
+import androidx.room.compiler.processing.ksp.KspJvmTypeResolutionScope
+import androidx.room.compiler.processing.ksp.KspMethodElement
+import androidx.room.compiler.processing.ksp.KspProcessingEnv
+import androidx.room.compiler.processing.ksp.KspType
+import com.google.devtools.ksp.symbol.KSTypeReference
+
+internal class KspSyntheticReceiverParameterElement(
+    val env: KspProcessingEnv,
+    override val enclosingMethodElement: KspMethodElement,
+    val receiverType: KSTypeReference,
+) : XExecutableParameterElement,
+    XEquality,
+    XAnnotated by KspAnnotated.create(
+        env = env,
+        delegate = null, // does not matter, this is synthetic and has no annotations.
+        filter = KspAnnotated.UseSiteFilter.NO_USE_SITE
+    ) {
+
+    override val name: String by lazy {
+        // KAPT uses `$this$<functionName>`
+        "$" + "this" + "$" + enclosingMethodElement.name
+    }
+
+    override val equalityItems: Array<out Any?> by lazy {
+        arrayOf(enclosingMethodElement, receiverType)
+    }
+
+    override val hasDefaultValue: Boolean
+        get() = false
+
+    private val jvmTypeResolutionScope by lazy {
+        KspJvmTypeResolutionScope.MethodParameter(
+            kspExecutableElement = enclosingMethodElement,
+            parameterIndex = 0, // Receiver param is the 1st one
+            annotated = enclosingMethodElement.declaration
+        )
+    }
+
+    override val type: XType by lazy {
+        env.wrap(receiverType).withJvmTypeResolver(jvmTypeResolutionScope)
+    }
+
+    override val fallbackLocationText: String
+        get() = "receiver parameter of ${enclosingMethodElement.fallbackLocationText}"
+
+    // Not applicable
+    override val docComment: String? get() = null
+
+    override fun asMemberOf(other: XType): KspType {
+        check(other is KspType)
+        val asMemberReceiverType = receiverType.resolve().let {
+            if (it.isError) {
+                return@let it
+            }
+            val asMember = enclosingMethodElement.declaration.asMemberOf(other.ksType)
+            checkNotNull(asMember.extensionReceiverType)
+        }
+        return env.wrap(
+            originatingReference = receiverType,
+            ksType = asMemberReceiverType,
+        ).withJvmTypeResolver(jvmTypeResolutionScope)
+    }
+
+    override fun kindName(): String {
+        return "synthetic receiver parameter"
+    }
+
+    override fun validate(): Boolean {
+        return true
+    }
+
+    override fun equals(other: Any?): Boolean {
+        return XEquality.equals(this, other)
+    }
+
+    override fun hashCode(): Int {
+        return XEquality.hashCode(equalityItems)
+    }
+}
diff --git a/room/room-compiler-processing/src/test/java/androidx/room/compiler/processing/XExecutableElementTest.kt b/room/room-compiler-processing/src/test/java/androidx/room/compiler/processing/XExecutableElementTest.kt
index 6c599c4..f069c60 100644
--- a/room/room-compiler-processing/src/test/java/androidx/room/compiler/processing/XExecutableElementTest.kt
+++ b/room/room-compiler-processing/src/test/java/androidx/room/compiler/processing/XExecutableElementTest.kt
@@ -31,6 +31,7 @@
 import com.squareup.javapoet.ClassName
 import com.squareup.javapoet.ParameterizedTypeName
 import com.squareup.javapoet.TypeName
+import com.squareup.javapoet.TypeVariableName
 import com.squareup.javapoet.WildcardTypeName
 import org.junit.Test
 import org.junit.runner.RunWith
@@ -963,6 +964,106 @@
         }
     }
 
+    @Test
+    fun extensionFun() {
+        fun buildSource(pkg: String) = Source.kotlin(
+            "Foo.kt",
+            """
+            package $pkg
+            abstract class Foo<T> {
+                fun String.ext1(): String = TODO()
+                fun String.ext2(inputParam: Int): String = TODO()
+                fun Foo<String>.ext3(): String = TODO()
+                fun Foo<T>.ext4(): String = TODO()
+                fun T.ext5(): String = TODO()
+                suspend fun String.ext6(): String = TODO()
+                abstract fun T.ext7(): String
+            }
+            class FooImpl : Foo<Int>() {
+                override fun Int.ext7(): String = TODO()
+            }
+            """.trimIndent()
+        )
+        runProcessorTest(
+            sources = listOf(buildSource(pkg = "app")),
+            classpath = compileFiles(listOf(buildSource(pkg = "lib")))
+        ) {
+            listOf("app", "lib").forEach { pkg ->
+                val element = it.processingEnv.requireTypeElement("$pkg.Foo")
+                element.getDeclaredMethodByJvmName("ext1").let { method ->
+                    assertThat(method.isExtensionFunction()).isTrue()
+                    assertThat(method.parameters.size).isEqualTo(1)
+                    assertThat(method.parameters[0].name).isEqualTo("\$this\$ext1")
+                    assertThat(method.parameters[0].type.typeName)
+                        .isEqualTo(String::class.typeName())
+                }
+                element.getDeclaredMethodByJvmName("ext2").let { method ->
+                    assertThat(method.parameters.size).isEqualTo(2)
+                    assertThat(method.parameters[0].name).isEqualTo("\$this\$ext2")
+                    assertThat(method.parameters[0].type.typeName)
+                        .isEqualTo(String::class.typeName())
+                    assertThat(method.parameters[1].name).isEqualTo("inputParam")
+                }
+                element.getDeclaredMethodByJvmName("ext3").let { method ->
+                    assertThat(method.parameters[0].type.typeName).isEqualTo(
+                        ParameterizedTypeName.get(
+                            ClassName.get(pkg, "Foo"),
+                            String::class.typeName()
+                        )
+                    )
+                }
+                element.getDeclaredMethodByJvmName("ext4").let { method ->
+                    assertThat(method.parameters[0].type.typeName).isEqualTo(
+                        ParameterizedTypeName.get(
+                            ClassName.get(pkg, "Foo"),
+                            TypeVariableName.get("T")
+                        )
+                    )
+                }
+                element.getDeclaredMethodByJvmName("ext5").let { method ->
+                    assertThat(method.parameters[0].type.typeName)
+                        .isEqualTo(TypeVariableName.get("T"))
+                }
+                element.getDeclaredMethodByJvmName("ext6").let { method ->
+                    assertThat(method.isSuspendFunction()).isTrue()
+                    assertThat(method.isExtensionFunction()).isTrue()
+                    assertThat(method.parameters.size).isEqualTo(2)
+                    assertThat(method.parameters[0].type.typeName)
+                        .isEqualTo(String::class.typeName())
+                    assertThat(method.parameters[1].type.typeName).isEqualTo(
+                        ParameterizedTypeName.get(
+                            ClassName.get("kotlin.coroutines", "Continuation"),
+                            WildcardTypeName.supertypeOf(String::class.typeName())
+                        )
+                    )
+                }
+                // Verify overridden Foo.ext7() asMemberOf FooImpl
+                element.getDeclaredMethodByJvmName("ext7").let { method ->
+                    assertThat(method.isAbstract()).isTrue()
+                    assertThat(method.isExtensionFunction()).isTrue()
+                    assertThat(method.parameters[0].type.typeName)
+                        .isEqualTo(TypeVariableName.get("T"))
+
+                    val fooImpl = it.processingEnv.requireTypeElement("$pkg.FooImpl")
+                    assertThat(method.parameters[0].asMemberOf(fooImpl.type).typeName)
+                        .isEqualTo(TypeName.INT.box())
+                }
+                // Verify non-overridden Foo.ext1() asMemberOf FooImpl
+                element.getDeclaredMethodByJvmName("ext1").let { method ->
+                    val fooImpl = it.processingEnv.requireTypeElement("$pkg.FooImpl")
+                    assertThat(method.parameters[0].asMemberOf(fooImpl.type).typeName)
+                        .isEqualTo(String::class.typeName())
+                }
+                // Verify non-overridden Foo.ext5() asMemberOf FooImpl
+                element.getDeclaredMethodByJvmName("ext5").let { method ->
+                    val fooImpl = it.processingEnv.requireTypeElement("$pkg.FooImpl")
+                    assertThat(method.parameters[0].asMemberOf(fooImpl.type).typeName)
+                        .isEqualTo(TypeName.INT.box())
+                }
+            }
+        }
+    }
+
     // see b/160258066
     private fun genericToPrimitiveOverrides(asMemberOf: Boolean) {
         val source = Source.kotlin(