WIP

plugins/compose-compiler-hosted/build.gradle.android

@@ -0,0 +1,54 @@
+/*
+ * Copyright 2019 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.
+ */
+
+
+import org.jetbrains.kotlin.gradle.tasks.KotlinCompile
+
+import static androidx.build.dependencies.DependenciesKt.*
+import androidx.build.CompilationTarget
+import androidx.build.LibraryGroups
+import androidx.build.LibraryVersions
+import androidx.build.Publish
+
+plugins {
+    id("AndroidXPlugin")
+    id("kotlin")
+}
+
+dependencies {
+    compileOnly(KOTLIN_STDLIB)
+    compileOnly("org.jetbrains.kotlin:kotlin-compiler:$KOTLIN_VERSION")
+    compileOnly("org.jetbrains.kotlin:kotlin-plugin:$KOTLIN_VERSION")
+    compileOnly("org.jetbrains.kotlin:kotlin-intellij-core:$KOTLIN_VERSION")
+    compileOnly("org.jetbrains.kotlin:kotlin-platform-api:$KOTLIN_VERSION")
+}
+
+tasks.withType(KotlinCompile).configureEach {
+    kotlinOptions {
+        jvmTarget = "1.8"
+    }
+}
+
+androidx {
+    name = "AndroidX Compose Hosted Compiler Plugin"
+    // Nobody should ever get this artifact from maven; just from studio or from source
+    publish = Publish.NONE
+    toolingProject = true
+    mavenGroup = LibraryGroups.COMPOSE
+    inceptionYear = "2019"
+    description = "Contains the Kotlin compiler plugin for Compose used in Android Studio and IDEA"
+    compilationTarget = CompilationTarget.HOST
+}

plugins/compose-compiler-hosted/build.gradle.kts

@@ -0,0 +1,47 @@
+description = "Kotlin Serialization Compiler Plugin"
+
+plugins {
+    kotlin("jvm")
+    id("jps-compatible")
+}
+
+dependencies {
+    compileOnly(intellijCoreDep()) { includeJars("intellij-core", "asm-all", rootProject = rootProject) }
+
+    compileOnly(project(":compiler:plugin-api"))
+    compileOnly(project(":compiler:cli-common"))
+    compileOnly(project(":compiler:frontend"))
+    compileOnly(project(":compiler:backend"))
+    compileOnly(project(":compiler:ir.backend.common"))
+    compileOnly(project(":compiler:backend.jvm"))
+    compileOnly(project(":js:js.frontend"))
+    compileOnly(project(":js:js.translator"))
+
+    runtime(kotlinStdlib())
+
+    testCompile(projectTests(":compiler:tests-common"))
+    testCompile(commonDep("junit:junit"))
+    testCompile("org.jetbrains.kotlinx:kotlinx-serialization-runtime:0.11.0")
+
+    testRuntimeOnly(intellijCoreDep()) { includeJars("intellij-core") }
+
+    Platform[192].orHigher {
+        testRuntimeOnly(intellijDep()) { includeJars("platform-concurrency") }
+    }
+}
+
+sourceSets {
+    "main" { projectDefault() }
+//    "test" { projectDefault() }
+}
+
+runtimeJar()
+sourcesJar()
+javadocJar()
+testsJar()
+
+projectTest(parallel = true) {
+    workingDir = rootDir
+}
+
+apply(from = "$rootDir/gradle/kotlinPluginPublication.gradle.kts")

plugins/compose-compiler-hosted/integration-tests/build.gradle

@@ -0,0 +1,85 @@
+/*
+ * Copyright 2019 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.
+ */
+
+
+import org.jetbrains.kotlin.gradle.tasks.KotlinCompile
+
+import static androidx.build.dependencies.DependenciesKt.*
+import androidx.build.LibraryGroups
+import androidx.build.LibraryVersions
+import androidx.build.Publish
+
+plugins {
+    id("AndroidXPlugin")
+    id("com.android.library")
+    id("kotlin-android")
+}
+
+dependencies {
+    implementation(KOTLIN_STDLIB)
+    testImplementation(JUNIT)
+    testImplementation(ROBOLECTRIC)
+
+    testImplementation("org.jetbrains.kotlin:kotlin-compiler:$KOTLIN_VERSION")
+
+    testImplementation(KOTLIN_STDLIB)
+    testImplementation(project(":compose:compose-runtime"))
+    testImplementation(project(":ui:ui-framework"))
+    testImplementation(project(":ui:ui-android-view"))
+    testImplementation(project(":compose:compose-compiler-hosted"))
+}
+
+android {
+    defaultConfig {
+        minSdkVersion 16
+        javaCompileOptions {
+            annotationProcessorOptions {
+                // The kotlin-compiler contains an annotation processor that we don't need to invoke
+                // The following ensure a warning isn't reported by the android plugin.
+                includeCompileClasspath false
+            }
+        }
+    }
+    lintOptions {
+        disable("SyntheticAccessor")
+    }
+    testOptions {
+        unitTests.all {
+            // There is only one version of the compose plugin built so the debug tests are
+            // sufficient as they test that one version
+            if (it.name == 'testReleaseUnitTest') {
+                filter {
+                    exclude '*'
+                }
+            }
+        }
+    }
+}
+
+tasks.withType(KotlinCompile).configureEach {
+    kotlinOptions {
+        jvmTarget = "1.8"
+    }
+}
+
+androidx {
+    name = "AndroidX Compiler CLI Tests"
+    publish = Publish.NONE
+    toolingProject = true
+    mavenGroup = LibraryGroups.COMPOSE
+    inceptionYear = "2019"
+    description = "Contains test for the compose compiler plugin"
+}

plugins/compose-compiler-hosted/integration-tests/src/main/AndroidManifest.xml

@@ -0,0 +1,23 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!--
+  Copyright 2019 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.
+  -->
+<manifest package="androidx.compose.runtime" xmlns:android="http://schemas.android.com/apk/res/android">
+
+    <application/>
+
+
+</manifest>
+

plugins/compose-compiler-hosted/integration-tests/src/test/java/androidx/compose/plugins/kotlin/AbstractCodegenSignatureTest.kt

@@ -0,0 +1,229 @@
+/*
+ * Copyright 2019 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.compose.plugins.kotlin
+
+import android.app.Activity
+import android.content.Context
+import android.os.Bundle
+import android.widget.LinearLayout
+import org.jetbrains.kotlin.backend.common.output.OutputFile
+import org.robolectric.Robolectric
+import java.net.URLClassLoader
+
+fun printPublicApi(classDump: String, name: String): String {
+    return classDump
+        .splitToSequence("\n")
+        .filter {
+            if (it.contains("INVOKESTATIC kotlin/internal/ir/Intrinsic")) {
+                // if instructions like this end up in our generated code, it means something
+                // went wrong. Usually it means that it just can't find the function to call,
+                // so it transforms it into this intrinsic call instead of failing. If this
+                // happens, we want to hard-fail the test as the code is definitely incorrect.
+                error(
+                    buildString {
+                        append("An unresolved call was found in the generated bytecode of '")
+                        append(name)
+                        append("'")
+                        appendln()
+                        appendln()
+                        appendln("Call was: $it")
+                        appendln()
+                        appendln("Entire class file output:")
+                        appendln(classDump)
+                    }
+                )
+            }
+            if (it.startsWith("  ")) {
+                if (it.startsWith("   ")) false
+                else it[2] != '/' && it[2] != '@'
+            } else {
+                it == "}" || it.endsWith("{")
+            }
+        }
+        .joinToString(separator = "\n")
+        .replace('$', '%') // replace $ to % to make comparing it to kotlin string literals easier
+}
+
+abstract class AbstractCodegenSignatureTest : AbstractCodegenTest() {
+
+    private var isSetup = false
+    override fun setUp() {
+        isSetup = true
+        super.setUp()
+    }
+
+    private fun <T> ensureSetup(block: () -> T): T {
+        if (!isSetup) setUp()
+        return block()
+    }
+
+    private fun OutputFile.printApi(): String {
+        return printPublicApi(asText(), relativePath)
+    }
+
+    fun checkApi(src: String, expected: String, dumpClasses: Boolean = false): Unit = ensureSetup {
+        val className = "Test_REPLACEME_${uniqueNumber++}"
+        val fileName = "$className.kt"
+
+        val loader = classLoader("""
+           import androidx.compose.*
+
+           $src
+        """, fileName, dumpClasses)
+
+        val apiString = loader
+            .allGeneratedFiles
+            .filter { it.relativePath.endsWith(".class") }
+            .map { it.printApi() }
+            .joinToString(separator = "\n")
+            .replace(className, "Test")
+
+        val expectedApiString = expected
+            .trimIndent()
+            .split("\n")
+            .filter { it.isNotBlank() }
+            .joinToString("\n")
+
+        assertEquals(expectedApiString, apiString)
+    }
+
+    fun validateBytecode(
+        src: String,
+        dumpClasses: Boolean = false,
+        validate: (String) -> Unit
+    ): Unit = ensureSetup {
+        val className = "Test_REPLACEME_${uniqueNumber++}"
+        val fileName = "$className.kt"
+
+        val loader = classLoader("""
+           import androidx.compose.*
+
+           $src
+        """, fileName, dumpClasses)
+
+        val apiString = loader
+            .allGeneratedFiles
+            .filter { it.relativePath.endsWith(".class") }
+            .map {
+                it.asText().replace('$', '%').replace(className, "Test")
+            }.joinToString("\n")
+
+        validate(apiString)
+    }
+
+    fun checkComposerParam(src: String, dumpClasses: Boolean = false): Unit = ensureSetup {
+        val className = "Test_REPLACEME_${uniqueNumber++}"
+        val compiledClasses = classLoader(
+            """
+                import androidx.compose.*
+                import android.widget.LinearLayout
+                import android.content.Context
+
+                $src
+
+                @Composable fun assertComposer(expected: Composer<*>?) {
+                    val actual = currentComposer
+                    assert(expected === actual)
+                }
+
+                private var __context: Context? = null
+
+                fun makeComposer(): Composer<*> {
+                    val container = LinearLayout(__context!!)
+                    return ViewComposer(
+                        __context!!,
+                        container,
+                        SlotTable(),
+                        Recomposer.current()
+                    )
+                }
+
+                fun invokeComposable(composer: Composer<*>?, fn: @Composable() () -> Unit) {
+                    if (composer == null) error("Composer was null")
+                    val realFn = fn as Function1<Composer<*>, Unit>
+                    composer.runWithComposing {
+                        composer.startRoot()
+                        realFn(composer)
+                        composer.endRoot()
+                    }
+                }
+
+                class Test {
+                  fun test(context: Context) {
+                    __context = context
+                    run()
+                    __context = null
+                  }
+                }
+            """,
+            fileName = className,
+            dumpClasses = dumpClasses
+        )
+
+        val allClassFiles = compiledClasses.allGeneratedFiles.filter {
+            it.relativePath.endsWith(".class")
+        }
+
+        val loader = URLClassLoader(emptyArray(), this.javaClass.classLoader)
+
+        val instanceClass = run {
+            var instanceClass: Class<*>? = null
+            var loadedOne = false
+            for (outFile in allClassFiles) {
+                val bytes = outFile.asByteArray()
+                val loadedClass = loadClass(loader, null, bytes)
+                if (loadedClass.name == "Test") instanceClass = loadedClass
+                loadedOne = true
+            }
+            if (!loadedOne) error("No classes loaded")
+            instanceClass ?: error("Could not find class $className in loaded classes")
+        }
+
+        val instanceOfClass = instanceClass.newInstance()
+        val testMethod = instanceClass.getMethod("test", Context::class.java)
+
+        val controller = Robolectric.buildActivity(TestActivity::class.java)
+        val activity = controller.create().get()
+        testMethod.invoke(instanceOfClass, activity)
+    }
+
+    private class TestActivity : Activity() {
+        override fun onCreate(savedInstanceState: Bundle?) {
+            super.onCreate(savedInstanceState)
+            setContentView(LinearLayout(this))
+        }
+    }
+
+    fun codegen(text: String, dumpClasses: Boolean = false): Unit = ensureSetup {
+        codegenNoImports(
+            """
+           import android.content.Context
+           import android.widget.*
+           import androidx.compose.*
+
+           $text
+
+        """, dumpClasses)
+    }
+
+    fun codegenNoImports(text: String, dumpClasses: Boolean = false): Unit = ensureSetup {
+        val className = "Test_${uniqueNumber++}"
+        val fileName = "$className.kt"
+
+        classLoader(text, fileName, dumpClasses)
+    }
+}

plugins/compose-compiler-hosted/integration-tests/src/test/java/androidx/compose/plugins/kotlin/AbstractCodegenTest.kt

@@ -0,0 +1,155 @@
+/*
+ * Copyright 2019 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.compose.plugins.kotlin
+
+import com.intellij.openapi.project.Project
+import com.intellij.openapi.util.text.StringUtilRt
+import com.intellij.openapi.vfs.CharsetToolkit
+import com.intellij.psi.PsiFileFactory
+import com.intellij.psi.impl.PsiFileFactoryImpl
+import com.intellij.testFramework.LightVirtualFile
+import org.jetbrains.kotlin.cli.jvm.compiler.EnvironmentConfigFiles
+import org.jetbrains.kotlin.cli.jvm.compiler.KotlinCoreEnvironment
+import org.jetbrains.kotlin.cli.jvm.config.addJvmClasspathRoots
+import org.jetbrains.kotlin.codegen.GeneratedClassLoader
+import org.jetbrains.kotlin.config.CompilerConfiguration
+import org.jetbrains.kotlin.config.JVMConfigurationKeys
+import org.jetbrains.kotlin.config.JvmTarget
+import org.jetbrains.kotlin.idea.KotlinLanguage
+import org.jetbrains.kotlin.psi.KtFile
+import org.jetbrains.kotlin.resolve.AnalyzingUtils
+import java.io.File
+
+abstract class AbstractCodegenTest : AbstractCompilerTest() {
+    override fun setUp() {
+        super.setUp()
+        val classPath = createClasspath() + additionalPaths
+
+        val configuration = newConfiguration()
+        configuration.addJvmClasspathRoots(classPath)
+        updateConfiguration(configuration)
+
+        myEnvironment = KotlinCoreEnvironment.createForTests(
+            myTestRootDisposable, configuration, EnvironmentConfigFiles.JVM_CONFIG_FILES
+        ).also { setupEnvironment(it) }
+    }
+
+    open fun updateConfiguration(configuration: CompilerConfiguration) {
+        configuration.put(JVMConfigurationKeys.IR, true)
+        configuration.put(JVMConfigurationKeys.JVM_TARGET, JvmTarget.JVM_1_8)
+    }
+
+    protected open fun helperFiles(): List<KtFile> = emptyList()
+
+    protected fun dumpClasses(loader: GeneratedClassLoader) {
+        for (file in loader.allGeneratedFiles.filter {
+            it.relativePath.endsWith(".class")
+        }) {
+            println("------\nFILE: ${file.relativePath}\n------")
+            println(file.asText())
+        }
+    }
+
+    protected fun classLoader(
+        source: String,
+        fileName: String,
+        dumpClasses: Boolean = false
+    ): GeneratedClassLoader {
+        val files = mutableListOf<KtFile>()
+        files.addAll(helperFiles())
+        files.add(sourceFile(fileName, source))
+        myFiles = CodegenTestFiles.create(files)
+        val loader = createClassLoader()
+        if (dumpClasses) dumpClasses(loader)
+        return loader
+    }
+
+    protected fun classLoader(
+        sources: Map<String, String>,
+        dumpClasses: Boolean = false
+    ): GeneratedClassLoader {
+        val files = mutableListOf<KtFile>()
+        files.addAll(helperFiles())
+        for ((fileName, source) in sources) {
+            files.add(sourceFile(fileName, source))
+        }
+        myFiles = CodegenTestFiles.create(files)
+        val loader = createClassLoader()
+        if (dumpClasses) dumpClasses(loader)
+        return loader
+    }
+
+    protected fun testFile(source: String, dumpClasses: Boolean = false) {
+        val files = mutableListOf<KtFile>()
+        files.addAll(helperFiles())
+        files.add(sourceFile("Test.kt", source))
+        myFiles = CodegenTestFiles.create(files)
+        val loader = createClassLoader()
+        if (dumpClasses) dumpClasses(loader)
+        val loadedClass = loader.loadClass("Test")
+        val instance = loadedClass.newInstance()
+        val instanceClass = instance::class.java
+        val testMethod = instanceClass.getMethod("test")
+        testMethod.invoke(instance)
+    }
+
+    protected fun testCompile(source: String, dumpClasses: Boolean = false) {
+        val files = mutableListOf<KtFile>()
+        files.addAll(helperFiles())
+        files.add(sourceFile("Test.kt", source))
+        myFiles = CodegenTestFiles.create(files)
+        val loader = createClassLoader()
+        if (dumpClasses) dumpClasses(loader)
+    }
+
+    protected fun sourceFile(name: String, source: String): KtFile {
+        val result =
+            createFile(name, source, myEnvironment!!.project)
+        val ranges = AnalyzingUtils.getSyntaxErrorRanges(result)
+        assert(ranges.isEmpty()) { "Syntax errors found in $name: $ranges" }
+        return result
+    }
+
+    protected fun loadClass(className: String, source: String): Class<*> {
+        myFiles = CodegenTestFiles.create(
+            "file.kt",
+            source,
+            myEnvironment!!.project
+        )
+        val loader = createClassLoader()
+        return loader.loadClass(className)
+    }
+
+    protected open val additionalPaths = emptyList<File>()
+}
+
+fun createFile(name: String, text: String, project: Project): KtFile {
+    var shortName = name.substring(name.lastIndexOf('/') + 1)
+    shortName = shortName.substring(shortName.lastIndexOf('\\') + 1)
+    val virtualFile = object : LightVirtualFile(
+        shortName,
+        KotlinLanguage.INSTANCE,
+        StringUtilRt.convertLineSeparators(text)
+    ) {
+        override fun getPath(): String = "/$name"
+    }
+
+    virtualFile.setCharset(CharsetToolkit.UTF8_CHARSET)
+    val factory = PsiFileFactory.getInstance(project) as PsiFileFactoryImpl
+
+    return factory.trySetupPsiForFile(virtualFile, KotlinLanguage.INSTANCE, true, false) as KtFile
+}

plugins/compose-compiler-hosted/integration-tests/src/test/java/androidx/compose/plugins/kotlin/AbstractCompilerTest.kt

@@ -0,0 +1,300 @@
+/*
+ * Copyright 2019 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.compose.plugins.kotlin
+
+import com.intellij.mock.MockProject
+import com.intellij.openapi.Disposable
+import com.intellij.openapi.util.Disposer
+import com.intellij.openapi.util.io.FileUtil
+import com.intellij.openapi.util.text.StringUtil
+import junit.framework.TestCase
+import org.jetbrains.annotations.Contract
+import org.jetbrains.kotlin.cli.common.CLIConfigurationKeys
+import org.jetbrains.kotlin.cli.common.messages.CompilerMessageLocation
+import org.jetbrains.kotlin.cli.common.messages.CompilerMessageSeverity
+import org.jetbrains.kotlin.cli.common.messages.MessageCollector
+import org.jetbrains.kotlin.cli.jvm.compiler.EnvironmentConfigFiles
+import org.jetbrains.kotlin.cli.jvm.compiler.KotlinCoreEnvironment
+import org.jetbrains.kotlin.cli.jvm.compiler.NoScopeRecordCliBindingTrace
+import org.jetbrains.kotlin.cli.jvm.config.addJvmClasspathRoots
+import org.jetbrains.kotlin.codegen.ClassBuilderFactories
+import org.jetbrains.kotlin.codegen.ClassFileFactory
+import org.jetbrains.kotlin.codegen.GeneratedClassLoader
+import org.jetbrains.kotlin.config.CommonConfigurationKeys
+import org.jetbrains.kotlin.config.CompilerConfiguration
+import org.jetbrains.kotlin.utils.rethrow
+import org.junit.After
+import java.io.File
+import java.net.MalformedURLException
+import java.net.URL
+import java.net.URLClassLoader
+
+private const val KOTLIN_RUNTIME_VERSION = "1.3.11"
+
+@Suppress("MemberVisibilityCanBePrivate")
+abstract class AbstractCompilerTest : TestCase() {
+    protected var myEnvironment: KotlinCoreEnvironment? = null
+    protected var myFiles: CodegenTestFiles? = null
+    protected var classFileFactory: ClassFileFactory? = null
+    protected var javaClassesOutputDirectory: File? = null
+    protected var additionalDependencies: List<File>? = null
+
+    override fun setUp() {
+        // Setup the environment for the analysis
+        System.setProperty("user.dir",
+            homeDir
+        )
+        myEnvironment = createEnvironment()
+        setupEnvironment(myEnvironment!!)
+        super.setUp()
+    }
+
+    override fun tearDown() {
+        myFiles = null
+        myEnvironment = null
+        javaClassesOutputDirectory = null
+        additionalDependencies = null
+        classFileFactory = null
+        Disposer.dispose(myTestRootDisposable)
+        super.tearDown()
+    }
+
+    fun ensureSetup(block: () -> Unit) {
+        setUp()
+        block()
+    }
+
+    @After
+    fun after() {
+        tearDown()
+    }
+
+    protected val defaultClassPath by lazy { systemClassLoaderJars() }
+
+    protected fun createClasspath() = defaultClassPath.filter {
+        !it.path.contains("robolectric")
+    }.toList()
+
+    val myTestRootDisposable = TestDisposable()
+
+    protected fun createEnvironment(): KotlinCoreEnvironment {
+        val classPath = createClasspath()
+
+        val configuration = newConfiguration()
+        configuration.addJvmClasspathRoots(classPath)
+
+        return KotlinCoreEnvironment.createForTests(
+            myTestRootDisposable,
+            configuration,
+            EnvironmentConfigFiles.JVM_CONFIG_FILES
+        )
+    }
+
+    protected open fun setupEnvironment(environment: KotlinCoreEnvironment) {
+        ComposeComponentRegistrar.registerProjectExtensions(
+            environment.project as MockProject,
+            environment.configuration
+        )
+    }
+
+    protected fun createClassLoader(): GeneratedClassLoader {
+        val classLoader = URLClassLoader(defaultClassPath.map {
+            it.toURI().toURL()
+        }.toTypedArray(), null)
+        return GeneratedClassLoader(
+            generateClassesInFile(),
+            classLoader,
+            *getClassPathURLs()
+        )
+    }
+
+    protected fun getClassPathURLs(): Array<URL> {
+        val files = mutableListOf<File>()
+        javaClassesOutputDirectory?.let { files.add(it) }
+        additionalDependencies?.let { files.addAll(it) }
+
+        try {
+            return files.map { it.toURI().toURL() }.toTypedArray()
+        } catch (e: MalformedURLException) {
+            throw rethrow(e)
+        }
+    }
+
+    private fun reportProblem(e: Throwable) {
+        e.printStackTrace()
+        System.err.println("Generating instructions as text...")
+        try {
+            System.err.println(classFileFactory?.createText()
+                ?: "Cannot generate text: exception was thrown during generation")
+        } catch (e1: Throwable) {
+            System.err.println(
+                "Exception thrown while trying to generate text, " +
+                        "the actual exception follows:"
+            )
+            e1.printStackTrace()
+            System.err.println(
+                "------------------------------------------------------------------" +
+                        "-----------"
+            )
+        }
+
+        System.err.println("See exceptions above")
+    }
+
+    protected fun generateClassesInFile(reportProblems: Boolean = true): ClassFileFactory {
+        return classFileFactory ?: run {
+            try {
+                val environment = myEnvironment ?: error("Environment not initialized")
+                val files = myFiles ?: error("Files not initialized")
+                val generationState = GenerationUtils.compileFiles(
+                    files.psiFiles, environment, ClassBuilderFactories.TEST,
+                    NoScopeRecordCliBindingTrace()
+                )
+                generationState.factory.also { classFileFactory = it }
+            } catch (e: TestsCompilerError) {
+                if (reportProblems) {
+                    reportProblem(e.original)
+                } else {
+                    System.err.println("Compilation failure")
+                }
+                throw e
+            } catch (e: Throwable) {
+                if (reportProblems) reportProblem(e)
+                throw TestsCompilerError(e)
+            }
+        }
+    }
+
+    protected fun getTestName(lowercaseFirstLetter: Boolean): String =
+        getTestName(this.name ?: "", lowercaseFirstLetter)
+    protected fun getTestName(name: String, lowercaseFirstLetter: Boolean): String {
+        val trimmedName = trimStart(name, "test")
+        return if (StringUtil.isEmpty(trimmedName)) "" else lowercaseFirstLetter(
+            trimmedName,
+            lowercaseFirstLetter
+        )
+    }
+
+    protected fun lowercaseFirstLetter(name: String, lowercaseFirstLetter: Boolean): String =
+        if (lowercaseFirstLetter && !isMostlyUppercase(name))
+            Character.toLowerCase(name[0]) + name.substring(1)
+        else name
+
+    protected fun isMostlyUppercase(name: String): Boolean {
+        var uppercaseChars = 0
+        for (i in 0 until name.length) {
+            if (Character.isLowerCase(name[i])) {
+                return false
+            }
+            if (Character.isUpperCase(name[i])) {
+                uppercaseChars++
+                if (uppercaseChars >= 3) return true
+            }
+        }
+        return false
+    }
+
+    inner class TestDisposable : Disposable {
+
+        override fun dispose() {}
+
+        override fun toString(): String {
+            val testName = this@AbstractCompilerTest.getTestName(false)
+            return this@AbstractCompilerTest.javaClass.name +
+                    if (StringUtil.isEmpty(testName)) "" else ".test$testName"
+        }
+    }
+
+    companion object {
+        val homeDir by lazy { File(computeHomeDirectory()).absolutePath }
+        val projectRoot by lazy { File(homeDir, "../../../../..").absolutePath }
+        val kotlinHome by lazy {
+            File(projectRoot, "prebuilts/androidx/external/org/jetbrains/kotlin/")
+        }
+        val outDir by lazy {
+            File(System.getenv("OUT_DIR") ?: File(projectRoot, "out").absolutePath)
+        }
+        val composePluginJar by lazy {
+
+            File(outDir, "ui/compose/compose-compiler/build/jarjar/compose-compiler.jar")
+        }
+
+        fun kotlinRuntimeJar(module: String) = File(
+            kotlinHome, "$module/$KOTLIN_RUNTIME_VERSION/$module-$KOTLIN_RUNTIME_VERSION.jar")
+
+        init {
+            System.setProperty("idea.home",
+                homeDir
+            )
+        }
+    }
+}
+
+private fun systemClassLoaderJars(): List<File> {
+    val result = (ClassLoader.getSystemClassLoader() as? URLClassLoader)?.urLs?.filter {
+        it.protocol == "file"
+    }?.map {
+        File(it.path)
+    }?.toList() ?: emptyList()
+    return result
+}
+
+private fun computeHomeDirectory(): String {
+    val userDir = System.getProperty("user.dir")
+    val dir = File(userDir ?: ".")
+    return FileUtil.toCanonicalPath(dir.absolutePath)
+}
+
+private const val TEST_MODULE_NAME = "test-module"
+
+fun newConfiguration(): CompilerConfiguration {
+    val configuration = CompilerConfiguration()
+    configuration.put(CommonConfigurationKeys.MODULE_NAME,
+        TEST_MODULE_NAME
+    )
+
+    configuration.put(CLIConfigurationKeys.MESSAGE_COLLECTOR_KEY, object : MessageCollector {
+        override fun clear() {}
+
+        override fun report(
+            severity: CompilerMessageSeverity,
+            message: String,
+            location: CompilerMessageLocation?
+        ) {
+            if (severity === CompilerMessageSeverity.ERROR) {
+                val prefix = if (location == null)
+                    ""
+                else
+                    "(" + location.path + ":" + location.line + ":" + location.column + ") "
+                throw AssertionError(prefix + message)
+            }
+        }
+
+        override fun hasErrors(): Boolean {
+            return false
+        }
+    })
+
+    return configuration
+}
+
+@Contract(pure = true)
+fun trimStart(s: String, prefix: String): String {
+    return if (s.startsWith(prefix)) {
+        s.substring(prefix.length)
+    } else s
+}

plugins/compose-compiler-hosted/integration-tests/src/test/java/androidx/compose/plugins/kotlin/AbstractComposeDiagnosticsTest.kt

@@ -0,0 +1,166 @@
+/*
+ * Copyright 2019 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.compose.plugins.kotlin
+
+import org.jetbrains.kotlin.checkers.utils.CheckerTestUtil
+import org.jetbrains.kotlin.checkers.DiagnosedRange
+import org.jetbrains.kotlin.cli.jvm.compiler.KotlinCoreEnvironment
+import org.jetbrains.kotlin.cli.jvm.compiler.TopDownAnalyzerFacadeForJVM
+import org.jetbrains.kotlin.cli.jvm.compiler.NoScopeRecordCliBindingTrace
+import org.jetbrains.kotlin.config.JVMConfigurationKeys
+import org.jetbrains.kotlin.config.JvmTarget
+import org.jetbrains.kotlin.diagnostics.Diagnostic
+import org.jetbrains.kotlin.diagnostics.DiagnosticWithParameters1
+import org.jetbrains.kotlin.diagnostics.RenderedDiagnostic
+import java.io.File
+
+abstract class AbstractComposeDiagnosticsTest : AbstractCompilerTest() {
+
+    fun doTest(expectedText: String) {
+        doTest(expectedText, myEnvironment!!)
+    }
+
+    fun doTest(expectedText: String, environment: KotlinCoreEnvironment) {
+        val diagnosedRanges: MutableList<DiagnosedRange> = ArrayList()
+        val clearText = CheckerTestUtil.parseDiagnosedRanges(expectedText, diagnosedRanges)
+        val file =
+            createFile("test.kt", clearText, environment.project)
+        val files = listOf(file)
+
+        // Use the JVM version of the analyzer to allow using classes in .jar files
+        val moduleTrace = NoScopeRecordCliBindingTrace()
+        val result = TopDownAnalyzerFacadeForJVM.analyzeFilesWithJavaIntegration(
+            environment.project,
+            files,
+            moduleTrace,
+            environment.configuration.copy().apply {
+                this.put(JVMConfigurationKeys.JVM_TARGET, JvmTarget.JVM_1_8)
+            },
+            environment::createPackagePartProvider
+        )
+
+        // Collect the errors
+        val errors = result.bindingContext.diagnostics.all().toMutableList()
+
+        val message = StringBuilder()
+
+        // Ensure all the expected messages are there
+        val found = mutableSetOf<Diagnostic>()
+        for (range in diagnosedRanges) {
+            for (diagnostic in range.getDiagnostics()) {
+                val reportedDiagnostics = errors.filter { it.factoryName == diagnostic.name }
+                if (reportedDiagnostics.isNotEmpty()) {
+                    val reportedDiagnostic =
+                        reportedDiagnostics.find {
+                            it.textRanges.find {
+                                it.startOffset == range.start && it.endOffset == range.end
+                            } != null
+                        }
+                    if (reportedDiagnostic == null) {
+                        val firstRange = reportedDiagnostics.first().textRanges.first()
+                        message.append("  Error ${diagnostic.name} reported at ${
+                            firstRange.startOffset
+                        }-${firstRange.endOffset} but expected at ${range.start}-${range.end}\n")
+                        message.append(
+                            sourceInfo(
+                                clearText,
+                                firstRange.startOffset, firstRange.endOffset,
+                                "  "
+                            )
+                        )
+                    } else {
+                        errors.remove(reportedDiagnostic)
+                        found.add(reportedDiagnostic)
+                    }
+                } else {
+                    message.append("  Diagnostic ${diagnostic.name} not reported, expected at ${
+                        range.start
+                    }\n")
+                    message.append(
+                        sourceInfo(
+                            clearText,
+                            range.start,
+                            range.end,
+                            "  "
+                        )
+                    )
+                }
+            }
+        }
+
+        // Ensure only the expected errors are reported
+        for (diagnostic in errors) {
+            if (diagnostic !in found) {
+                val range = diagnostic.textRanges.first()
+                message.append(
+                    "  Unexpected diagnostic ${diagnostic.factoryName} reported at ${
+                        range.startOffset
+                    }\n"
+                )
+                message.append(
+                    sourceInfo(
+                        clearText,
+                        range.startOffset,
+                        range.endOffset,
+                        "  "
+                    )
+                )
+            }
+        }
+
+        // Throw an error if anything was found that was not expected
+        if (message.length > 0) throw Exception("Mismatched errors:\n$message")
+    }
+}
+
+fun assertExists(file: File): File {
+    if (!file.exists()) {
+        throw IllegalStateException("'$file' does not exist. Run test from gradle")
+    }
+    return file
+}
+
+// Normalize the factory's name to find the name supplied by a plugin
+@Suppress("UNCHECKED_CAST")
+val Diagnostic.factoryName: String
+    inline get() {
+        if (factory.name == "PLUGIN_ERROR")
+            return (this as
+                    DiagnosticWithParameters1<*, RenderedDiagnostic<*>>).a.diagnostic.factory.name
+        if (factory.name == "PLUGIN_WARNING")
+            return (this as
+                    DiagnosticWithParameters1<*, RenderedDiagnostic<*>>).a.diagnostic.factory.name
+        return factory.name
+    }
+
+fun String.lineStart(offset: Int): Int {
+    return this.lastIndexOf('\n', offset) + 1
+}
+
+fun String.lineEnd(offset: Int): Int {
+    val result = this.indexOf('\n', offset)
+    return if (result < 0) this.length else result
+}
+
+// Return the source line that contains the given range with the range underlined with '~'s
+fun sourceInfo(clearText: String, start: Int, end: Int, prefix: String = ""): String {
+    val lineStart = clearText.lineStart(start)
+    val lineEnd = clearText.lineEnd(start)
+    val displayEnd = if (end > lineEnd) lineEnd else end
+    return prefix + clearText.substring(lineStart, lineEnd) + "\n" +
+            prefix + " ".repeat(start - lineStart) + "~".repeat(displayEnd - start) + "\n"
+}

plugins/compose-compiler-hosted/integration-tests/src/test/java/androidx/compose/plugins/kotlin/AbstractLoweringTests.kt

@@ -0,0 +1,128 @@
+/*
+ * Copyright 2020 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.compose.plugins.kotlin
+
+import androidx.compose.Composer
+import org.jetbrains.kotlin.resolve.calls.model.ResolvedCall
+import java.net.URLClassLoader
+
+abstract class AbstractLoweringTests : AbstractCodegenTest() {
+
+    fun codegen(text: String, dumpClasses: Boolean = false) {
+        codegenNoImports(
+            """
+           import android.content.Context
+           import android.widget.*
+           import androidx.compose.*
+
+           $text
+
+        """, dumpClasses)
+    }
+
+    fun codegenNoImports(text: String, dumpClasses: Boolean = false) {
+        val className = "Test_${uniqueNumber++}"
+        val fileName = "$className.kt"
+
+        classLoader(text, fileName, dumpClasses)
+    }
+
+    fun compose(
+        supportingCode: String,
+        composeCode: String,
+        valuesFactory: () -> Map<String, Any> = { emptyMap() },
+        dumpClasses: Boolean = false
+    ): RobolectricComposeTester {
+        val className = "TestFCS_${uniqueNumber++}"
+        val fileName = "$className.kt"
+
+        val candidateValues = valuesFactory()
+
+        @Suppress("NO_REFLECTION_IN_CLASS_PATH")
+        val parameterList = candidateValues.map {
+            if (it.key.contains(':')) {
+                it.key
+            } else "${it.key}: ${it.value::class.qualifiedName}"
+        }.joinToString()
+        val parameterTypes = candidateValues.map {
+            it.value::class.javaPrimitiveType ?: it.value::class.javaObjectType
+        }.toTypedArray()
+
+        val compiledClasses = classLoader(
+            """
+       import android.content.Context
+       import android.widget.*
+       import androidx.compose.*
+       import androidx.ui.androidview.adapters.*
+
+       $supportingCode
+
+       class $className {
+
+         @Composable
+         fun test($parameterList) {
+           $composeCode
+         }
+       }
+    """, fileName, dumpClasses
+        )
+
+        val allClassFiles = compiledClasses.allGeneratedFiles.filter {
+            it.relativePath.endsWith(".class")
+        }
+
+        val loader = URLClassLoader(emptyArray(), this.javaClass.classLoader)
+
+        val instanceClass = run {
+            var instanceClass: Class<*>? = null
+            var loadedOne = false
+            for (outFile in allClassFiles) {
+                val bytes = outFile.asByteArray()
+                val loadedClass = loadClass(loader, null, bytes)
+                if (loadedClass.name == className) instanceClass = loadedClass
+                loadedOne = true
+            }
+            if (!loadedOne) error("No classes loaded")
+            instanceClass ?: error("Could not find class $className in loaded classes")
+        }
+
+        val instanceOfClass = instanceClass.newInstance()
+        val testMethod = instanceClass.getMethod("test", *parameterTypes, Composer::class.java)
+
+        return compose {
+            val values = valuesFactory()
+            val arguments = values.map { it.value as Any }.toTypedArray()
+            testMethod.invoke(instanceOfClass, *arguments, it)
+        }
+    }
+
+    private fun ResolvedCall<*>.isEmit(): Boolean = candidateDescriptor is ComposableEmitDescriptor
+    private fun ResolvedCall<*>.isCall(): Boolean =
+        candidateDescriptor is ComposableFunctionDescriptor
+
+    private val callPattern = Regex("(<normal>)|(<emit>)|(<call>)")
+    private fun extractCarets(text: String): Pair<String, List<Pair<Int, String>>> {
+        val indices = mutableListOf<Pair<Int, String>>()
+        var offset = 0
+        val src = callPattern.replace(text) {
+            indices.add(it.range.first - offset to it.value)
+            offset += it.range.last - it.range.first + 1
+            ""
+        }
+        return src to indices
+    }
+}

plugins/compose-compiler-hosted/integration-tests/src/test/java/androidx/compose/plugins/kotlin/AbstractMultiPlatformIntegrationTest.kt

@@ -0,0 +1,166 @@
+/*
+ * Copyright 2020 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.compose.plugins.kotlin
+
+import com.intellij.openapi.util.io.FileUtil
+import org.jetbrains.kotlin.cli.common.CLICompiler
+import org.jetbrains.kotlin.cli.common.CLITool
+import org.jetbrains.kotlin.cli.common.ExitCode
+import org.jetbrains.kotlin.cli.jvm.K2JVMCompiler
+import org.jetbrains.org.objectweb.asm.ClassReader
+import org.jetbrains.org.objectweb.asm.util.TraceClassVisitor
+import java.io.ByteArrayOutputStream
+import java.io.File
+import java.io.PrintStream
+import java.io.PrintWriter
+
+// KotlinTestUtils
+private fun tmpDir(name: String): File {
+    return FileUtil.createTempDirectory(name, "", false).canonicalFile
+}
+// AbstractCliTest
+private fun executeCompilerGrabOutput(
+    compiler: CLITool<*>,
+    args: List<String>
+): Pair<String, ExitCode> {
+    val output = StringBuilder()
+
+    var index = 0
+    do {
+        var next = args.subList(index, args.size).indexOf("---")
+        if (next == -1) {
+            next = args.size
+        }
+        val (first, second) = executeCompiler(compiler, args.subList(index, next))
+        output.append(first)
+        if (second != ExitCode.OK) {
+            return Pair(output.toString(), second)
+        }
+        index = next + 1
+    } while (index < args.size)
+
+    return Pair(output.toString(), ExitCode.OK)
+}
+// CompilerTestUtil
+private fun executeCompiler(compiler: CLITool<*>, args: List<String>): Pair<String, ExitCode> {
+    val bytes = ByteArrayOutputStream()
+    val origErr = System.err
+    try {
+        System.setErr(PrintStream(bytes))
+        val exitCode = CLITool.doMainNoExit(compiler, args.toTypedArray())
+        return Pair(String(bytes.toByteArray()), exitCode)
+    } finally {
+        System.setErr(origErr)
+    }
+}
+// jetTestUtils
+fun String.trimTrailingWhitespaces(): String =
+    this.split('\n').joinToString(separator = "\n") { it.trimEnd() }
+// jetTestUtils
+fun String.trimTrailingWhitespacesAndAddNewlineAtEOF(): String =
+    this.trimTrailingWhitespaces().let {
+            result -> if (result.endsWith("\n")) result else result + "\n"
+    }
+
+abstract class AbstractMultiPlatformIntegrationTest : AbstractCompilerTest() {
+    fun multiplatform(
+        common: String,
+        jvm: String,
+        output: String
+    ) {
+        setUp()
+        val tmpdir = tmpDir(getTestName(true))
+
+        assert(composePluginJar.exists())
+
+        val optionalArgs = arrayOf(
+            "-cp",
+            defaultClassPath
+                .filter { it.exists() }
+                .joinToString(File.pathSeparator) { it.absolutePath },
+            "-kotlin-home",
+            AbstractCompilerTest.kotlinHome.absolutePath,
+            "-P", "plugin:androidx.compose.plugins.idea:enabled=true",
+            "-Xplugin=${composePluginJar.absolutePath}",
+            "-Xuse-ir"
+        )
+
+        val jvmOnlyArgs = arrayOf("-no-stdlib")
+
+        val srcDir = File(tmpdir, "srcs").absolutePath
+        val commonSrc = File(srcDir, "common.kt")
+        val jvmSrc = File(srcDir, "jvm.kt")
+
+        FileUtil.writeToFile(commonSrc, common)
+        FileUtil.writeToFile(jvmSrc, jvm)
+
+        val jvmDest = File(tmpdir, "jvm").absolutePath
+
+        val result = K2JVMCompiler().compile(
+            jvmSrc,
+            commonSrc,
+            "-d", jvmDest,
+            *optionalArgs,
+            *jvmOnlyArgs
+        )
+
+        val files = File(jvmDest).listFiles()
+
+        if (files == null || files.isEmpty()) {
+            assertEquals(output.trimIndent(), result)
+            return
+        }
+
+        val sb = StringBuilder()
+
+        files
+            .filter { it.extension == "class" }
+            .sortedBy { it.absolutePath }
+            .distinctBy { it.name }
+            .forEach {
+                val os = ByteArrayOutputStream()
+                val printWriter = PrintWriter(os)
+                val writer = TraceClassVisitor(printWriter)
+                val reader = ClassReader(it.inputStream())
+                reader.accept(
+                    writer,
+                    ClassReader.SKIP_CODE or ClassReader.SKIP_DEBUG or ClassReader.SKIP_FRAMES
+                )
+                sb.append(os.toString())
+                sb.appendln()
+            }
+
+        assertEquals(output.trimIndent(), printPublicApi(sb.toString(), "test"))
+    }
+
+    private fun CLICompiler<*>.compile(
+        sources: File,
+        commonSources: File?,
+        vararg mainArguments: String
+    ): String = buildString {
+        val (output, exitCode) = executeCompilerGrabOutput(
+            this@compile,
+            listOfNotNull(sources.absolutePath,
+                commonSources?.absolutePath,
+                commonSources?.absolutePath?.let("-Xcommon-sources="::plus)) +
+                    "-Xmulti-platform" + mainArguments
+        )
+        appendln("Exit code: $exitCode")
+        appendln("Output:")
+        appendln(output)
+    }.trimTrailingWhitespacesAndAddNewlineAtEOF().trimEnd('\r', '\n')
+}

plugins/compose-compiler-hosted/integration-tests/src/test/java/androidx/compose/plugins/kotlin/CodegenTestFiles.kt

@@ -0,0 +1,49 @@
+/*
+ * Copyright 2019 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.compose.plugins.kotlin
+
+import com.intellij.openapi.project.Project
+import org.jetbrains.kotlin.checkers.utils.CheckerTestUtil
+import org.jetbrains.kotlin.psi.KtFile
+
+class CodegenTestFiles private constructor(
+    val psiFiles: List<KtFile>
+) {
+
+    companion object {
+        fun create(ktFiles: List<KtFile>): CodegenTestFiles {
+            assert(!ktFiles.isEmpty()) { "List should have at least one file" }
+            return CodegenTestFiles(ktFiles)
+        }
+
+        fun create(
+            fileName: String,
+            contentWithDiagnosticMarkup: String,
+            project: Project
+        ): CodegenTestFiles {
+            val content = CheckerTestUtil.parseDiagnosedRanges(
+                contentWithDiagnosticMarkup,
+                ArrayList(),
+                null
+            )
+
+            val file = createFile(fileName, content, project)
+
+            return create(listOf(file))
+        }
+    }
+}

plugins/compose-compiler-hosted/integration-tests/src/test/java/androidx/compose/plugins/kotlin/ComposeCallResolutionDiagnosticsTests.kt

@@ -0,0 +1,78 @@
+/*
+ * Copyright 2019 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.compose.plugins.kotlin
+
+class ComposeCallResolutionDiagnosticsTests : AbstractComposeDiagnosticsTest() {
+
+    private var isSetup = false
+    override fun setUp() {
+        isSetup = true
+        super.setUp()
+    }
+
+    private fun <T> ensureSetup(block: () -> T): T {
+        if (!isSetup) setUp()
+        return block()
+    }
+
+    private fun setupAndDoTest(text: String) = ensureSetup { doTest(text) }
+
+    fun testImplicitlyPassedReceiverScope1() = setupAndDoTest(
+        """
+            import androidx.compose.*
+            import android.widget.*
+            import android.os.Bundle
+            import android.app.Activity
+            import android.widget.FrameLayout
+
+            val x: Any? = null
+
+            fun Activity.setViewContent(composable: @Composable() () -> Unit): Composition? {
+                assert(composable != x)
+                return null
+            }
+
+            open class WebComponentActivity : Activity() {
+                override fun onCreate(savedInstanceState: Bundle?) {
+                    super.onCreate(savedInstanceState)
+
+                    setViewContent {
+                        FrameLayout {
+                        }
+                    }
+                }
+            }
+        """
+    )
+
+    fun testSimpleReceiverScope() = setupAndDoTest(
+        """
+            import android.widget.FrameLayout
+            import androidx.compose.Composable
+            import androidx.compose.ViewComposer
+
+            class SomeScope {
+             val composer: ViewComposer get() = error("should not be called")
+            }
+
+            @Composable fun SomeScope.foo() {
+                FrameLayout { }
+            }
+
+        """
+    )
+}

plugins/compose-compiler-hosted/integration-tests/src/test/java/androidx/compose/plugins/kotlin/ComposeCallResolverTests.kt

@@ -0,0 +1,294 @@
+/*
+ * Copyright 2019 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.compose.plugins.kotlin
+
+import com.intellij.psi.PsiElement
+import com.intellij.psi.util.PsiTreeUtil
+import org.jetbrains.kotlin.psi.KtBlockExpression
+import org.jetbrains.kotlin.psi.KtDeclaration
+import org.jetbrains.kotlin.psi.KtElement
+import org.jetbrains.kotlin.psi.KtFile
+import org.jetbrains.kotlin.psi.KtPsiFactory
+import org.jetbrains.kotlin.resolve.BindingContext
+import org.jetbrains.kotlin.resolve.calls.callUtil.getResolvedCall
+import org.jetbrains.kotlin.resolve.calls.model.ResolvedCall
+import kotlin.reflect.KClass
+
+class ComposeCallResolverTests : AbstractCodegenTest() {
+
+    fun testProperties() = assertInterceptions(
+        """
+            import androidx.compose.*
+
+            @Composable val foo get() = 123
+
+            class A {
+                @Composable val bar get() = 123
+            }
+
+            @Composable val A.bam get() = 123
+
+            @Composable
+            fun test() {
+                val a = A()
+                <call>foo
+                a.<call>bar
+                a.<call>bam
+            }
+        """
+    )
+
+    fun testBasicCallTypes() = assertInterceptions(
+        """
+            import androidx.compose.*
+            import android.widget.TextView
+
+            @Composable fun Foo() {}
+
+            fun Bar() {}
+
+            @Composable
+            fun test() {
+                <call>Foo()
+                <emit>TextView(text="text")
+                <normal>Bar()
+            }
+        """
+    )
+
+    fun testReceiverScopeCall() = assertInterceptions(
+        """
+            import androidx.compose.*
+
+            @Composable fun Int.Foo() {}
+
+            @Composable
+            fun test() {
+                val x = 1
+                x.<call>Foo()
+
+                with(x) {
+                    <call>Foo()
+                }
+            }
+        """
+    )
+
+    fun testInvokeOperatorCall() = assertInterceptions(
+        """
+            import androidx.compose.*
+
+            @Composable operator fun Int.invoke(y: Int) {}
+
+            @Composable
+            fun test() {
+                val x = 1
+                <call>x(y=10)
+            }
+        """
+    )
+
+    fun testComposableLambdaCall() = assertInterceptions(
+        """
+            import androidx.compose.*
+
+            @Composable
+            fun test(children: @Composable() () -> Unit) {
+                <call>children()
+            }
+        """
+    )
+
+    fun testComposableLambdaCallWithGenerics() = assertInterceptions(
+        """
+            import androidx.compose.*
+
+            @Composable fun <T> A(value: T, block: @Composable() (T) -> Unit) {
+                <call>block(value)
+            }
+
+            @Composable fun <T> B(
+                value: T,
+                block: @Composable() (@Composable() (T) -> Unit) -> Unit
+            ) {
+                <call>block({ })
+            }
+
+            @Composable
+            fun test() {
+                <call>A(123) { it ->
+                    println(it)
+                }
+                <call>B(123) { it ->
+                    <call>it(456)
+                }
+            }
+        """
+    )
+
+    // TODO(chuckj): Replace with another nested function call.
+    fun xtestMethodInvocations() = assertInterceptions(
+        """
+            import androidx.compose.*
+
+            val x = Ambient.of<Int> { 123 }
+
+            @Composable
+            fun test() {
+                x.<call>Provider(456) {
+
+                }
+            }
+        """
+    )
+
+    fun testReceiverLambdaInvocation() = assertInterceptions(
+        """
+            import androidx.compose.*
+
+            class TextSpanScope internal constructor(val composer: ViewComposer)
+
+            @Composable fun Foo(
+                scope: TextSpanScope, 
+                composable: @Composable() TextSpanScope.() -> Unit
+            ) {
+                with(scope) {
+                    <call>composable()
+                }
+            }
+        """
+    )
+
+    fun testReceiverLambda2() = assertInterceptions(
+        """
+            import androidx.compose.*
+
+            class DensityScope(val density: Density)
+
+            class Density
+
+            val DensityAmbient = Ambient.of<Density>()
+
+            @Composable
+            fun ambientDensity() = ambient(DensityAmbient)
+
+            @Composable
+            fun WithDensity(block: @Composable DensityScope.() -> Unit) {
+                DensityScope(ambientDensity()).<call>block()
+            }
+        """
+    )
+
+    fun testInlineChildren() = assertInterceptions(
+        """
+            import androidx.compose.*
+            import android.widget.LinearLayout
+
+            @Composable
+            inline fun PointerInputWrapper(
+                crossinline children: @Composable() () -> Unit
+            ) {
+                // Hide the internals of PointerInputNode
+                <emit>LinearLayout {
+                    <call>children()
+                }
+            }
+        """
+    )
+
+    private fun <T> setup(block: () -> T): T {
+        return block()
+    }
+
+    fun assertInterceptions(srcText: String) = setup {
+        val (text, carets) = extractCarets(srcText)
+
+        val environment = myEnvironment ?: error("Environment not initialized")
+
+        val ktFile = KtPsiFactory(environment.project).createFile(text)
+        val bindingContext = JvmResolveUtil.analyze(
+            ktFile,
+            environment
+        ).bindingContext
+
+        carets.forEachIndexed { index, (offset, calltype) ->
+            val resolvedCall = resolvedCallAtOffset(bindingContext, ktFile, offset)
+                ?: error("No resolved call found at index: $index, offset: $offset. Expected " +
+                    "$calltype.")
+
+            when (calltype) {
+                "<normal>" -> assert(!resolvedCall.isCall() && !resolvedCall.isEmit())
+                "<emit>" -> assert(resolvedCall.isEmit())
+                "<call>" -> assert(resolvedCall.isCall())
+                else -> error("Call type of $calltype not recognized.")
+            }
+        }
+    }
+
+    private fun ResolvedCall<*>.isEmit(): Boolean = candidateDescriptor is ComposableEmitDescriptor
+    private fun ResolvedCall<*>.isCall(): Boolean =
+        when (candidateDescriptor) {
+            is ComposableFunctionDescriptor -> true
+            is ComposablePropertyDescriptor -> true
+            else -> false
+        }
+
+    private val callPattern = Regex("(<normal>)|(<emit>)|(<call>)")
+    private fun extractCarets(text: String): Pair<String, List<Pair<Int, String>>> {
+        val indices = mutableListOf<Pair<Int, String>>()
+        var offset = 0
+        val src = callPattern.replace(text) {
+            indices.add(it.range.first - offset to it.value)
+            offset += it.range.last - it.range.first + 1
+            ""
+        }
+        return src to indices
+    }
+
+    private fun resolvedCallAtOffset(
+        bindingContext: BindingContext,
+        jetFile: KtFile,
+        index: Int
+    ): ResolvedCall<*>? {
+        val element = jetFile.findElementAt(index)!!
+        return element.getNearestResolvedCall(bindingContext)
+    }
+}
+
+fun PsiElement?.getNearestResolvedCall(bindingContext: BindingContext): ResolvedCall<*>? {
+    var node: PsiElement? = this
+    while (node != null) {
+        when (node) {
+            is KtBlockExpression,
+            is KtDeclaration -> return null
+            is KtElement -> {
+                val resolvedCall = node.getResolvedCall(bindingContext)
+                if (resolvedCall != null) {
+                    return resolvedCall
+                }
+            }
+        }
+        node = node.parent
+    }
+    return null
+}
+
+private inline fun <reified T : PsiElement> PsiElement.parentOfType(): T? = parentOfType(T::class)
+
+private fun <T : PsiElement> PsiElement.parentOfType(vararg classes: KClass<out T>): T? {
+    return PsiTreeUtil.getParentOfType(this, *classes.map { it.java }.toTypedArray())
+}

plugins/compose-compiler-hosted/integration-tests/src/test/java/androidx/compose/plugins/kotlin/ComposeCodegenTestUtils.kt

@@ -0,0 +1,33 @@
+/*
+ * Copyright 2019 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.
+ */
+
+@file:Suppress("MemberVisibilityCanBePrivate")
+
+package androidx.compose.plugins.kotlin
+
+var uniqueNumber = 0
+
+fun loadClass(loader: ClassLoader, name: String?, bytes: ByteArray): Class<*> {
+    val defineClassMethod = ClassLoader::class.javaObjectType.getDeclaredMethod(
+        "defineClass",
+        String::class.javaObjectType,
+        ByteArray::class.javaObjectType,
+        Int::class.javaPrimitiveType,
+        Int::class.javaPrimitiveType
+    )
+    defineClassMethod.isAccessible = true
+    return defineClassMethod.invoke(loader, name, bytes, 0, bytes.size) as Class<*>
+}

plugins/compose-compiler-hosted/integration-tests/src/test/java/androidx/compose/plugins/kotlin/ComposeMultiPlatformTests.kt

@@ -0,0 +1,70 @@
+/*
+ * Copyright 2020 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.compose.plugins.kotlin
+
+import org.junit.Test
+
+class ComposeMultiPlatformTests : AbstractMultiPlatformIntegrationTest() {
+    @Test
+    fun testBasicMpp() = ensureSetup {
+        multiplatform(
+            """
+            expect val foo: String
+        """,
+            """
+            actual val foo = ""
+        """,
+            """
+            public final class JvmKt {
+              private final static Ljava/lang/String; foo = ""
+              public final static getFoo()Ljava/lang/String;
+              public final static <clinit>()V
+            }
+        """
+        )
+    }
+
+    @Test
+    fun testBasicComposable() = ensureSetup {
+        multiplatform(
+        """
+            import androidx.compose.Composable
+
+            expect @Composable fun Test()
+        """,
+        """
+            import androidx.compose.Composable
+
+            actual @Composable fun Test() {}
+        """,
+        """
+        final class JvmKt%Test%1 extends kotlin/jvm/internal/Lambda implements kotlin/jvm/functions/Function1 {
+          OUTERCLASS JvmKt Test (Landroidx/compose/Composer;)V
+          final static INNERCLASS JvmKt%Test%1 null null
+          synthetic <init>()V
+          public final invoke(Landroidx/compose/Composer;)V
+          public synthetic bridge invoke(Ljava/lang/Object;)Ljava/lang/Object;
+        }
+        public final class JvmKt {
+          final static INNERCLASS JvmKt%Test%1 null null
+          public final static Test(Landroidx/compose/Composer;)V
+          public final static synthetic Test()V
+        }
+        """
+        )
+    }
+}

plugins/compose-compiler-hosted/integration-tests/src/test/java/androidx/compose/plugins/kotlin/ComposeRobolectricTestRunner.kt

@@ -0,0 +1,30 @@
+/*
+ * Copyright 2019 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.compose.plugins.kotlin
+
+import org.junit.runners.model.FrameworkMethod
+import org.robolectric.RobolectricTestRunner
+import org.robolectric.internal.bytecode.InstrumentationConfiguration
+
+class ComposeRobolectricTestRunner(testClass: Class<*>) : RobolectricTestRunner(testClass) {
+    override fun createClassLoaderConfig(method: FrameworkMethod?): InstrumentationConfiguration {
+        val builder = InstrumentationConfiguration.Builder(super.createClassLoaderConfig(method))
+        builder.doNotInstrumentPackage("androidx.compose")
+        builder.doNotInstrumentPackage("androidx.ui")
+        return builder.build()
+    }
+}

plugins/compose-compiler-hosted/integration-tests/src/test/java/androidx/compose/plugins/kotlin/FcsModelCodeGenTests.kt

@@ -0,0 +1,362 @@
+/*
+ * Copyright 2019 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.compose.plugins.kotlin
+
+import android.app.Activity
+import android.view.ViewGroup
+import android.widget.TextView
+import org.junit.Before
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.robolectric.RuntimeEnvironment
+import org.robolectric.annotation.Config
+import androidx.compose.FrameManager
+import org.junit.Ignore
+
+@RunWith(ComposeRobolectricTestRunner::class)
+@Config(
+    manifest = Config.NONE,
+    minSdk = 23,
+    maxSdk = 23
+)
+@Ignore("b/142799174 - Flaky tests")
+class FcsModelCodeGenTests : AbstractCodegenTest() {
+
+    @Before
+    fun before() {
+        val scheduler = RuntimeEnvironment.getMasterScheduler()
+        scheduler.pause()
+    }
+
+    override fun setUp() {
+        isSetup = true
+        super.setUp()
+    }
+
+    private var isSetup = false
+    private inline fun <T> ensureSetup(crossinline block: () -> T): T {
+        if (!isSetup) setUp()
+        return FrameManager.isolated { block() }
+    }
+
+    @Test
+    fun testCGModelView_PersonModel(): Unit = ensureSetup {
+        val tvNameId = 384
+        val tvAgeId = 385
+
+        var name = PRESIDENT_NAME_1
+        var age = PRESIDENT_AGE_1
+        compose(
+            """
+            @Model
+            class FcsPerson4(var name: String, var age: Int)
+
+            @Composable
+            fun PersonView4(person: FcsPerson4) {
+              Observe {
+                TextView(text=person.name, id=$tvNameId)
+                TextView(text=person.age.toString(), id=$tvAgeId)
+              }
+            }
+
+            val president = FcsPerson4("$PRESIDENT_NAME_1", $PRESIDENT_AGE_1)
+            """, { mapOf("name" to name, "age" to age) }, """
+               president.name = name
+               president.age = age
+            """, """
+                PersonView4(person=president)
+            """).then { activity ->
+            val tvName = activity.findViewById(tvNameId) as TextView
+            val tvAge = activity.findViewById(tvAgeId) as TextView
+            assertEquals(PRESIDENT_NAME_1, tvName.text)
+            assertEquals(PRESIDENT_AGE_1.toString(), tvAge.text)
+
+            name = PRESIDENT_NAME_16
+            age = PRESIDENT_AGE_16
+        }.then { activity ->
+            val tvName = activity.findViewById(tvNameId) as TextView
+            val tvAge = activity.findViewById(tvAgeId) as TextView
+            assertEquals(PRESIDENT_NAME_16, tvName.text)
+            assertEquals(PRESIDENT_AGE_16.toString(), tvAge.text)
+        }
+    }
+
+    @Test // b/120843442
+    fun testCGModelView_ObjectModel(): Unit = ensureSetup {
+        val tvNameId = 384
+        val tvAgeId = 385
+
+        var name = PRESIDENT_NAME_1
+        var age = PRESIDENT_AGE_1
+        compose(
+            """
+            @Model
+            object fcs_president {
+                var name: String = "$PRESIDENT_NAME_1"
+                var age: Int = $PRESIDENT_AGE_1
+            }
+
+            @Composable
+            fun PresidentView() {
+              Observe {
+                TextView(text=fcs_president.name, id=$tvNameId)
+                TextView(text=fcs_president.age.toString(), id=$tvAgeId)
+              }
+            }
+
+            """, { mapOf("name" to name, "age" to age) }, """
+               fcs_president.name = name
+               fcs_president.age = age
+            """, """
+                PresidentView()
+            """).then { activity ->
+            val tvName = activity.findViewById(tvNameId) as TextView
+            val tvAge = activity.findViewById(tvAgeId) as TextView
+            assertEquals(PRESIDENT_NAME_1, tvName.text)
+            assertEquals(PRESIDENT_AGE_1.toString(), tvAge.text)
+
+            name = PRESIDENT_NAME_16
+            age = PRESIDENT_AGE_16
+        }.then { activity ->
+            val tvName = activity.findViewById(tvNameId) as TextView
+            val tvAge = activity.findViewById(tvAgeId) as TextView
+            assertEquals(PRESIDENT_NAME_16, tvName.text)
+            assertEquals(PRESIDENT_AGE_16.toString(), tvAge.text)
+        }
+    }
+
+    @Test // b/120836313
+    fun testCGModelView_DataModel(): Unit = ensureSetup {
+        val tvNameId = 384
+        val tvAgeId = 385
+
+        var name = PRESIDENT_NAME_1
+        var age = PRESIDENT_AGE_1
+        compose(
+            """
+            @Model
+            data class FcsPersonB(var name: String, var age: Int)
+
+            @Composable
+            fun PersonViewB(person: FcsPersonB) {
+              Observe {
+                TextView(text=person.name, id=$tvNameId)
+                TextView(text=person.age.toString(), id=$tvAgeId)
+              }
+            }
+
+            val president = FcsPersonB("$PRESIDENT_NAME_1", $PRESIDENT_AGE_1)
+            """, { mapOf("name" to name, "age" to age) }, """
+               president.name = name
+               president.age = age
+            """, """
+                PersonViewB(person=president)
+            """).then { activity ->
+            val tvName = activity.findViewById(tvNameId) as TextView
+            val tvAge = activity.findViewById(tvAgeId) as TextView
+            assertEquals(PRESIDENT_NAME_1, tvName.text)
+            assertEquals(PRESIDENT_AGE_1.toString(), tvAge.text)
+
+            name = PRESIDENT_NAME_16
+            age = PRESIDENT_AGE_16
+        }.then { activity ->
+            val tvName = activity.findViewById(tvNameId) as TextView
+            val tvAge = activity.findViewById(tvAgeId) as TextView
+            assertEquals(PRESIDENT_NAME_16, tvName.text)
+            assertEquals(PRESIDENT_AGE_16.toString(), tvAge.text)
+        }
+    }
+
+    @Test // b/120843442
+    fun testCGModelView_ZeroFrame(): Unit = ensureSetup {
+        val tvNameId = 384
+        val tvAgeId = 385
+
+        var name = PRESIDENT_NAME_1
+        var age = PRESIDENT_AGE_1
+        compose(
+            """
+            @Model
+            class FcsPersonC(var name: String, var age: Int)
+
+            @Composable
+            fun PersonViewC(person: FcsPersonC) {
+              Observe {
+                TextView(text=person.name, id=$tvNameId)
+                TextView(text=person.age.toString(), id=$tvAgeId)
+              }
+            }
+
+            val president = FrameManager.unframed { FcsPersonC("$PRESIDENT_NAME_1", ${
+                PRESIDENT_AGE_1}) }
+            """, { mapOf("name" to name, "age" to age) }, """
+               president.name = name
+               president.age = age
+            """, """
+                PersonViewC(person=president)
+            """).then { activity ->
+            val tvName = activity.findViewById(tvNameId) as TextView
+            val tvAge = activity.findViewById(tvAgeId) as TextView
+            assertEquals(name, tvName.text)
+            assertEquals(age.toString(), tvAge.text)
+
+            name = PRESIDENT_NAME_16
+            age = PRESIDENT_AGE_16
+        }.then { activity ->
+            val tvName = activity.findViewById(tvNameId) as TextView
+            val tvAge = activity.findViewById(tvAgeId) as TextView
+            assertEquals(PRESIDENT_NAME_16, tvName.text)
+            assertEquals(PRESIDENT_AGE_16.toString(), tvAge.text)
+        }
+    }
+
+    @Test // b/120843442
+    fun testCGModelView_ZeroFrame_Modification(): Unit = ensureSetup {
+        val tvNameId = 384
+        val tvAgeId = 385
+
+        var name = PRESIDENT_NAME_1
+        var age = PRESIDENT_AGE_1
+        compose(
+            """
+            @Model
+            class FcsPersonD(var name: String, var age: Int)
+
+            @Composable
+            fun PersonViewD(person: FcsPersonD) {
+              Observe {
+                TextView(text=person.name, id=$tvNameId)
+                TextView(text=person.age.toString(), id=$tvAgeId)
+              }
+            }
+
+            val president = FrameManager.framed { FcsPersonD("$PRESIDENT_NAME_1", ${
+            PRESIDENT_AGE_1}).apply { age = $PRESIDENT_AGE_1 } }
+            """, { mapOf("name" to name, "age" to age) }, """
+               president.name = name
+               president.age = age
+            """, """
+                PersonViewD(person=president)
+            """).then { activity ->
+            val tvName = activity.findViewById(tvNameId) as TextView
+            val tvAge = activity.findViewById(tvAgeId) as TextView
+            assertEquals(PRESIDENT_NAME_1, tvName.text)
+            assertEquals(PRESIDENT_AGE_1.toString(), tvAge.text)
+
+            name = PRESIDENT_NAME_16
+            age = PRESIDENT_AGE_16
+        }.then { activity ->
+            val tvName = activity.findViewById(tvNameId) as TextView
+            val tvAge = activity.findViewById(tvAgeId) as TextView
+            assertEquals(PRESIDENT_NAME_16, tvName.text)
+            assertEquals(PRESIDENT_AGE_16.toString(), tvAge.text)
+        }
+    }
+
+    @Test
+    fun testCGModelView_LambdaInInitializer(): Unit = ensureSetup {
+        // Check that the lambda gets moved correctly.
+        compose("""
+            @Model
+            class Database(val name: String) {
+              var queries = (1..10).map { "query ${'$'}it" }
+            }
+
+            @Composable
+            fun View() {
+                val d = Database("some")
+                TextView(text = "query = ${'$'}{d.queries[0]}", id=100)
+            }
+        """, { emptyMap<String, Any>() }, "", "View()").then {
+            val tv = it.findViewById<TextView>(100)
+            assertEquals("query = query 1", tv.text)
+        }
+    }
+
+    fun compose(
+        prefix: String,
+        valuesFactory: () -> Map<String, Any>,
+        advance: String,
+        composition: String,
+        dumpClasses: Boolean = false
+    ): RobolectricComposeTester {
+        val className = "Test_${uniqueNumber++}"
+        val fileName = "$className.kt"
+
+        val candidateValues = valuesFactory()
+
+        @Suppress("NO_REFLECTION_IN_CLASS_PATH")
+        val parameterList = candidateValues.map {
+            "${it.key}: ${it.value::class.qualifiedName}"
+        }.joinToString()
+        val parameterTypes = candidateValues.map {
+            it.value::class.javaPrimitiveType ?: it.value::class.javaObjectType
+        }.toTypedArray()
+
+        val compiledClasses = classLoader("""
+           import android.content.Context
+           import android.widget.*
+           import androidx.compose.*
+
+           $prefix
+
+           class $className {
+
+             fun compose() {
+               $composition
+             }
+
+             fun advance($parameterList) {
+               $advance
+             }
+           }
+        """, fileName, dumpClasses)
+
+        val allClassFiles = compiledClasses.allGeneratedFiles.filter {
+            it.relativePath.endsWith(".class")
+        }
+
+        val instanceClass = run {
+            var instanceClass: Class<*>? = null
+            var loadedOne = false
+            for (outFile in allClassFiles) {
+                val bytes = outFile.asByteArray()
+                val loadedClass = loadClass(
+                    this.javaClass.classLoader!!,
+                    null,
+                    bytes
+                )
+                if (loadedClass.name == className) instanceClass = loadedClass
+                loadedOne = true
+            }
+            if (!loadedOne) error("No classes loaded")
+            instanceClass ?: error("Could not find class $className in loaded classes")
+        }
+
+        val instanceOfClass = instanceClass.newInstance()
+
+        return composeMulti({
+            val composeMethod = instanceClass.getMethod("compose")
+            composeMethod.invoke(instanceOfClass)
+        }) {
+            val values = valuesFactory()
+            val arguments = values.map { it.value }.toTypedArray()
+            val advanceMethod = instanceClass.getMethod("advance", *parameterTypes)
+            advanceMethod.invoke(instanceOfClass, *arguments)
+        }
+    }
+}

plugins/compose-compiler-hosted/integration-tests/src/test/java/androidx/compose/plugins/kotlin/FcsTypeResolutionTests.kt

@@ -0,0 +1,512 @@
+/*
+ * Copyright 2019 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.compose.plugins.kotlin
+
+class FcsTypeResolutionTests : AbstractComposeDiagnosticsTest() {
+
+    fun testImplicitlyPassedReceiverScope1() = doTest(
+        """
+            import androidx.compose.*
+
+            @Composable
+            fun Int.Foo(children: @Composable Int.() -> Unit) {
+                children()
+            }
+        """
+    )
+
+    fun testImplicitlyPassedReceiverScope2() = doTest(
+        """
+            import androidx.compose.*
+
+            @Composable
+            fun Int.Foo(children: @Composable Int.(foo: String) -> Unit) {
+                children(<!NO_VALUE_FOR_PARAMETER, NO_VALUE_FOR_PARAMETER!>)<!>
+            }
+
+            @Composable
+            fun Bar(children: @Composable Int.() -> Unit) {
+                children(<!NO_VALUE_FOR_PARAMETER!>)<!>
+            }
+        """
+    )
+
+    fun testSmartCastsAndPunning() = doTest(
+        """
+            import androidx.compose.*
+
+            @Composable
+            fun Foo(bar: String) { print(bar) }
+
+            @Composable
+            fun test(bar: String?) {
+                Foo(<!TYPE_MISMATCH!>bar<!>)
+                if (bar != null) {
+                    Foo(bar)
+                    Foo(bar=bar)
+                }
+            }
+        """
+    )
+
+    fun testExtensionInvoke() = doTest(
+        """
+            import androidx.compose.*
+
+            class Foo {}
+            @Composable operator fun Foo.invoke() {}
+
+            @Composable fun test() {
+                Foo()
+            }
+        """
+    )
+
+    fun testResolutionInsideWhenExpression() = doTest(
+        """
+            import androidx.compose.*
+            import android.widget.TextView
+
+            @Composable fun doSomething(foo: Boolean) {
+                when (foo) {
+                    true -> TextView(text="Started...")
+                    false -> TextView(text="Continue...")
+                }
+            }
+        """
+    )
+
+    fun testUsedParameters() = doTest(
+        """
+            import androidx.compose.*
+            import android.widget.LinearLayout
+
+            @Composable fun Foo(x: Int, composeItem: @Composable() () -> Unit = {}) {
+                println(x)
+                print(composeItem == {})
+            }
+
+            @Composable fun test(
+                children: @Composable() () -> Unit,
+                value: Int,
+                x: Int,
+                children2: @Composable() () -> Unit,
+                value2: Int
+            ) {
+                LinearLayout {
+                    // named argument
+                    Foo(x=value)
+
+                    // indexed argument
+                    Foo(x)
+
+                    // tag
+                    children()
+                }
+                Foo(x=123, composeItem={
+                    val abc = 123
+
+                    // attribute value
+                    Foo(x=abc)
+
+                    // attribute value
+                    Foo(x=value2)
+
+                    // tag
+                    children2()
+                })
+            }
+        """
+    )
+
+    fun testDispatchInvoke() = doTest(
+        """
+            import androidx.compose.*
+
+            class Bam {
+                @Composable fun Foo() {}
+            }
+
+            @Composable fun test() {
+                with(Bam()) {
+                    Foo()
+                }
+            }
+        """
+    )
+
+    fun testDispatchAndExtensionReceiver() = doTest(
+        """
+            import androidx.compose.*
+
+            class Bam {
+                inner class Foo {}
+            }
+
+            @Composable operator fun Bam.Foo.invoke() {}
+
+            @Composable fun test() {
+                with(Bam()) {
+                    Foo()
+                }
+            }
+        """
+    )
+
+    fun testDispatchAndExtensionReceiverLocal() = doTest(
+        """
+            import androidx.compose.*
+
+            class Foo {}
+
+            class Bam {
+                @Composable operator fun Foo.invoke() {}
+                @Composable operator fun invoke() {
+                    Foo()
+                }
+            }
+
+        """
+    )
+
+    fun testMissingAttributes() = doTest(
+        """
+            import androidx.compose.*
+
+            data class Foo(val value: Int)
+
+            @Composable fun A(x: Foo) { println(x) }
+
+            // NOTE: It's important that the diagnostic be only over the call target, and not the
+            // entire element so that a single error doesn't end up making a huge part of an 
+            // otherwise correct file "red".
+            @Composable fun Test(F: @Composable() (x: Foo) -> Unit) {
+                // NOTE: constructor attributes and fn params get a "missing parameter" diagnostic
+                A(<!NO_VALUE_FOR_PARAMETER!>)<!>
+
+                // local
+                F(<!NO_VALUE_FOR_PARAMETER!>)<!>
+
+                val x = Foo(123)
+
+                A(x)
+                F(x)
+                A(x=x)
+                F(x=x)
+            }
+
+        """.trimIndent()
+    )
+
+    fun testDuplicateAttributes() = doTest(
+        """
+            import androidx.compose.*
+
+            data class Foo(val value: Int)
+
+            @Composable fun A(x: Foo) { println(x) }
+
+            @Composable fun Test() {
+                val x = Foo(123)
+
+                // NOTE: It's important that the diagnostic be only over the attribute key, so that
+                // we don't make a large part of the elements red when the type is otherwise correct
+                A(x=x, <!ARGUMENT_PASSED_TWICE!>x<!>=x)
+            }
+
+        """.trimIndent()
+    )
+
+    fun testChildrenNamedAndBodyDuplicate() = doTest(
+        """
+            import androidx.compose.*
+
+            @Composable fun A(children: @Composable() () -> Unit) { children() }
+
+            @Composable fun Test() {
+                A(children={}) <!TOO_MANY_ARGUMENTS!>{ }<!>
+            }
+
+        """.trimIndent()
+    )
+
+    fun testAbstractClassTags() = doTest(
+        """
+            import androidx.compose.*
+            import android.content.Context
+            import android.widget.LinearLayout
+
+            abstract class Foo {}
+
+            abstract class Bar(context: Context) : LinearLayout(context) {}
+
+            @Composable fun Test() {
+                <!CREATING_AN_INSTANCE_OF_ABSTRACT_CLASS!>Foo()<!>
+                <!CREATING_AN_INSTANCE_OF_ABSTRACT_CLASS!>Bar(<!NO_VALUE_FOR_PARAMETER!>)<!><!>
+            }
+
+        """.trimIndent()
+    )
+
+    fun testGenerics() = doTest(
+        """
+            import androidx.compose.*
+
+            class A { fun a() {} }
+            class B { fun b() {} }
+
+            @Composable fun <T> Bar(x: Int, value: T, f: (T) -> Unit) { println(value); println(f); println(x) }
+
+            @Composable fun Test() {
+
+                val fa: (A) -> Unit = { it.a() }
+                val fb: (B) -> Unit = { it.b() }
+
+                Bar(x=1, value=A(), f={ it.a() })
+                Bar(x=1, value=B(), f={ it.b() })
+                Bar(x=1, value=A(), f=fa)
+                Bar(x=1, value=B(), f=fb)
+                Bar(x=1, value=B(), f={ it.<!UNRESOLVED_REFERENCE!>a<!>() })
+                Bar(x=1, value=A(), f={ it.<!UNRESOLVED_REFERENCE!>b<!>() })
+                <!TYPE_INFERENCE_CONFLICTING_SUBSTITUTIONS!>Bar<!>(
+                  x=1, 
+                  value=A(), 
+                  f=fb
+                )
+                <!TYPE_INFERENCE_CONFLICTING_SUBSTITUTIONS!>Bar<!>(
+                  x=1,
+                  value=B(), 
+                  f=fa
+                )
+            }
+
+        """.trimIndent()
+    )
+
+    fun testUnresolvedAttributeValueResolvedTarget() = doTest(
+        """
+            import androidx.compose.*
+
+            @Composable fun Fam(bar: Int, x: Int) {
+                print(bar)
+                print(x)
+            }
+
+            @Composable fun Test() {
+                Fam(
+                  bar=<!UNRESOLVED_REFERENCE!>undefined<!>,
+                  x=1
+                )
+                Fam(
+                  bar=1,
+                  x=<!UNRESOLVED_REFERENCE!>undefined<!>
+                )
+                Fam(
+                  <!UNRESOLVED_REFERENCE!>bar<!>,
+                  <!UNRESOLVED_REFERENCE!>x<!>
+                )
+
+                Fam(
+                  bar=<!TYPE_MISMATCH!>""<!>,
+                  x=<!TYPE_MISMATCH!>""<!>
+                )
+            }
+
+        """.trimIndent()
+    )
+
+    // TODO(lmr): this triggers an exception!
+    fun testEmptyAttributeValue() = doTest(
+        """
+            import androidx.compose.*
+
+            @Composable fun Foo(abc: Int, xyz: Int) {
+                print(abc)
+                print(xyz)
+            }
+
+            @Composable fun Test() {
+                Foo(abc=<!NO_VALUE_FOR_PARAMETER!>)<!>
+
+                // NOTE(lmr): even though there is NO diagnostic here, there *is* a parse
+                // error. This is intentional and done to mimic how kotlin handles function
+                // calls with no value expression in a call parameter list (ie, `Foo(123,)`)
+                Foo(abc=123, xyz=)
+            }
+
+        """.trimIndent()
+    )
+
+    fun testMismatchedAttributes() = doTest(
+        """
+            import androidx.compose.*
+
+            open class A {}
+            class B : A() {}
+
+            @Composable fun Foo(x: A = A(), y: A = B(), z: B = B()) {
+                print(x)
+                print(y)
+                print(z)
+            }
+
+            @Composable fun Test() {
+                Foo(
+                    x=A(),
+                    y=A(),
+                    z=<!TYPE_MISMATCH!>A()<!>
+                )
+                Foo(
+                    x=B(),
+                    y=B(),
+                    z=B()
+                )
+                Foo(
+                    x=<!CONSTANT_EXPECTED_TYPE_MISMATCH!>1<!>,
+                    y=<!CONSTANT_EXPECTED_TYPE_MISMATCH!>1<!>,
+                    z=<!CONSTANT_EXPECTED_TYPE_MISMATCH!>1<!>
+                )
+            }
+
+        """.trimIndent()
+    )
+
+    fun testErrorAttributeValue() = doTest(
+        """
+            import androidx.compose.*
+
+            @Composable fun Foo(x: Int = 1) { print(x) }
+
+            @Composable fun Test() {
+                Foo(
+                    x=<!UNRESOLVED_REFERENCE!>someUnresolvedValue<!>,
+                    <!NAMED_PARAMETER_NOT_FOUND!>y<!>=<!UNRESOLVED_REFERENCE!>someUnresolvedValue<!>
+                )
+            }
+
+        """.trimIndent()
+    )
+
+    fun testUnresolvedQualifiedTag() = doTest(
+        """
+            import androidx.compose.*
+
+            object MyNamespace {
+                @Composable fun Bar(children: @Composable() () -> Unit = {}) { 
+                    children() 
+                }
+
+                var Baz = @Composable { }
+
+                var someString = ""
+                class NonComponent {}
+            }
+
+            class Boo {
+                @Composable fun Wat() { }
+            }
+
+            @Composable fun Test() {
+
+                MyNamespace.Bar()
+                MyNamespace.Baz()
+                MyNamespace.<!UNRESOLVED_REFERENCE!>Qoo<!>()
+                MyNamespace.<!FUNCTION_EXPECTED!>someString<!>()
+                MyNamespace.NonComponent()
+                MyNamespace.Bar {}
+                MyNamespace.Baz <!TOO_MANY_ARGUMENTS!>{}<!>
+
+                val obj = Boo()
+                Boo.<!UNRESOLVED_REFERENCE!>Wat<!>()
+                obj.Wat()
+
+                MyNamespace.<!UNRESOLVED_REFERENCE!>Bam<!>()
+                <!UNRESOLVED_REFERENCE!>SomethingThatDoesntExist<!>.Foo()
+
+                obj.Wat <!TOO_MANY_ARGUMENTS!>{
+                }<!>
+
+                MyNamespace.<!UNRESOLVED_REFERENCE!>Qoo<!> {
+                }
+
+                MyNamespace.<!FUNCTION_EXPECTED!>someString<!> {
+                }
+
+                <!UNRESOLVED_REFERENCE!>SomethingThatDoesntExist<!>.Foo {
+                }
+
+                MyNamespace.NonComponent <!TOO_MANY_ARGUMENTS!>{}<!>
+
+                MyNamespace.<!UNRESOLVED_REFERENCE!>Bam<!> {}
+
+            }
+
+        """.trimIndent()
+    )
+
+    // TODO(lmr): overloads creates resolution exception
+    fun testChildren() = doTest(
+        """
+            import androidx.compose.*
+            import android.widget.Button
+            import android.widget.LinearLayout
+
+            @Composable fun ChildrenRequired2(children: @Composable() () -> Unit) { children() }
+
+            @Composable fun ChildrenOptional3(children: @Composable() () -> Unit = {}){ children() }
+
+            @Composable fun NoChildren2() {}
+
+            @Composable 
+            fun MultiChildren(c: @Composable() (x: Int) -> Unit = {}) { c(1) }
+
+            @Composable 
+            fun MultiChildren(c: @Composable() (x: Int, y: Int) -> Unit = { x, y ->println(x + y) }) { c(1,1) }
+
+            @Composable fun Test() {
+                ChildrenRequired2 {}
+                ChildrenRequired2(<!NO_VALUE_FOR_PARAMETER!>)<!>
+
+                ChildrenOptional3 {}
+                ChildrenOptional3()
+
+                NoChildren2 <!TOO_MANY_ARGUMENTS!>{}<!>
+                NoChildren2()
+
+                <!OVERLOAD_RESOLUTION_AMBIGUITY!>MultiChildren<!> {}
+                MultiChildren { x ->
+                    println(x)
+                }
+                MultiChildren { x, y ->
+                    println(x + y)
+                }
+                <!NONE_APPLICABLE!>MultiChildren<!> { <!CANNOT_INFER_PARAMETER_TYPE!>x<!>, 
+                <!CANNOT_INFER_PARAMETER_TYPE!>y<!>, <!CANNOT_INFER_PARAMETER_TYPE!>z<!> ->
+                    println(x + y + z)
+                }
+
+                Button()
+                LinearLayout()
+
+                LinearLayout {}
+
+                Button <!TOO_MANY_ARGUMENTS!>{}<!>
+            }
+
+        """.trimIndent()
+    )
+}

plugins/compose-compiler-hosted/integration-tests/src/test/java/androidx/compose/plugins/kotlin/GenerationUtils.kt

@@ -0,0 +1,105 @@
+/*
+ * Copyright 2019 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.compose.plugins.kotlin
+
+import com.intellij.psi.search.GlobalSearchScope
+import org.jetbrains.kotlin.backend.common.phaser.PhaseConfig
+import org.jetbrains.kotlin.backend.jvm.JvmIrCodegenFactory
+import org.jetbrains.kotlin.backend.jvm.defaultJvmPhases
+import org.jetbrains.kotlin.backend.jvm.withPluginPhases
+import org.jetbrains.kotlin.cli.common.CLIConfigurationKeys
+import org.jetbrains.kotlin.cli.jvm.compiler.KotlinCoreEnvironment
+import org.jetbrains.kotlin.cli.jvm.compiler.NoScopeRecordCliBindingTrace
+import org.jetbrains.kotlin.codegen.ClassBuilderFactories
+import org.jetbrains.kotlin.codegen.ClassBuilderFactory
+import org.jetbrains.kotlin.codegen.CompilationErrorHandler
+import org.jetbrains.kotlin.codegen.DefaultCodegenFactory
+import org.jetbrains.kotlin.codegen.KotlinCodegenFacade
+import org.jetbrains.kotlin.codegen.state.GenerationState
+import org.jetbrains.kotlin.config.CompilerConfiguration
+import org.jetbrains.kotlin.config.JVMConfigurationKeys
+import org.jetbrains.kotlin.load.kotlin.PackagePartProvider
+import org.jetbrains.kotlin.psi.KtFile
+import org.jetbrains.kotlin.resolve.AnalyzingUtils
+import org.jetbrains.kotlin.resolve.BindingTrace
+
+object GenerationUtils {
+    @JvmStatic
+    @JvmOverloads
+    fun compileFiles(
+        files: List<KtFile>,
+        environment: KotlinCoreEnvironment,
+        classBuilderFactory: ClassBuilderFactory = ClassBuilderFactories.TEST,
+        trace: BindingTrace = NoScopeRecordCliBindingTrace()
+    ): GenerationState =
+        compileFiles(
+            files,
+            environment.configuration,
+            classBuilderFactory,
+            environment::createPackagePartProvider,
+            trace
+        )
+
+    @JvmStatic
+    @JvmOverloads
+    fun compileFiles(
+        files: List<KtFile>,
+        configuration: CompilerConfiguration,
+        classBuilderFactory: ClassBuilderFactory,
+        packagePartProvider: (GlobalSearchScope) -> PackagePartProvider,
+        trace: BindingTrace = NoScopeRecordCliBindingTrace()
+    ): GenerationState {
+        val analysisResult =
+            JvmResolveUtil.analyzeAndCheckForErrors(
+                files.first().project,
+                files,
+                configuration,
+                packagePartProvider,
+                trace
+            )
+        analysisResult.throwIfError()
+
+        val state = GenerationState.Builder(
+            files.first().project,
+            classBuilderFactory,
+            analysisResult.moduleDescriptor,
+            analysisResult.bindingContext,
+            files,
+            configuration
+        ).codegenFactory(
+            if (configuration.getBoolean(JVMConfigurationKeys.IR))
+                JvmIrCodegenFactory(
+                    configuration.get(CLIConfigurationKeys.PHASE_CONFIG)
+                        ?: PhaseConfig(defaultJvmPhases).withPluginPhases(files.first().project)
+                )
+            else DefaultCodegenFactory
+        ).build()
+        if (analysisResult.shouldGenerateCode) {
+            KotlinCodegenFacade.compileCorrectFiles(state, CompilationErrorHandler
+                .THROW_EXCEPTION)
+        }
+
+        // For JVM-specific errors
+        try {
+            AnalyzingUtils.throwExceptionOnErrors(state.collectedExtraJvmDiagnostics)
+        } catch (e: Throwable) {
+            throw TestsCompilerError(e)
+        }
+
+        return state
+    }
+}

plugins/compose-compiler-hosted/integration-tests/src/test/java/androidx/compose/plugins/kotlin/JvmResolveUtil.kt

@@ -0,0 +1,99 @@
+/*
+ * Copyright 2019 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.compose.plugins.kotlin
+
+import com.intellij.openapi.project.Project
+import com.intellij.psi.search.GlobalSearchScope
+import org.jetbrains.kotlin.analyzer.AnalysisResult
+import org.jetbrains.kotlin.cli.jvm.compiler.CliBindingTrace
+import org.jetbrains.kotlin.cli.jvm.compiler.KotlinCoreEnvironment
+import org.jetbrains.kotlin.cli.jvm.compiler.TopDownAnalyzerFacadeForJVM
+import org.jetbrains.kotlin.config.CompilerConfiguration
+import org.jetbrains.kotlin.load.kotlin.PackagePartProvider
+import org.jetbrains.kotlin.psi.KtFile
+import org.jetbrains.kotlin.resolve.AnalyzingUtils
+import org.jetbrains.kotlin.resolve.BindingTrace
+
+object JvmResolveUtil {
+    @JvmStatic
+    fun analyzeAndCheckForErrors(
+        project: Project,
+        files: Collection<KtFile>,
+        configuration: CompilerConfiguration,
+        packagePartProvider: (GlobalSearchScope) -> PackagePartProvider,
+        trace: BindingTrace = CliBindingTrace()
+    ): AnalysisResult {
+        for (file in files) {
+            try {
+                AnalyzingUtils.checkForSyntacticErrors(file)
+            } catch (e: Exception) {
+                throw TestsCompilerError(e)
+            }
+        }
+
+        return analyze(
+            project,
+            files,
+            configuration,
+            packagePartProvider,
+            trace
+        ).apply {
+            try {
+                AnalyzingUtils.throwExceptionOnErrors(bindingContext)
+            } catch (e: Exception) {
+                throw TestsCompilerError(e)
+            }
+        }
+    }
+
+    @JvmStatic
+    fun analyze(file: KtFile, environment: KotlinCoreEnvironment): AnalysisResult =
+        analyze(setOf(file), environment)
+
+    @JvmStatic
+    fun analyze(files: Collection<KtFile>, environment: KotlinCoreEnvironment): AnalysisResult =
+        analyze(
+            files,
+            environment,
+            environment.configuration
+        )
+
+    @JvmStatic
+    fun analyze(
+        files: Collection<KtFile>,
+        environment: KotlinCoreEnvironment,
+        configuration: CompilerConfiguration
+    ): AnalysisResult =
+        analyze(
+            environment.project,
+            files,
+            configuration,
+            environment::createPackagePartProvider
+        )
+
+    private fun analyze(
+        project: Project,
+        files: Collection<KtFile>,
+        configuration: CompilerConfiguration,
+        packagePartProviderFactory: (GlobalSearchScope) -> PackagePartProvider,
+        trace: BindingTrace = CliBindingTrace()
+    ): AnalysisResult {
+        return TopDownAnalyzerFacadeForJVM.analyzeFilesWithJavaIntegration(
+            project, files, trace, configuration, packagePartProviderFactory
+        )
+    }
+}

plugins/compose-compiler-hosted/integration-tests/src/test/java/androidx/compose/plugins/kotlin/KtxCrossModuleTests.kt

@@ -0,0 +1,896 @@
+/*
+ * Copyright 2019 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.compose.plugins.kotlin
+
+import android.widget.TextView
+import androidx.compose.Composer
+import com.intellij.openapi.util.io.FileUtil
+import org.jetbrains.kotlin.backend.common.output.OutputFile
+import org.jetbrains.kotlin.cli.jvm.compiler.KotlinCoreEnvironment
+import org.jetbrains.kotlin.config.CompilerConfiguration
+import org.jetbrains.kotlin.config.JVMConfigurationKeys
+import org.junit.Ignore
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.robolectric.annotation.Config
+import java.io.File
+import java.net.URLClassLoader
+
+@Ignore("b/148457543")
+@RunWith(ComposeRobolectricTestRunner::class)
+@Config(
+    manifest = Config.NONE,
+    minSdk = 23,
+    maxSdk = 23
+)
+class KtxCrossModuleTests : AbstractCodegenTest() {
+
+    @Test
+    fun testCrossinlineEmittable(): Unit = ensureSetup {
+        compile(
+            "TestG", mapOf(
+                "library module" to mapOf(
+                    "x/A.kt" to """
+                    package x
+
+                    import androidx.compose.Composable
+                    import android.widget.LinearLayout
+
+                    @Composable inline fun row(crossinline children: @Composable() () -> Unit) {
+                        LinearLayout {
+                            children()
+                        }
+                    }
+                 """
+                ),
+                "Main" to mapOf(
+                    "b/B.kt" to """
+                    package b
+
+                    import androidx.compose.Composable
+                    import x.row
+
+                    @Composable fun Test() {
+                        row { }
+                    }
+                """
+                )
+            )
+        )
+    }
+
+    @Test
+    fun testConstCrossModule(): Unit = ensureSetup {
+        compile(
+            "TestG", mapOf(
+                "library module" to mapOf(
+                    "x/A.kt" to """
+                    package x
+
+                    const val MyConstant: String = ""
+                 """
+                ),
+                "Main" to mapOf(
+                    "b/B.kt" to """
+                    package b
+
+                    import x.MyConstant
+
+                    fun Test(foo: String = MyConstant) {
+                        print(foo)
+                    }
+                """
+                )
+            )
+        ) {
+            assert(it.contains("LDC \"\""))
+            assert(!it.contains("INVOKESTATIC x/AKt.getMyConstant"))
+        }
+    }
+
+    @Test
+    fun testNonCrossinlineComposable(): Unit = ensureSetup {
+        compile(
+            "TestG", mapOf(
+                "library module" to mapOf(
+                    "x/A.kt" to """
+                    package x
+
+                    import androidx.compose.Composable
+                    import androidx.compose.Pivotal
+
+                    @Composable
+                    inline fun <T> key(
+                        block: @Composable() () -> T
+                    ): T = block()
+                 """
+                ),
+                "Main" to mapOf(
+                    "b/B.kt" to """
+                    package b
+
+                    import androidx.compose.Composable
+                    import x.key
+
+                    @Composable fun Test() {
+                        key { }
+                    }
+                """
+                )
+            )
+        )
+    }
+
+    @Test
+    fun testNonCrossinlineComposableNoGenerics(): Unit = ensureSetup {
+        compile(
+            "TestG", mapOf(
+                "library module" to mapOf(
+                    "x/A.kt" to """
+                    package x
+
+                    import androidx.compose.Composable
+                    import androidx.compose.Pivotal
+
+                    @Composable
+                    inline fun key(
+                        @Suppress("UNUSED_PARAMETER")
+                        @Pivotal
+                        v1: Int,
+                        block: @Composable() () -> Int
+                    ): Int = block()
+                 """
+                ),
+                "Main" to mapOf(
+                    "b/B.kt" to """
+                    package b
+
+                    import androidx.compose.Composable
+                    import x.key
+
+                    @Composable fun Test() {
+                        key(123) { 456 }
+                    }
+                """
+                )
+            )
+        )
+    }
+
+    @Test
+    fun testRemappedTypes(): Unit = ensureSetup {
+        compile(
+            "TestG", mapOf(
+                "library module" to mapOf(
+                    "x/A.kt" to """
+                    package x
+
+                    class A {
+                        fun makeA(): A { return A() }
+                        fun makeB(): B { return B() }
+                        class B() {
+                        }
+                    }
+                 """
+                ),
+                "Main" to mapOf(
+                    "b/B.kt" to """
+                    package b
+
+                    import x.A
+
+                    class C {
+                        fun useAB() {
+                            val a = A()
+                            a.makeA()
+                            a.makeB()
+                            val b = A.B()
+                        }
+                    }
+                """
+                )
+            )
+        )
+    }
+
+    @Test
+    fun testInlineIssue(): Unit = ensureSetup {
+        compile(
+            "TestG", mapOf(
+                "library module" to mapOf(
+                    "x/C.kt" to """
+                    fun ghi() {
+                        abc {}
+                    }
+                    """,
+                    "x/A.kt" to """
+                    inline fun abc(fn: () -> Unit) {
+                        fn()
+                    }
+                    """,
+                    "x/B.kt" to """
+                    fun def() {
+                        abc {}
+                    }
+                    """
+                ),
+                "Main" to mapOf(
+                    "b/B.kt" to """
+                """
+                )
+            )
+        )
+    }
+
+    @Test
+    fun testInlineComposableProperty(): Unit = ensureSetup {
+        compile(
+            "TestG", mapOf(
+                "library module" to mapOf(
+                    "x/A.kt" to """
+                    package x
+
+                    import androidx.compose.Composable
+
+                    class Foo {
+                      @Composable val value: Int get() = 123
+                    }
+                 """
+                ),
+                "Main" to mapOf(
+                    "b/B.kt" to """
+                    package b
+
+                    import androidx.compose.Composable
+                    import x.Foo
+
+                    val foo = Foo()
+
+                    @Composable fun Test() {
+                        print(foo.value)
+                    }
+                """
+                )
+            )
+        )
+    }
+
+    @Test
+    fun testNestedInlineIssue(): Unit = ensureSetup {
+        compile(
+            "TestG", mapOf(
+                "library module" to mapOf(
+                    "x/C.kt" to """
+                    fun ghi() {
+                        abc {
+                            abc {
+
+                            }
+                        }
+                    }
+                    """,
+                    "x/A.kt" to """
+                    inline fun abc(fn: () -> Unit) {
+                        fn()
+                    }
+                    """,
+                    "x/B.kt" to """
+                    fun def() {
+                        abc {
+                            abc {
+
+                            }
+                        }
+                    }
+                    """
+                ),
+                "Main" to mapOf(
+                    "b/B.kt" to """
+                """
+                )
+            )
+        )
+    }
+
+    @Test
+    fun testComposerIntrinsicInline(): Unit = ensureSetup {
+        compile(
+            "TestG", mapOf(
+                "library module" to mapOf(
+                    "x/C.kt" to """
+                    import androidx.compose.Composable
+
+                    @Composable
+                    fun ghi() {
+                        val x = abc()
+                        print(x)
+                    }
+                    """,
+                    "x/A.kt" to """
+                    import androidx.compose.Composable
+                    import androidx.compose.currentComposer
+
+                    @Composable
+                    inline fun abc(): Any {
+                        return currentComposer
+                    }
+                    """,
+                    "x/B.kt" to """
+                    import androidx.compose.Composable
+
+                    @Composable
+                    fun def() {
+                        val x = abc()
+                        print(x)
+                    }
+                    """
+                ),
+                "Main" to mapOf(
+                    "b/B.kt" to """
+                """
+                )
+            )
+        )
+    }
+
+    @Test
+    fun testComposableOrderIssue(): Unit = ensureSetup {
+        compile(
+            "TestG", mapOf(
+                "library module" to mapOf(
+                    "C.kt" to """
+                    import androidx.compose.*
+
+                    @Composable
+                    fun b() {
+                        a()
+                    }
+                    """,
+                    "A.kt" to """
+                    import androidx.compose.*
+
+                    @Composable
+                    fun a() {
+
+                    }
+                    """,
+                    "B.kt" to """
+                    import androidx.compose.*
+
+                    @Composable
+                    fun c() {
+                        a()
+                    }
+
+                    """
+                ),
+                "Main" to mapOf(
+                    "b/B.kt" to """
+                """
+                )
+            )
+        )
+    }
+
+    @Test
+    fun testSimpleXModuleCall(): Unit = ensureSetup {
+        compile(
+            "TestG", mapOf(
+                "library module" to mapOf(
+                    "a/A.kt" to """
+                    package a
+
+                    import androidx.compose.*
+
+                    @Composable
+                    fun FromA() {}
+                 """
+                ),
+                "Main" to mapOf(
+                    "b/B.kt" to """
+                    package b
+
+                    import a.FromA
+                    import androidx.compose.*
+
+                    @Composable
+                    fun FromB() {
+                        FromA()
+                    }
+                """
+                )
+            )
+        )
+    }
+
+    @Test
+    fun testJvmFieldIssue(): Unit = ensureSetup {
+        compile(
+            "TestG", mapOf(
+                "library module" to mapOf(
+                    "x/C.kt" to """
+                    fun Test2() {
+                      bar = 10
+                      print(bar)
+                    }
+                    """,
+                    "x/A.kt" to """
+                      @JvmField var bar: Int = 0
+                    """,
+                    "x/B.kt" to """
+                    fun Test() {
+                      bar = 10
+                      print(bar)
+                    }
+                    """
+                ),
+                "Main" to mapOf(
+                    "b/B.kt" to """
+                """
+                )
+            )
+        )
+    }
+
+    @Test
+    fun testInstanceXModuleCall(): Unit = ensureSetup {
+        compile(
+            "TestH", mapOf(
+                "library module" to mapOf(
+                    "a/Foo.kt" to """
+                    package a
+
+                    import androidx.compose.*
+
+                    class Foo {
+                        @Composable
+                        fun FromA() {}
+                    }
+                 """
+                ),
+                "Main" to mapOf(
+                    "B.kt" to """
+                    import a.Foo
+                    import androidx.compose.*
+
+                    @Composable
+                    fun FromB() {
+                        Foo().FromA()
+                    }
+                """
+                )
+            )
+        )
+    }
+
+    @Test
+    fun testXModuleProperty(): Unit = ensureSetup {
+        compile(
+            "TestI", mapOf(
+                "library module" to mapOf(
+                    "a/Foo.kt" to """
+                    package a
+
+                    import androidx.compose.*
+
+                    @Composable val foo: Int get() { return 123 }
+                 """
+                ),
+                "Main" to mapOf(
+                    "B.kt" to """
+                    import a.foo
+                    import androidx.compose.*
+
+                    @Composable
+                    fun FromB() {
+                        foo
+                    }
+                """
+                )
+            )
+        )
+    }
+
+    @Test
+    fun testXModuleInterface(): Unit = ensureSetup {
+        compile(
+            "TestJ", mapOf(
+                "library module" to mapOf(
+                    "a/Foo.kt" to """
+                    package a
+
+                    import androidx.compose.*
+
+                    interface Foo {
+                        @Composable fun foo()
+                    }
+                 """
+                ),
+                "Main" to mapOf(
+                    "B.kt" to """
+                    import a.Foo
+                    import androidx.compose.*
+
+                    class B : Foo {
+                        @Composable override fun foo() {}
+                    }
+
+                    @Composable fun Example(inst: Foo) {
+                        B().foo()
+                        inst.foo()
+                    }
+                """
+                )
+            )
+        )
+    }
+
+    @Test
+    fun testXModuleCtorComposableParam(): Unit = ensureSetup {
+        compile(
+            "TestJ", mapOf(
+                "library module" to mapOf(
+                    "a/Foo.kt" to """
+                    package a
+
+                    import androidx.compose.*
+
+                    class Foo(val bar: @Composable() () -> Unit)
+                 """
+                ),
+                "Main" to mapOf(
+                    "B.kt" to """
+                    import a.Foo
+                    import androidx.compose.*
+
+                    @Composable fun Example(bar: @Composable() () -> Unit) {
+                        val foo = Foo(bar)
+                    }
+                """
+                )
+            )
+        )
+    }
+
+    @Test
+    fun testCrossModule_SimpleComposition(): Unit = ensureSetup {
+        val tvId = 29
+
+        compose(
+            "TestF", mapOf(
+                "library module" to mapOf(
+                    "my/test/lib/InternalComp.kt" to """
+                    package my.test.lib
+
+                    import androidx.compose.*
+
+                    @Composable fun InternalComp(block: @Composable() () -> Unit) {
+                        block()
+                    }
+                 """
+                ),
+                "Main" to mapOf(
+                    "my/test/app/Main.kt" to """
+                   package my.test.app
+
+                   import android.widget.*
+                   import androidx.compose.*
+                   import my.test.lib.*
+
+                   var bar = 0
+                   var doRecompose: () -> Unit = {}
+
+                   class TestF {
+                       @Composable
+                       fun compose() {
+                         Recompose { recompose ->
+                           doRecompose = recompose
+                           Foo(bar)
+                         }
+                       }
+
+                       fun advance() {
+                         bar++
+                         doRecompose()
+                       }
+                   }
+
+                   @Composable
+                   fun Foo(bar: Int) {
+                     InternalComp {
+                       TextView(text="${'$'}bar", id=$tvId)
+                     }
+                   }
+                """
+                )
+            )
+        ).then { activity ->
+            val tv = activity.findViewById(tvId) as TextView
+            assertEquals("0", tv.text)
+        }.then { activity ->
+            val tv = activity.findViewById(tvId) as TextView
+            assertEquals("1", tv.text)
+        }
+    }
+
+    @Test
+    fun testCrossModule_ComponentFunction(): Unit = ensureSetup {
+        val tvName = 101
+        val tvAge = 102
+
+        compose(
+            "TestF", mapOf(
+                "library KTX module" to mapOf(
+                    "my/test2/lib/ktx/ComponentFunction.kt" to """
+                       package my.test2.ktx
+
+                       import android.widget.*
+                       import androidx.compose.*
+
+                       @Composable
+                       fun ComponentFunction(name: String, age: Int) {
+                         LinearLayout {
+                           TextView(text=name, id=$tvName)
+                           TextView(text="${'$'}age", id=$tvAge)
+                         }
+                       }
+                 """
+                ),
+                "Main" to mapOf(
+                    "my/test2/app/Test.kt" to """
+                       package my.test2.app
+
+                       import android.widget.*
+                       import androidx.compose.*
+                       import my.test2.ktx.*
+
+                       var age = $PRESIDENT_AGE_1
+                       var name = "$PRESIDENT_NAME_1"
+                       var doRecompose: () -> Unit = {}
+
+                       class TestF {
+                           @Composable
+                           fun compose() {
+                             Recompose { recompose ->
+                               doRecompose = recompose
+                               Foo(name=name, age=age)
+                             }
+                           }
+
+                           fun advance() {
+                             age = $PRESIDENT_AGE_16
+                             name = "$PRESIDENT_NAME_16"
+                             doRecompose()
+                           }
+                       }
+
+                       @Composable
+                       fun Foo(name: String, age: Int) {
+                         ComponentFunction(name, age)
+                       }
+                    """
+                )
+            )
+        ).then { activity ->
+            val name = activity.findViewById(tvName) as TextView
+            assertEquals(PRESIDENT_NAME_1, name.text)
+            val age = activity.findViewById(tvAge) as TextView
+            assertEquals("$PRESIDENT_AGE_1", age.text)
+        }.then { activity ->
+            val name = activity.findViewById(tvName) as TextView
+            assertEquals(PRESIDENT_NAME_16, name.text)
+            val age = activity.findViewById(tvAge) as TextView
+            assertEquals("$PRESIDENT_AGE_16", age.text)
+        }
+    }
+
+    @Test
+    fun testCrossModule_ObjectFunction(): Unit = ensureSetup {
+        val tvName = 101
+        val tvAge = 102
+
+        compose(
+            "TestF", mapOf(
+                "library KTX module" to mapOf(
+                    "my/test2/lib/ktx/ObjectFunction.kt" to """
+                       package my.test2.ktx
+
+                       import android.widget.*
+                       import androidx.compose.*
+
+                       object Container {
+                           @Composable
+                           fun ComponentFunction(name: String, age: Int) {
+                             LinearLayout {
+                               TextView(text=name, id=$tvName)
+                               TextView(text="${'$'}age", id=$tvAge)
+                             }
+                           }
+                       }
+                 """
+                ),
+                "Main" to mapOf(
+                    "my/test2/app/Test.kt" to """
+                       package my.test2.app
+
+                       import android.widget.*
+                       import androidx.compose.*
+                       import my.test2.ktx.*
+
+                       var age = $PRESIDENT_AGE_1
+                       var name = "$PRESIDENT_NAME_1"
+                       var doRecompose: () -> Unit = {}
+
+                       class TestF {
+                           @Composable
+                           fun compose() {
+                             Recompose { recompose ->
+                               doRecompose = recompose
+                               Foo(name, age)
+                             }
+                           }
+
+                           fun advance() {
+                             age = $PRESIDENT_AGE_16
+                             name = "$PRESIDENT_NAME_16"
+                             doRecompose()
+                           }
+                       }
+
+                       @Composable
+                       fun Foo(name: String, age: Int) {
+                         Container.ComponentFunction(name, age)
+                       }
+                    """
+                )
+            )
+        ).then { activity ->
+            val name = activity.findViewById(tvName) as TextView
+            assertEquals(PRESIDENT_NAME_1, name.text)
+            val age = activity.findViewById(tvAge) as TextView
+            assertEquals("$PRESIDENT_AGE_1", age.text)
+        }.then { activity ->
+            val name = activity.findViewById(tvName) as TextView
+            assertEquals(PRESIDENT_NAME_16, name.text)
+            val age = activity.findViewById(tvAge) as TextView
+            assertEquals("$PRESIDENT_AGE_16", age.text)
+        }
+    }
+
+    fun compile(
+        mainClassName: String,
+        modules: Map<String, Map<String, String>>,
+        dumpClasses: Boolean = false,
+        validate: ((String) -> Unit)? = null
+    ): List<OutputFile> {
+        val libraryClasses = (modules.filter { it.key != "Main" }.map {
+            // Setup for compile
+            this.classFileFactory = null
+            this.myEnvironment = null
+            setUp(it.key.contains("--ktx=false"))
+
+            classLoader(it.value, dumpClasses).allGeneratedFiles.also {
+                // Write the files to the class directory so they can be used by the next module
+                // and the application
+                it.writeToDir(classesDirectory)
+            }
+        } + emptyList()).reduce { acc, mutableList -> acc + mutableList }
+
+        // Setup for compile
+        this.classFileFactory = null
+        this.myEnvironment = null
+        setUp()
+
+        // compile the next one
+        val appClasses = classLoader(modules["Main"]
+            ?: error("No Main module specified"), dumpClasses).allGeneratedFiles
+
+        // Load the files looking for mainClassName
+        val outputFiles = (libraryClasses + appClasses).filter {
+            it.relativePath.endsWith(".class")
+        }
+
+        if (validate != null) {
+            validate(outputFiles.joinToString("\n") { it.asText().replace('$', '%') })
+        }
+
+        return outputFiles
+    }
+
+    fun compose(
+        mainClassName: String,
+        modules: Map<String, Map<String, String>>,
+        dumpClasses: Boolean = false
+    ): RobolectricComposeTester {
+        val allClasses = compile(mainClassName, modules, dumpClasses)
+        val loader = URLClassLoader(emptyArray(), this.javaClass.classLoader)
+        val instanceClass = run {
+            var instanceClass: Class<*>? = null
+            var loadedOne = false
+            for (outFile in allClasses) {
+                val bytes = outFile.asByteArray()
+                val loadedClass = loadClass(loader, null, bytes)
+                if (loadedClass.name.endsWith(mainClassName)) instanceClass = loadedClass
+                loadedOne = true
+            }
+            if (!loadedOne) error("No classes loaded")
+            instanceClass ?: error("Could not find class $mainClassName in loaded classes")
+        }
+
+        val instanceOfClass = instanceClass.newInstance()
+        val advanceMethod = instanceClass.getMethod("advance")
+        val composeMethod = instanceClass.getMethod("compose", Composer::class.java)
+
+        return composeMulti({
+            composeMethod.invoke(instanceOfClass, it)
+        }) {
+            advanceMethod.invoke(instanceOfClass)
+        }
+    }
+
+    fun setUp(disable: Boolean = false) {
+        if (disable) {
+            this.disableIrAndKtx = true
+            try {
+                setUp()
+            } finally {
+                this.disableIrAndKtx = false
+            }
+        } else {
+            setUp()
+        }
+    }
+
+    override fun setUp() {
+        if (disableIrAndKtx) {
+            super.setUp()
+        } else {
+            super.setUp()
+        }
+    }
+
+    override fun setupEnvironment(environment: KotlinCoreEnvironment) {
+        if (!disableIrAndKtx) {
+            super.setupEnvironment(environment)
+        }
+    }
+
+    private var disableIrAndKtx = false
+
+    override fun updateConfiguration(configuration: CompilerConfiguration) {
+        super.updateConfiguration(configuration)
+        if (disableIrAndKtx) {
+            configuration.put(JVMConfigurationKeys.IR, false)
+        }
+    }
+
+    private var testLocalUnique = 0
+    private var classesDirectory = tmpDir(
+        "kotlin-${testLocalUnique++}-classes"
+    )
+
+    override val additionalPaths: List<File> = listOf(classesDirectory)
+}
+
+fun OutputFile.writeToDir(directory: File) =
+    FileUtil.writeToFile(File(directory, relativePath), asByteArray())
+
+fun Collection<OutputFile>.writeToDir(directory: File) = forEach { it.writeToDir(directory) }
+
+private fun tmpDir(name: String): File {
+    return FileUtil.createTempDirectory(name, "", false).canonicalFile
+}

plugins/compose-compiler-hosted/integration-tests/src/test/java/androidx/compose/plugins/kotlin/KtxModelCodeGenTests.kt

@@ -0,0 +1,346 @@
+/*
+ * Copyright 2019 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.compose.plugins.kotlin
+
+import android.widget.TextView
+import androidx.compose.FrameManager
+import androidx.compose.Composer
+import org.junit.Before
+import org.junit.Ignore
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.robolectric.Robolectric
+import org.robolectric.RuntimeEnvironment
+import org.robolectric.annotation.Config
+
+val PRESIDENT_NAME_1 = "George Washington"
+val PRESIDENT_AGE_1 = 57
+val PRESIDENT_NAME_16 = "Abraham Lincoln"
+val PRESIDENT_AGE_16 = 52
+
+@RunWith(ComposeRobolectricTestRunner::class)
+@Config(
+    manifest = Config.NONE,
+    minSdk = 23,
+    maxSdk = 23
+)
+class KtxModelCodeGenTests : AbstractCodegenTest() {
+
+    @Before
+    fun before() {
+        val scheduler = RuntimeEnvironment.getMasterScheduler()
+        scheduler.pause()
+    }
+
+    override fun setUp() {
+        isSetup = true
+        super.setUp()
+    }
+
+    private var isSetup = false
+    private inline fun <T> ensureSetup(crossinline block: () -> T): T {
+        if (!isSetup) setUp()
+        return FrameManager.isolated { block() }
+    }
+
+    @Test
+    @Ignore("TODO(b/138720405): Investigate synchronisation issues in tests")
+    fun testCGModelView_PersonModel(): Unit = ensureSetup {
+        val tvNameId = 384
+        val tvAgeId = 385
+
+        var name = PRESIDENT_NAME_1
+        var age = PRESIDENT_AGE_1
+        compose(
+            """
+            @Model
+            class Person4(var name: String, var age: Int)
+
+            @Composable
+            fun PersonView4(person: Person4) {
+              Observe {
+                TextView(text=person.name, id=$tvNameId)
+                TextView(text=person.age.toString(), id=$tvAgeId)
+              }
+            }
+
+            val president = Person4("$PRESIDENT_NAME_1", $PRESIDENT_AGE_1)
+            """, { mapOf("name" to name, "age" to age) }, """
+               president.name = name
+               president.age = age
+            """, """
+                PersonView4(person=president)
+            """).then { activity ->
+            val tvName = activity.findViewById(tvNameId) as TextView
+            val tvAge = activity.findViewById(tvAgeId) as TextView
+            assertEquals(PRESIDENT_NAME_1, tvName.text)
+            assertEquals(PRESIDENT_AGE_1.toString(), tvAge.text)
+
+            name = PRESIDENT_NAME_16
+            age = PRESIDENT_AGE_16
+        }.then { activity ->
+            val tvName = activity.findViewById(tvNameId) as TextView
+            val tvAge = activity.findViewById(tvAgeId) as TextView
+            assertEquals(PRESIDENT_NAME_16, tvName.text)
+            assertEquals(PRESIDENT_AGE_16.toString(), tvAge.text)
+        }
+    }
+
+    @Test // b/120843442
+    fun testCGModelView_ObjectModel(): Unit = ensureSetup {
+        val tvNameId = 384
+        val tvAgeId = 385
+
+        var name = PRESIDENT_NAME_1
+        var age = PRESIDENT_AGE_1
+        compose(
+            """
+            @Model
+            object president {
+                var name: String = "$PRESIDENT_NAME_1"
+                var age: Int = $PRESIDENT_AGE_1
+            }
+
+            @Composable
+            fun PresidentView() {
+              Observe {
+                TextView(text=president.name, id=$tvNameId)
+                TextView(text=president.age.toString(), id=$tvAgeId)
+              }
+            }
+
+            """, { mapOf("name" to name, "age" to age) }, """
+               president.name = name
+               president.age = age
+            """, """
+                PresidentView()
+            """).then { activity ->
+            val tvName = activity.findViewById(tvNameId) as TextView
+            val tvAge = activity.findViewById(tvAgeId) as TextView
+            assertEquals(PRESIDENT_NAME_1, tvName.text)
+            assertEquals(PRESIDENT_AGE_1.toString(), tvAge.text)
+
+            name = PRESIDENT_NAME_16
+            age = PRESIDENT_AGE_16
+        }.then { activity ->
+            val tvName = activity.findViewById(tvNameId) as TextView
+            val tvAge = activity.findViewById(tvAgeId) as TextView
+            assertEquals(PRESIDENT_NAME_16, tvName.text)
+            assertEquals(PRESIDENT_AGE_16.toString(), tvAge.text)
+        }
+    }
+
+    @Test // b/120836313
+    fun testCGModelView_DataModel(): Unit = ensureSetup {
+        val tvNameId = 384
+        val tvAgeId = 385
+
+        var name = PRESIDENT_NAME_1
+        var age = PRESIDENT_AGE_1
+        compose(
+            """
+            @Model
+            data class PersonB(var name: String, var age: Int)
+
+            @Composable
+            fun PersonViewB(person: PersonB) {
+              Observe {
+                TextView(text=person.name, id=$tvNameId)
+                TextView(text=person.age.toString(), id=$tvAgeId)
+              }
+            }
+
+            val president = PersonB("$PRESIDENT_NAME_1", $PRESIDENT_AGE_1)
+            """, { mapOf("name" to name, "age" to age) }, """
+               president.name = name
+               president.age = age
+            """, """
+                PersonViewB(person=president)
+            """).then { activity ->
+            val tvName = activity.findViewById(tvNameId) as TextView
+            val tvAge = activity.findViewById(tvAgeId) as TextView
+            assertEquals(PRESIDENT_NAME_1, tvName.text)
+            assertEquals(PRESIDENT_AGE_1.toString(), tvAge.text)
+
+            name = PRESIDENT_NAME_16
+            age = PRESIDENT_AGE_16
+        }.then { activity ->
+            val tvName = activity.findViewById(tvNameId) as TextView
+            val tvAge = activity.findViewById(tvAgeId) as TextView
+            assertEquals(PRESIDENT_NAME_16, tvName.text)
+            assertEquals(PRESIDENT_AGE_16.toString(), tvAge.text)
+        }
+    }
+
+    @Test // b/120843442
+    fun testCGModelView_ZeroFrame(): Unit = ensureSetup {
+        val tvNameId = 384
+        val tvAgeId = 385
+
+        var name = PRESIDENT_NAME_1
+        var age = PRESIDENT_AGE_1
+        compose(
+            """
+            @Model
+            class PersonC(var name: String, var age: Int)
+
+            @Composable
+            fun PersonViewC(person: PersonC) {
+              Observe {
+                TextView(text=person.name, id=$tvNameId)
+                TextView(text=person.age.toString(), id=$tvAgeId)
+              }
+            }
+
+            val president = FrameManager.unframed { PersonC("$PRESIDENT_NAME_1", $PRESIDENT_AGE_1) }
+            """, { mapOf("name" to name, "age" to age) }, """
+               president.name = name
+               president.age = age
+            """, """
+                PersonViewC(person=president)
+            """).then { activity ->
+            val tvName = activity.findViewById(tvNameId) as TextView
+            val tvAge = activity.findViewById(tvAgeId) as TextView
+            assertEquals(name, tvName.text)
+            assertEquals(age.toString(), tvAge.text)
+
+            name = PRESIDENT_NAME_16
+            age = PRESIDENT_AGE_16
+        }.then { activity ->
+            val tvName = activity.findViewById(tvNameId) as TextView
+            val tvAge = activity.findViewById(tvAgeId) as TextView
+            assertEquals(PRESIDENT_NAME_16, tvName.text)
+            assertEquals(PRESIDENT_AGE_16.toString(), tvAge.text)
+        }
+    }
+
+    @Test // b/120843442
+    fun testCGModelView_ZeroFrame_Modification(): Unit = ensureSetup {
+        val tvNameId = 384
+        val tvAgeId = 385
+
+        var name = PRESIDENT_NAME_1
+        var age = PRESIDENT_AGE_1
+        compose(
+            """
+            @Model
+            class PersonD(var name: String, var age: Int)
+
+            @Composable
+            fun PersonViewD(person: PersonD) {
+              Observe {
+                TextView(text=person.name, id=$tvNameId)
+                TextView(text=person.age.toString(), id=$tvAgeId)
+              }
+            }
+
+            val president = FrameManager.framed { PersonD("$PRESIDENT_NAME_1", $PRESIDENT_AGE_1).apply { age = $PRESIDENT_AGE_1 } }
+            """, { mapOf("name" to name, "age" to age) }, """
+               president.name = name
+               president.age = age
+            """, """
+                PersonViewD(person=president)
+            """).then { activity ->
+            val tvName = activity.findViewById(tvNameId) as TextView
+            val tvAge = activity.findViewById(tvAgeId) as TextView
+            assertEquals(PRESIDENT_NAME_1, tvName.text)
+            assertEquals(PRESIDENT_AGE_1.toString(), tvAge.text)
+
+            name = PRESIDENT_NAME_16
+            age = PRESIDENT_AGE_16
+        }.then { activity ->
+            val tvName = activity.findViewById(tvNameId) as TextView
+            val tvAge = activity.findViewById(tvAgeId) as TextView
+            assertEquals(PRESIDENT_NAME_16, tvName.text)
+            assertEquals(PRESIDENT_AGE_16.toString(), tvAge.text)
+        }
+    }
+
+    fun compose(
+        prefix: String,
+        valuesFactory: () -> Map<String, Any>,
+        advance: String,
+        composition: String,
+        dumpClasses: Boolean = false
+    ): RobolectricComposeTester {
+        val className = "Test_${uniqueNumber++}"
+        val fileName = "$className.kt"
+
+        val candidateValues = valuesFactory()
+
+        @Suppress("NO_REFLECTION_IN_CLASS_PATH")
+        val parameterList = candidateValues.map {
+            "${it.key}: ${it.value::class.qualifiedName}"
+        }.joinToString()
+        val parameterTypes = candidateValues.map {
+            it.value::class.javaPrimitiveType ?: it.value::class.javaObjectType
+        }.toTypedArray()
+
+        val compiledClasses = classLoader("""
+           import android.content.Context
+           import android.widget.*
+           import androidx.compose.*
+
+           $prefix
+
+           class $className {
+
+             @Composable
+             fun compose() {
+               $composition
+             }
+
+             fun advance($parameterList) {
+               $advance
+             }
+           }
+        """, fileName, dumpClasses)
+
+        val allClassFiles = compiledClasses.allGeneratedFiles.filter {
+            it.relativePath.endsWith(".class")
+        }
+
+        val instanceClass = run {
+            var instanceClass: Class<*>? = null
+            var loadedOne = false
+            for (outFile in allClassFiles) {
+                val bytes = outFile.asByteArray()
+                val loadedClass = loadClass(
+                    this.javaClass.classLoader!!,
+                    null,
+                    bytes
+                )
+                if (loadedClass.name == className) instanceClass = loadedClass
+                loadedOne = true
+            }
+            if (!loadedOne) error("No classes loaded")
+            instanceClass ?: error("Could not find class $className in loaded classes")
+        }
+
+        val instanceOfClass = instanceClass.newInstance()
+        val advanceMethod = instanceClass.getMethod("advance", *parameterTypes)
+        val composeMethod = instanceClass.getMethod("compose", Composer::class.java)
+
+        return composeMulti({
+            composeMethod.invoke(instanceOfClass, it)
+        }) {
+            val values = valuesFactory()
+            val arguments = values.map { it.value }.toTypedArray()
+            advanceMethod.invoke(instanceOfClass, *arguments)
+        }
+    }
+}

plugins/compose-compiler-hosted/integration-tests/src/test/java/androidx/compose/plugins/kotlin/KtxTransformationTest.kt

@@ -0,0 +1,772 @@
+/*
+ * Copyright 2019 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.compose.plugins.kotlin
+
+import org.junit.Before
+
+class KtxTransformationTest : AbstractCodegenTest() {
+
+    fun testObserveLowering() = ensureSetup { testCompile(
+        """
+            import android.widget.Button
+            import androidx.compose.*
+            import androidx.ui.androidview.adapters.setOnClick
+
+            @Model
+            class FancyButtonData() {
+                var x = 0
+            }
+
+            @Composable
+            fun SimpleComposable() {
+                FancyButton(state=FancyButtonData())
+            }
+
+            @Composable
+            fun FancyButton(state: FancyButtonData) {
+               Button(text=("Clicked "+state.x+" times"), onClick={state.x++}, id=42)
+            }
+        """
+    ) }
+
+    fun testEmptyComposeFunction() = ensureSetup { testCompile(
+        """
+        import androidx.compose.*
+
+        class Foo {
+            @Composable
+            operator fun invoke() {}
+        }
+        """
+    ) }
+
+    fun testSingleViewCompose() = ensureSetup { testCompile(
+        """
+        import androidx.compose.*
+        import android.widget.*
+
+        class Foo {
+            @Composable
+            operator fun invoke() {
+                TextView()
+            }
+        }
+        """
+    ) }
+
+    fun testMultipleRootViewCompose() = ensureSetup { testCompile(
+        """
+        import androidx.compose.*
+        import android.widget.*
+
+        class Foo {
+            @Composable
+            operator fun invoke() {
+                TextView()
+                TextView()
+                TextView()
+            }
+        }
+        """
+    ) }
+
+    fun testNestedViewCompose() = ensureSetup { testCompile(
+        """
+        import androidx.compose.*
+        import android.widget.*
+
+        class Foo {
+            @Composable
+            operator fun invoke() {
+                LinearLayout {
+                    TextView()
+                    LinearLayout {
+                        TextView()
+                        TextView()
+                    }
+                }
+            }
+        }
+        """
+    ) }
+
+    fun testSingleComposite() = ensureSetup { testCompile(
+        """
+         import androidx.compose.*
+
+        @Composable
+        fun Bar() {}
+
+        class Foo {
+            @Composable
+            operator fun invoke() {
+                Bar()
+            }
+        }
+        """
+    ) }
+
+    fun testMultipleRootComposite() = ensureSetup { testCompile(
+        """
+         import androidx.compose.*
+
+        @Composable
+        fun Bar() {}
+
+        class Foo {
+            @Composable
+            operator fun invoke() {
+                Bar()
+                Bar()
+                Bar()
+            }
+        }
+        """
+    ) }
+
+    fun testViewAndComposites() = ensureSetup { testCompile(
+        """
+        import androidx.compose.*
+        import android.widget.*
+
+        @Composable
+        fun Bar() {}
+
+        class Foo {
+            @Composable
+            operator fun invoke() {
+                LinearLayout {
+                    Bar()
+                }
+            }
+        }
+        """
+    ) }
+
+    fun testForEach() = ensureSetup { testCompile(
+        """
+         import androidx.compose.*
+
+        @Composable
+        fun Bar() {}
+
+        class Foo {
+            @Composable
+            operator fun invoke() {
+                listOf(1, 2, 3).forEach {
+                    Bar()
+                }
+            }
+        }
+        """
+    ) }
+
+    fun testForLoop() = ensureSetup { testCompile(
+        """
+         import androidx.compose.*
+
+        @Composable
+        fun Bar() {}
+
+        class Foo {
+            @Composable
+            operator fun invoke() {
+                for (i in listOf(1, 2, 3)) {
+                    Bar()
+                }
+            }
+        }
+        """
+    ) }
+
+    fun testEarlyReturns() = ensureSetup { testCompile(
+        """
+         import androidx.compose.*
+
+        @Composable
+        fun Bar() {}
+
+        class Foo {
+            var visible: Boolean = false
+            @Composable
+            operator fun invoke() {
+                if (!visible) return
+                else "" // TODO: Remove this line when fixed upstream
+                Bar()
+            }
+        }
+        """
+    ) }
+
+    fun testConditionalRendering() = ensureSetup { testCompile(
+        """
+         import androidx.compose.*
+
+        @Composable
+        fun Bar() {}
+
+        @Composable
+        fun Bam() {}
+
+        class Foo {
+            var visible: Boolean = false
+            @Composable
+            operator fun invoke() {
+                if (!visible) {
+                    Bar()
+                } else {
+                    Bam()
+                }
+            }
+        }
+        """
+    ) }
+
+    fun testExtensionFunctions() = ensureSetup { testCompile(
+        """
+
+        import androidx.compose.*
+        import android.widget.*
+
+        fun LinearLayout.setSomeExtension(x: Int) {
+        }
+        class X {
+            @Composable
+            operator fun invoke() {
+                LinearLayout(someExtension=123)
+            }
+        }
+        """
+    ) }
+
+    fun testChildrenWithTypedParameters() = ensureSetup { testCompile(
+        """
+        import android.widget.*
+        import androidx.compose.*
+
+        @Composable fun HelperComponent(
+            children: @Composable() (title: String, rating: Int) -> Unit
+        ) {
+            children("Hello World!", 5)
+            children("Kompose is awesome!", 5)
+            children("Bitcoin!", 4)
+        }
+
+        class MainComponent {
+            var name = "World"
+            @Composable
+            operator fun invoke() {
+                HelperComponent { title: String, rating: Int ->
+                    TextView(text=(title+" ("+rating+" stars)"))
+                }
+            }
+        }
+        """
+    ) }
+
+    fun testChildrenCaptureVariables() = ensureSetup { testCompile(
+        """
+        import android.widget.*
+        import androidx.compose.*
+
+        @Composable fun HelperComponent(children: @Composable() () -> Unit) {
+        }
+
+        class MainComponent {
+            var name = "World"
+            @Composable
+            operator fun invoke() {
+                val childText = "Hello World!"
+                HelperComponent {
+                    TextView(text=childText + name)
+                }
+            }
+        }
+        """
+    ) }
+
+    fun testChildrenDeepCaptureVariables() = ensureSetup { testCompile(
+        """
+        import android.widget.*
+        import androidx.compose.*
+
+        @Composable fun A(children: @Composable() () -> Unit) {
+            children()
+        }
+
+        @Composable fun B(children: @Composable() () -> Unit) {
+            children()
+        }
+
+        class MainComponent {
+            var name = "World"
+            @Composable
+            operator fun invoke() {
+                val childText = "Hello World!"
+                A {
+                    B {
+                        println(childText + name)
+                    }
+                }
+            }
+        }
+        """
+    ) }
+
+    fun testChildrenDeepCaptureVariablesWithParameters() = ensureSetup {
+        testCompile(
+        """
+        import android.widget.*
+        import androidx.compose.*
+
+        @Composable fun A(children: @Composable() (x: String) -> Unit) {
+            children("")
+        }
+
+        @Composable fun B(children: @Composable() (y: String) -> Unit) {
+            children("")
+        }
+
+        class MainComponent {
+            var name = "World"
+            @Composable
+            operator fun invoke() {
+                val childText = "Hello World!"
+                A { x ->
+                    B { y ->
+                        println(childText + name + x + y)
+                    }
+                }
+            }
+        }
+        """
+        )
+    }
+
+    fun testChildrenOfNativeView() = ensureSetup { testCompile(
+        """
+        import android.widget.*
+        import androidx.compose.*
+
+        class MainComponent {
+            @Composable
+            operator fun invoke() {
+                LinearLayout {
+                    TextView(text="some child content2!")
+                    TextView(text="some child content!3")
+                }
+            }
+        }
+        """
+    ) }
+
+    fun testIrSpecial() = ensureSetup { testCompile(
+        """
+        import android.widget.*
+        import androidx.compose.*
+
+        @Composable fun HelperComponent(children: @Composable() () -> Unit) {}
+
+        class MainComponent {
+            @Composable
+            operator fun invoke() {
+                val x = "Hello"
+                val y = "World"
+                HelperComponent {
+                    for(i in 1..100) {
+                        TextView(text=x+y+i)
+                    }
+                }
+            }
+        }
+        """
+    ) }
+
+    fun testGenericsInnerClass() = ensureSetup { testCompile(
+        """
+        import androidx.compose.*
+
+        class A<T>(val value: T) {
+            @Composable fun Getter(x: T? = null) {
+            }
+        }
+
+        @Composable
+        fun doStuff() {
+            val a = A(123)
+
+            // a.Getter() here has a bound type argument through A
+            a.Getter(x=456)
+        }
+        """
+    ) }
+
+    fun testXGenericConstructorParams() = ensureSetup { testCompile(
+        """
+        import androidx.compose.*
+
+        @Composable fun <T> A(
+            value: T,
+            list: List<T>? = null
+        ) {
+
+        }
+
+        @Composable
+        fun doStuff() {
+            val x = 123
+
+            // we can create element with just value, no list
+            A(value=x)
+
+            // if we add a list, it can infer the type
+            A(
+                value=x,
+                list=listOf(234, x)
+            )
+        }
+        """
+    ) }
+
+    fun testSimpleNoArgsComponent() = ensureSetup { testCompile(
+        """
+        import androidx.compose.*
+
+        @Composable
+        fun Simple() {}
+
+        @Composable
+        fun run() {
+            Simple()
+        }
+        """
+    ) }
+
+    fun testDotQualifiedObjectToClass() = ensureSetup { testCompile(
+        """
+        import androidx.compose.*
+
+        object Obj {
+            @Composable
+            fun B() {}
+        }
+
+        @Composable
+        fun run() {
+            Obj.B()
+        }
+        """
+    ) }
+
+    fun testPackageQualifiedTags() = ensureSetup { testCompile(
+        """
+        import androidx.compose.*
+
+        @Composable
+        fun run() {
+            android.widget.TextView(text="bar")
+        }
+        """
+    ) }
+
+    fun testLocalLambda() = ensureSetup { testCompile(
+        """
+        import androidx.compose.*
+
+        @Composable
+        fun Simple() {}
+
+        @Composable
+        fun run() {
+            val foo = @Composable { Simple() }
+            foo()
+        }
+        """
+    ) }
+
+    fun testPropertyLambda() = ensureSetup { testCompile(
+        """
+        import androidx.compose.*
+
+        class Test(var children: @Composable () () -> Unit) {
+            @Composable
+            operator fun invoke() {
+                children()
+            }
+        }
+        """
+    ) }
+
+    fun testLambdaWithArgs() = ensureSetup { testCompile(
+        """
+        import androidx.compose.*
+
+        class Test(var children: @Composable() (x: Int) -> Unit) {
+            @Composable
+            operator fun invoke() {
+                children(x=123)
+            }
+        }
+        """
+    ) }
+
+    fun testLocalMethod() = ensureSetup { testCompile(
+        """
+        import androidx.compose.*
+
+        class Test {
+            @Composable
+            fun doStuff() {}
+            @Composable
+            operator fun invoke() {
+                doStuff()
+            }
+        }
+        """
+    ) }
+
+    fun testSimpleLambdaChildren() = ensureSetup { testCompile(
+        """
+        import androidx.compose.*
+        import android.widget.*
+        import android.content.*
+
+        @Composable fun Example(children: @Composable() () -> Unit) {
+
+        }
+
+        @Composable
+        fun run(text: String) {
+            Example {
+                println("hello ${"$"}text")
+            }
+        }
+        """
+    ) }
+
+    fun testFunctionComponentsWithChildrenSimple() = ensureSetup { testCompile(
+        """
+        import androidx.compose.*
+
+        @Composable
+        fun Example(children: @Composable() () -> Unit) {}
+
+        @Composable
+        fun run(text: String) {
+            Example {
+                println("hello ${"$"}text")
+            }
+        }
+        """
+    ) }
+
+    fun testFunctionComponentWithChildrenOneArg() = ensureSetup { testCompile(
+        """
+        import androidx.compose.*
+
+        @Composable
+        fun Example(children: @Composable() (String) -> Unit) {}
+
+        @Composable
+        fun run(text: String) {
+            Example { x ->
+                println("hello ${"$"}x")
+            }
+        }
+        """
+    ) }
+
+    fun testKtxLambdaInForLoop() = ensureSetup { testCompile(
+        """
+        import androidx.compose.*
+        import android.widget.TextView
+
+        @Composable
+        fun foo() {
+            val lambda = @Composable {  }
+            for(x in 1..5) {
+                lambda()
+                lambda()
+            }
+        }
+        """
+    ) }
+
+    fun testKtxLambdaInIfElse() = ensureSetup { testCompile(
+        """
+        import androidx.compose.*
+        import android.widget.TextView
+
+        @Composable
+        fun foo(x: Boolean) {
+            val lambda = @Composable { TextView(text="Hello World") }
+            if(true) {
+                lambda()
+                lambda()
+                lambda()
+            } else {
+                lambda()
+            }
+        }
+        """
+    ) }
+
+    fun testMultiplePivotalAttributesOdd() = ensureSetup { testCompile(
+        """
+        import androidx.compose.*
+
+        @Composable fun Foo(
+            @Pivotal a: Int,
+            @Pivotal b: Int,
+            @Pivotal c: Int,
+            @Pivotal d: Int,
+            @Pivotal e: Int
+        ) {
+
+        }
+
+        class Bar {
+            @Composable
+            operator fun invoke() {
+                Foo(
+                    a=1,
+                    b=2,
+                    c=3,
+                    d=4,
+                    e=5
+                )
+            }
+        }
+        """
+    ) }
+
+    fun testSinglePivotalAttribute() = ensureSetup { testCompile(
+        """
+        import androidx.compose.*
+
+        @Composable fun Foo(
+            @Pivotal a: Int
+        ) {
+
+        }
+
+        class Bar {
+            @Composable
+            operator fun invoke() {
+                Foo(
+                    a=1
+                )
+            }
+        }
+        """
+    ) }
+
+    fun testKtxVariableTagsProperlyCapturedAcrossKtxLambdas() = ensureSetup {
+        testCompile(
+        """
+        import androidx.compose.*
+        import androidx.ui.androidview.adapters.*
+
+        @Composable fun Foo(children: @Composable() (sub: @Composable() () -> Unit) -> Unit) {
+
+        }
+
+        @Composable fun Boo(children: @Composable() () -> Unit) {
+
+        }
+
+        class Bar {
+            @Composable
+            operator fun invoke() {
+                Foo { sub ->
+                    Boo {
+                        sub()
+                    }
+                }
+            }
+        }
+        """
+        )
+    }
+
+    fun testKtxEmittable() = ensureSetup { testCompile(
+        """
+        import androidx.compose.*
+
+        open class MockEmittable: Emittable {
+          override fun emitInsertAt(index: Int, instance: Emittable) {}
+          override fun emitRemoveAt(index: Int, count: Int) {}
+          override fun emitMove(from: Int, to: Int, count: Int) {}
+        }
+
+        class MyEmittable: MockEmittable() {
+          var a: Int = 1
+        }
+
+        class Comp {
+          @Composable
+            operator fun invoke() {
+            MyEmittable(a=2)
+          }
+        }
+        """
+    ) }
+
+    fun testKtxCompoundEmittable() = ensureSetup { testCompile(
+        """
+        import androidx.compose.*
+
+        open class MockEmittable: Emittable {
+          override fun emitInsertAt(index: Int, instance: Emittable) {}
+          override fun emitRemoveAt(index: Int, count: Int) {}
+          override fun emitMove(from: Int, to: Int, count: Int) {}
+        }
+
+        class MyEmittable: MockEmittable() {
+          var a: Int = 1
+        }
+
+        class Comp {
+          @Composable
+            operator fun invoke() {
+            MyEmittable(a=1) {
+              MyEmittable(a=2)
+              MyEmittable(a=3)
+              MyEmittable(a=4)
+              MyEmittable(a=5)
+            }
+          }
+        }
+        """
+    ) }
+
+    fun testInvocableObject() = ensureSetup { testCompile(
+        """
+        import androidx.compose.*
+
+        class Foo { }
+        @Composable
+        operator fun Foo.invoke() {  }
+
+        @Composable
+        fun test() {
+            val foo = Foo()
+            foo()
+        }
+        """
+    ) }
+}

plugins/compose-compiler-hosted/integration-tests/src/test/java/androidx/compose/plugins/kotlin/LambdaMemoizationTests.kt

@@ -0,0 +1,541 @@
+/*
+ * Copyright 2020 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.compose.plugins.kotlin
+
+import android.widget.Button
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.robolectric.annotation.Config
+
+@RunWith(ComposeRobolectricTestRunner::class)
+@Config(
+    manifest = Config.NONE,
+    minSdk = 23,
+    maxSdk = 23
+)
+class LambdaMemoizationTests : AbstractLoweringTests() {
+
+    @Test
+    fun nonCapturingEventLambda() = skipping("""
+            fun eventFired() { }
+
+            @Composable
+            fun EventHolder(event: () -> Unit, block: @Composable() () -> Unit) {
+              block()
+            }
+
+            @Composable
+            fun Example(model: String) {
+                EventHolder(event = { eventFired() }) {
+                  workToBeRepeated()
+                  ValidateModel(text = model)
+                }
+            }
+        """)
+
+    @Test
+    fun lambdaInClassInitializer() = skipping("""
+            @Composable
+            fun EventHolder(event: () -> Unit) {
+              workToBeRepeated()
+            }
+
+            @Composable
+            fun Example(model: String) {
+              class Nested {
+                // Should not memoize the initializer
+                val lambda: () -> Unit = { }
+              }
+              val n = Nested()
+              ValidateModel(model)
+              EventHolder(event = n.lambda)
+            }
+        """)
+
+    @Test
+    fun methodReferenceEvent() = skipping("""
+            fun eventFired() { }
+
+            @Composable
+            fun EventHolder(event: () -> Unit) {}
+
+            @Composable
+            fun Example(model: String) {
+              workToBeRepeated()
+              EventHolder(event = ::eventFired)
+            }
+        """)
+
+    @Test
+    fun methodReferenceOnValue() = skipping("""
+        fun eventFired(value: String) { }
+
+        @Composable
+        fun ValidateEvent(expected: String, event: () -> String) {
+          val value = event()
+          require(expected == value) {
+            "Expected '${'$'}expected', received '${'$'}value'"
+          }
+        }
+
+        @Composable
+        fun Test(model: String, unchanged: String) {
+          ValidateEvent(unchanged, unchanged::toString)
+          ValidateEvent(model, model::toString)
+        }
+
+        @Composable
+        fun Example(model: String) {
+          Test(model, "unchanged")
+        }
+    """)
+
+    @Test
+    fun extensionMethodReferenceOnValue() = skipping("""
+        fun eventFired(value: String) { }
+
+        fun String.self() = this
+
+        @Composable
+        fun ValidateEvent(expected: String, event: () -> String) {
+          val value = event()
+          require(expected == value) {
+            "Expected '${'$'}expected', received '${'$'}value'"
+          }
+        }
+
+        @Composable
+        fun Test(model: String, unchanged: String) {
+          ValidateEvent(unchanged, unchanged::self)
+          ValidateEvent(model, model::self)
+        }
+
+        @Composable
+        fun Example(model: String) {
+          Test(model, "unchanged")
+        }
+    """)
+
+    @Test
+    fun doNotMemoizeCallsToInlines() = skipping("""
+            fun eventFired(data: String) { }
+
+            @Composable
+            fun EventHolder(event: () -> Unit, block: @Composable() () -> Unit) {
+               workToBeRepeated()
+               block()
+            }
+
+            @Composable
+            inline fun <T, V1> inlined(value: V1, block: () -> T) = block()
+
+            @Composable
+            fun Example(model: String) {
+              val e1 = inlined(model) { { eventFired(model) } }
+              EventHolder(event = e1) {
+                workToBeRepeated()
+                ValidateModel(model)
+              }
+              val e2 = remember(model) { { eventFired(model) } }
+              EventHolder(event = e2) {
+                workToBeRepeated()
+                ValidateModel(model)
+              }
+            }
+        """)
+
+    @Test
+    fun captureParameterDirectEventLambda() = skipping("""
+            fun eventFired(data: String) { }
+
+            @Composable
+            fun EventHolder(event: () -> Unit, block: @Composable() () -> Unit) {
+               workToBeRepeated()
+               block()
+            }
+
+            @Composable
+            fun Example(model: String) {
+                EventHolder(event = { eventFired(model) }) {
+                  workToBeRepeated()
+                  ValidateModel(text = model)
+                }
+            }
+        """)
+
+    @Test
+    fun shouldNotRememberDirectLambdaParameter() = skipping("""
+        fun eventFired(data: String) {
+          println("Validating ${'$'}data")
+          validateModel(data)
+        }
+
+        @Composable
+        fun EventHolder(event: () -> Unit) {
+          workToBeRepeated()
+          event()
+        }
+
+        @Composable
+        fun EventWrapper(event: () -> Unit) {
+          EventHolder(event)
+        }
+
+        @Composable
+        fun Example(model: String) {
+          EventWrapper(event = { eventFired(model) })
+        }
+    """)
+
+    @Test
+    fun narrowCaptureValidation() = skipping("""
+        fun eventFired(data: String) { }
+
+        @Composable
+        fun ExpectUnmodified(event: () -> Unit) {
+          workToBeAvoided()
+        }
+
+        @Composable
+        fun ExpectModified(event: () -> Unit) {
+          workToBeRepeated()
+        }
+
+        @Composable
+        fun Example(model: String) {
+          // Unmodified
+          val unmodified = model.substring(0, 4)
+          val modified = model + " abc"
+
+          ExpectUnmodified(event = { eventFired(unmodified) })
+          ExpectModified(event = { eventFired(modified) })
+          ExpectModified(event = { eventFired(model) })
+        }
+    """)
+
+    @Test
+    fun captureInANestedScope() = skipping("""
+        fun eventFired(data: String) { }
+
+        @Composable
+        fun ExpectUnmodified(event: () -> Unit) {
+          workToBeAvoided()
+        }
+
+        @Composable
+        fun ExpectModified(event: () -> Unit) {
+          workToBeRepeated()
+        }
+
+        @Composable
+        fun Wrapped(block: @Composable() () -> Unit) {
+          block()
+        }
+
+        @Composable
+        fun Example(model: String) {
+          val unmodified = model.substring(0, 4)
+          val modified = model + " abc"
+
+          Wrapped {
+            ExpectUnmodified(event = { eventFired(unmodified) })
+            ExpectModified(event = { eventFired(modified) })
+            ExpectModified(event = { eventFired(model) })
+          }
+          Wrapped {
+            Wrapped {
+              ExpectUnmodified(event = { eventFired(unmodified) })
+              ExpectModified(event = { eventFired(modified) })
+              ExpectModified(event = { eventFired(model) })
+            }
+          }
+        }
+    """)
+
+    @Test
+    fun twoCaptures() = skipping("""
+        fun eventFired(data: String) { }
+
+        @Composable
+        fun ExpectUnmodified(event: () -> Unit) {
+          workToBeAvoided()
+        }
+
+        @Composable
+        fun ExpectModified(event: () -> Unit) {
+          workToBeRepeated()
+        }
+
+        @Composable
+        fun Wrapped(block: @Composable() () -> Unit) {
+          block()
+        }
+
+        @Composable
+        fun Example(model: String) {
+          val unmodified1 = model.substring(0, 4)
+          val unmodified2 = model.substring(0, 5)
+          val modified1 = model + " abc"
+          val modified2 = model + " abcd"
+
+          ExpectUnmodified(event = { eventFired(unmodified1 + unmodified2) })
+          ExpectModified(event = { eventFired(modified1 + unmodified1) })
+          ExpectModified(event = { eventFired(unmodified2 + modified2) })
+          ExpectModified(event = { eventFired(modified1 + modified2) })
+        }
+    """)
+
+    @Test
+    fun threeCaptures() = skipping("""
+        fun eventFired(data: String) { }
+
+        @Composable
+        fun ExpectUnmodified(event: () -> Unit) {
+          workToBeAvoided()
+        }
+
+        @Composable
+        fun ExpectModified(event: () -> Unit) {
+          workToBeRepeated()
+        }
+
+        @Composable
+        fun Example(model: String) {
+          val unmodified1 = model.substring(0, 4)
+          val unmodified2 = model.substring(0, 5)
+          val unmodified3 = model.substring(0, 6)
+          val modified1 = model + " abc"
+          val modified2 = model + " abcd"
+          val modified3 = model + " abcde"
+
+          ExpectUnmodified(event = { eventFired(unmodified1 + unmodified2 + unmodified3) })
+          ExpectModified(event = { eventFired(unmodified1 + unmodified2 + modified3) })
+          ExpectModified(event = { eventFired(unmodified1 + modified2 + unmodified3) })
+          ExpectModified(event = { eventFired(unmodified1 + modified2 + modified3) })
+          ExpectModified(event = { eventFired(modified1 + unmodified2 + unmodified3) })
+          ExpectModified(event = { eventFired(modified1 + unmodified2 + modified3) })
+          ExpectModified(event = { eventFired(modified1 + modified2 + unmodified3) })
+          ExpectModified(event = { eventFired(modified1 + modified2 + modified3) })
+        }
+    """)
+
+    @Test
+    fun fiveCaptures() = skipping("""
+        fun eventFired(data: String) { }
+
+        @Composable
+        fun ExpectUnmodified(event: () -> Unit) {
+          workToBeAvoided()
+        }
+
+        @Composable
+        fun ExpectModified(event: () -> Unit) {
+          workToBeRepeated()
+        }
+
+        @Composable
+        fun Example(model: String) {
+          val modified = model
+          val unmodified1 = model.substring(0, 1)
+          val unmodified2 = model.substring(0, 2)
+          val unmodified3 = model.substring(0, 3)
+          val unmodified4 = model.substring(0, 4)
+          val unmodified5 = model.substring(0, 5)
+
+          ExpectUnmodified(event = { eventFired(
+              unmodified1 + unmodified2 + unmodified3 + unmodified4 + unmodified1
+            ) })
+
+          ExpectModified(event = { eventFired(
+              unmodified1 + unmodified2 + unmodified3 + unmodified4 + unmodified1 + modified
+            ) })
+        }
+    """)
+
+    @Test
+    fun doNotMemoizeNonStableCaptures() = skipping("""
+        val unmodifiedUnstable = Any()
+        val unmodifiedString = "unmodified"
+
+        fun eventFired(data: String) { }
+
+        @Composable
+        fun ExpectUnmodified(event: () -> Unit) {
+          workToBeAvoided()
+        }
+
+        @Composable
+        fun ExpectModified(event: () -> Unit) {
+          workToBeRepeated()
+        }
+
+        @Composable
+        fun NonStable(model: String, nonStable: Any, unmodified: String) {
+          workToBeRepeated()
+          ExpectModified(event = { eventFired(nonStable.toString()) })
+          ExpectModified(event = { eventFired(model) })
+          ExpectUnmodified(event = { eventFired(unmodified) })
+        }
+
+        @Composable
+        fun Example(model: String) {
+          NonStable(model, unmodifiedUnstable, unmodifiedString)
+        }
+    """)
+
+    @Test
+    fun doNotMemoizeVarCapures() = skipping("""
+        fun eventFired(data: Int) { }
+
+        @Composable
+        fun ExpectUnmodified(event: () -> Unit) {
+          workToBeAvoided()
+        }
+
+        @Composable
+        fun ExpectModified(event: () -> Unit) {
+          workToBeRepeated()
+        }
+
+        @Composable
+        fun Wrap(block: @Composable() () -> Unit) {
+          block()
+        }
+
+        @Composable
+        fun Test(model: String, b: Int) {
+          var a = 1
+          var c = false
+          ExpectModified(event = { a++ })
+          ExpectModified(event = { eventFired(a) })
+          ExpectModified(event = { c = true })
+          ExpectUnmodified(event = { eventFired(b) })
+          Wrap {
+            ExpectModified(event = { a++ })
+            ExpectModified(event = { eventFired(a) })
+            ExpectModified(event = { c = true })
+            ExpectUnmodified(event = { eventFired(b) })
+          }
+        }
+
+        @Composable
+        fun Example(model: String) {
+          Test(model, 1)
+        }
+    """)
+
+    @Test
+    fun considerNonComposableCaptures() = skipping("""
+        fun eventFired(data: Int) {}
+
+        @Composable
+        fun ExpectUnmodified(event: () -> Unit) {
+          workToBeAvoided()
+        }
+
+        @Composable
+        fun ExpectModified(event: () -> Unit) {
+          workToBeRepeated()
+        }
+
+        inline fun wrap(value: Int, block: (value: Int) -> Unit) {
+          block(value)
+        }
+
+        @Composable
+        fun Example(model: String) {
+           wrap(iterations) { number ->
+             ExpectModified(event = { eventFired(number) })
+             ExpectUnmodified(event = { eventFired(5) })
+           }
+        }
+    """)
+
+    private fun skipping(text: String, dumpClasses: Boolean = false) =
+        ensureSetup {
+            compose("""
+                var avoidedWorkCount = 0
+                var repeatedWorkCount = 0
+                var expectedAvoidedWorkCount = 0
+                var expectedRepeatedWorkCount = 0
+
+                fun workToBeAvoided(msg: String = "") {
+                   avoidedWorkCount++
+                   println("Work to be avoided ${'$'}avoidedWorkCount ${'$'}msg")
+                }
+                fun workToBeRepeated(msg: String = "") {
+                   repeatedWorkCount++
+                   println("Work to be repeated ${'$'}repeatedWorkCount ${'$'}msg")
+                }
+
+                $text
+
+                @Composable
+                fun Display(text: String) {}
+
+                fun validateModel(text: String) {
+                  require(text == "Iteration ${'$'}iterations")
+                }
+
+                @Composable
+                fun ValidateModel(text: String) {
+                  validateModel(text)
+                }
+
+                @Composable
+                fun TestHost() {
+                   println("START: Iteration - ${'$'}iterations")
+                   Button(id=42, onClick=invalidate)
+                   Example("Iteration ${'$'}iterations")
+                   println("END  : Iteration - ${'$'}iterations")
+                   validate()
+                }
+
+                var iterations = 0
+
+                fun validate() {
+                  if (iterations++ == 0) {
+                    expectedAvoidedWorkCount = avoidedWorkCount
+                    expectedRepeatedWorkCount = repeatedWorkCount
+                    repeatedWorkCount = 0
+                  } else {
+                    require(expectedAvoidedWorkCount == avoidedWorkCount) {
+                      "Executed avoided work unexpectedly, expected " +
+                      "${'$'}expectedAvoidedWorkCount" +
+                      ", received ${'$'}avoidedWorkCount"
+                    }
+                    require(expectedRepeatedWorkCount == repeatedWorkCount) {
+                      "Expected more repeated work, expected ${'$'}expectedRepeatedWorkCount" +
+                      ", received ${'$'}repeatedWorkCount"
+                    }
+                    repeatedWorkCount = 0
+                  }
+                }
+
+            """, """
+                TestHost()
+            """, dumpClasses = dumpClasses).then { activity ->
+                val button = activity.findViewById(42) as Button
+                button.performClick()
+            }.then { activity ->
+                val button = activity.findViewById(42) as Button
+                button.performClick()
+            }.then {
+                // Wait for test to complete
+            }
+        }
+}

plugins/compose-compiler-hosted/integration-tests/src/test/java/androidx/compose/plugins/kotlin/RobolectricComposeTester.kt

@@ -0,0 +1,93 @@
+/*
+ * Copyright 2020 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.compose.plugins.kotlin
+
+import android.app.Activity
+import android.os.Bundle
+import android.view.ViewGroup
+import android.widget.LinearLayout
+import androidx.compose.Compose
+import androidx.compose.Composer
+import androidx.compose.Composition
+import androidx.compose.Recomposer
+import org.robolectric.Robolectric
+import org.robolectric.RuntimeEnvironment
+import kotlin.reflect.full.findParameterByName
+import kotlin.reflect.full.functions
+import kotlin.reflect.full.isSubtypeOf
+import kotlin.reflect.full.starProjectedType
+
+const val ROOT_ID = 18284847
+
+private class TestActivity : Activity() {
+    override fun onCreate(savedInstanceState: Bundle?) {
+        super.onCreate(savedInstanceState)
+        setContentView(LinearLayout(this).apply { id = ROOT_ID })
+    }
+}
+
+private val Activity.root get() = findViewById(ROOT_ID) as ViewGroup
+
+fun compose(composable: (Composer<*>) -> Unit) =
+    RobolectricComposeTester(composable)
+fun composeMulti(composable: (Composer<*>) -> Unit, advance: () -> Unit) =
+    RobolectricComposeTester(composable, advance)
+
+class RobolectricComposeTester internal constructor(
+    val composable: (Composer<*>) -> Unit,
+    val advance: (() -> Unit)? = null
+) {
+    inner class ActiveTest(
+        val activity: Activity,
+        val advance: () -> Unit
+    ) {
+        fun then(block: (activity: Activity) -> Unit): ActiveTest {
+            val scheduler = RuntimeEnvironment.getMasterScheduler()
+            scheduler.advanceToLastPostedRunnable()
+            advance()
+            scheduler.advanceToLastPostedRunnable()
+            block(activity)
+            return this
+        }
+    }
+
+    fun then(block: (activity: Activity) -> Unit): ActiveTest {
+        val scheduler = RuntimeEnvironment.getMasterScheduler()
+        scheduler.pause()
+        val controller = Robolectric.buildActivity(TestActivity::class.java)
+        val activity = controller.create().get()
+        val root = activity.root
+        scheduler.advanceToLastPostedRunnable()
+        val composeInto = Compose::class.java.methods.first {
+            if (it.name != "composeInto") false
+            else {
+                val param = it.parameters.getOrNull(2)
+                param?.type == Function1::class.java
+            }
+        }
+        val composition = composeInto.invoke(
+            Compose,
+            root,
+            null,
+            composable
+        ) as Composition
+        scheduler.advanceToLastPostedRunnable()
+        block(activity)
+        val advanceFn = advance ?: { composition.compose() }
+        return ActiveTest(activity, advanceFn)
+    }
+}

plugins/compose-compiler-hosted/integration-tests/src/test/java/androidx/compose/plugins/kotlin/ScopeComposabilityTests.kt

@@ -0,0 +1,227 @@
+/*
+ * Copyright 2019 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.compose.plugins.kotlin
+
+import androidx.compose.plugins.kotlin.analysis.ComposeWritableSlices
+import androidx.compose.plugins.kotlin.ComposableAnnotationChecker.Composability
+import com.intellij.psi.PsiElement
+import org.jetbrains.kotlin.psi.KtElement
+import org.jetbrains.kotlin.psi.KtFile
+import org.jetbrains.kotlin.psi.KtFunction
+import org.jetbrains.kotlin.psi.KtFunctionLiteral
+import org.jetbrains.kotlin.psi.KtLambdaExpression
+import org.jetbrains.kotlin.psi.KtProperty
+import org.jetbrains.kotlin.psi.KtPropertyAccessor
+import org.jetbrains.kotlin.psi.KtPsiFactory
+import org.jetbrains.kotlin.resolve.BindingContext
+import org.jetbrains.kotlin.resolve.calls.model.ResolvedCall
+import org.junit.Ignore
+
+class ScopeComposabilityTests : AbstractCodegenTest() {
+
+    fun testNormalFunctions() = assertComposability(
+        """
+            import androidx.compose.*
+
+            fun Foo() {
+                <normal>
+            }
+            class Bar {
+                fun bam() { <normal> }
+                val baz: Int get() { <normal>return 123 }
+            }
+        """
+    )
+
+    fun testPropGetter() = assertComposability(
+        """
+            import androidx.compose.*
+
+            val baz: Int get() { <normal>return 123 }
+        """
+    )
+
+    fun testBasicComposable() = assertComposability(
+        """
+            import androidx.compose.*
+
+            @Composable
+            fun Foo() {
+                <marked>
+            }
+        """
+    )
+
+    fun testBasicComposable2() = assertComposability(
+        """
+            import androidx.compose.*
+
+            val foo = @Composable { <marked> }
+
+            @Composable
+            fun Bar() {
+                <marked>
+                fun bam() { <normal> }
+                val x = { <normal> }
+                val y = @Composable { <marked> }
+                @Composable fun z() { <marked> }
+            }
+        """
+    )
+
+    // TODO(b/147250515): get inlined lambdas to analyze correctly
+    fun xtestBasicComposable3() = assertComposability(
+        """
+            import androidx.compose.*
+
+            @Composable
+            fun Bar() {
+                <marked>
+                listOf(1, 2, 3).forEach {
+                    <inferred> // should be inferred, but is normal
+                }
+            }
+        """
+    )
+
+    fun testBasicComposable4() = assertComposability(
+        """
+            import androidx.compose.*
+
+            @Composable fun Wrap(block: @Composable() () -> Unit) { block() }
+
+            @Composable
+            fun Bar() {
+                <marked>
+                Wrap {
+                    <marked>
+                    Wrap {
+                        <marked>
+                    }
+                }
+            }
+        """
+    )
+
+    fun testBasicComposable5() = assertComposability(
+        """
+            import androidx.compose.*
+
+            @Composable fun Callback(block: () -> Unit) { block() }
+
+            @Composable
+            fun Bar() {
+                <marked>
+                Callback {
+                    <normal>
+                }
+            }
+        """
+    )
+
+    fun testBasicComposable6() = assertComposability(
+        """
+            import androidx.compose.*
+
+            fun kickOff(block: @Composable() () -> Unit) {  }
+
+            fun Bar() {
+                <normal>
+                kickOff {
+                    <marked>
+                }
+            }
+        """
+    )
+
+    private fun <T> setup(block: () -> T): T {
+        return block()
+    }
+
+    fun assertComposability(srcText: String) = setup {
+        val (text, carets) = extractCarets(srcText)
+
+        val environment = myEnvironment ?: error("Environment not initialized")
+
+        val ktFile = KtPsiFactory(environment.project).createFile(text)
+        val bindingContext = JvmResolveUtil.analyze(
+            ktFile,
+            environment
+        ).bindingContext
+
+        carets.forEachIndexed { index, (offset, marking) ->
+            val composability = composabiliityAtOffset(bindingContext, ktFile, offset)
+                ?: error("composability not found for index: $index, offset: $offset. Expected " +
+                        "$marking.")
+
+            when (marking) {
+                "<marked>" -> assertEquals("index: $index", Composability.MARKED, composability)
+                "<inferred>" -> assertEquals("index: $index", Composability.INFERRED, composability)
+                "<normal>" -> assertEquals(
+                    "index: $index",
+                    Composability.NOT_COMPOSABLE,
+                    composability
+                )
+                else -> error("Composability of $marking not recognized.")
+            }
+        }
+    }
+
+    private val callPattern = Regex("(<marked>)|(<inferred>)|(<normal>)")
+    private fun extractCarets(text: String): Pair<String, List<Pair<Int, String>>> {
+        val indices = mutableListOf<Pair<Int, String>>()
+        var offset = 0
+        val src = callPattern.replace(text) {
+            indices.add(it.range.first - offset to it.value)
+            offset += it.range.last - it.range.first + 1
+            ""
+        }
+        return src to indices
+    }
+
+    private fun composabiliityAtOffset(
+        bindingContext: BindingContext,
+        jetFile: KtFile,
+        index: Int
+    ): Composability? {
+        val element = jetFile.findElementAt(index)!!
+        return element.getNearestComposability(bindingContext)
+    }
+}
+
+fun PsiElement?.getNearestComposability(
+    bindingContext: BindingContext
+): Composability? {
+    var node: PsiElement? = this
+    while (node != null) {
+        when (node) {
+            is KtFunctionLiteral -> {
+                // keep going, as this is a "KtFunction", but we actually want the
+                // KtLambdaExpression
+            }
+            is KtLambdaExpression,
+            is KtFunction,
+            is KtPropertyAccessor,
+            is KtProperty -> {
+                val el = node as KtElement
+                return bindingContext.get(ComposeWritableSlices.COMPOSABLE_ANALYSIS, el)
+            }
+        }
+        node = node.parent as? KtElement
+    }
+    return null
+}

plugins/compose-compiler-hosted/integration-tests/src/test/java/androidx/compose/plugins/kotlin/TestsCompilerError.kt

@@ -0,0 +1,19 @@
+/*
+ * Copyright 2019 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.compose.plugins.kotlin
+
+class TestsCompilerError(val original: Throwable) : Throwable(original)

plugins/compose-compiler-hosted/integration-tests/src/test/java/androidx/compose/plugins/kotlin/analysis/ComposableCheckerTests.kt

@@ -0,0 +1,985 @@
+package androidx.compose.plugins.kotlin.analysis
+
+import org.jetbrains.kotlin.cli.jvm.compiler.EnvironmentConfigFiles
+import org.jetbrains.kotlin.cli.jvm.compiler.KotlinCoreEnvironment
+import org.jetbrains.kotlin.cli.jvm.config.addJvmClasspathRoots
+import androidx.compose.plugins.kotlin.AbstractComposeDiagnosticsTest
+import androidx.compose.plugins.kotlin.newConfiguration
+import com.intellij.openapi.util.Disposer
+
+class ComposableCheckerTests : AbstractComposeDiagnosticsTest() {
+
+    override fun setUp() {
+        // intentionally don't call super.setUp() here since we are recreating an environment
+        // every test
+        System.setProperty("user.dir",
+            homeDir
+        )
+    }
+
+    fun doTest(text: String, expectPass: Boolean) {
+        val disposable = TestDisposable()
+        val classPath = createClasspath()
+        val configuration = newConfiguration()
+        configuration.addJvmClasspathRoots(classPath)
+
+        val environment =
+            KotlinCoreEnvironment.createForTests(
+                disposable,
+                configuration,
+                EnvironmentConfigFiles.JVM_CONFIG_FILES
+            )
+        setupEnvironment(environment)
+
+        try {
+            doTest(text, environment)
+            if (!expectPass) {
+                throw ExpectedFailureException(
+                    "Test unexpectedly passed, but SHOULD FAIL"
+                )
+            }
+        } catch (e: ExpectedFailureException) {
+            throw e
+        } catch (e: Exception) {
+            if (expectPass) throw Exception(e)
+        } finally {
+            Disposer.dispose(disposable)
+        }
+    }
+
+    class ExpectedFailureException(message: String) : Exception(message)
+
+    fun check(expectedText: String) {
+        doTest(expectedText, true)
+    }
+
+    fun checkFail(expectedText: String) {
+        doTest(expectedText, false)
+    }
+
+    fun testComposableReporting001() {
+        checkFail("""
+            import androidx.compose.*;
+            import android.widget.TextView;
+
+            fun myStatelessFunctionalComponent() {
+                TextView(text="Hello World!")
+            }
+
+            fun foo() {
+                myStatelessFunctionalComponent()
+            }
+        """)
+        check("""
+            import androidx.compose.*;
+            import android.widget.TextView;
+
+            @Composable
+            fun myStatelessFunctionalComponent() {
+                TextView(text="Hello World!")
+            }
+
+            @Composable
+            fun foo() {
+                myStatelessFunctionalComponent()
+            }
+        """)
+    }
+
+    fun testComposableReporting002() {
+        checkFail("""
+            import androidx.compose.*;
+            import android.widget.TextView;
+
+            val myLambda1 = { TextView(text="Hello World!") }
+            val myLambda2: () -> Unit = { TextView(text="Hello World!") }
+        """)
+        check("""
+            import androidx.compose.*;
+            import android.widget.TextView;
+
+            val myLambda1 = @Composable() { TextView(text="Hello World!") }
+            val myLambda2: @Composable() ()->Unit = { TextView(text="Hello World!") }
+        """)
+    }
+
+    fun testComposableReporting003() {
+        check("""
+            import androidx.compose.*;
+            import android.widget.TextView;
+
+            fun myRandomFunction() {
+                <!NONE_APPLICABLE!>TextView<!>(text="Hello World!")
+            }
+        """)
+    }
+
+    fun testComposableReporting004() {
+        check("""
+            import androidx.compose.*;
+            import android.widget.TextView;
+
+            @Composable
+            fun foo() {
+                val myRandomLambda = { <!NONE_APPLICABLE!>TextView<!>(text="Hello World!") }
+                System.out.println(myRandomLambda)
+            }
+        """)
+    }
+
+    fun testComposableReporting006() {
+        checkFail("""
+            import androidx.compose.*;
+            import android.widget.TextView;
+
+            fun foo() {
+                val bar = {
+                    TextView()
+                }
+                bar()
+                System.out.println(bar)
+            }
+        """)
+        check("""
+            import androidx.compose.*;
+            import android.widget.TextView;
+
+            @Composable
+            fun foo() {
+                val bar = @Composable {
+                    TextView()
+                }
+                bar()
+                System.out.println(bar)
+            }
+        """)
+    }
+
+    fun testComposableReporting007() {
+        checkFail("""
+            import androidx.compose.*;
+            import android.widget.TextView;
+
+            fun foo(children: @Composable() ()->Unit) {
+                <!SVC_INVOCATION!>children<!>()
+            }
+        """)
+    }
+
+    fun testComposableReporting008() {
+        checkFail("""
+            import androidx.compose.*;
+            import android.widget.TextView;
+
+            fun foo() {
+                val bar: @Composable() ()->Unit = @Composable {
+                    TextView()
+                }
+                <!SVC_INVOCATION!>bar<!>()
+                System.out.println(bar)
+            }
+        """)
+    }
+
+    fun testComposableReporting009() {
+        checkFail("""
+            import androidx.compose.*;
+            import android.widget.TextView;
+
+            fun myStatelessFunctionalComponent() {
+                TextView(text="Hello World!")
+            }
+
+            fun noise() {
+                <!SVC_INVOCATION!>myStatelessFunctionalComponent<!>()
+            }
+        """)
+        checkFail("""
+            import androidx.compose.*;
+            import android.widget.TextView;
+
+            @Composable
+            fun myStatelessFunctionalComponent() {
+                TextView(text="Hello World!")
+            }
+
+            fun noise() {
+                <!SVC_INVOCATION!>myStatelessFunctionalComponent<!>()
+            }
+        """)
+    }
+
+    fun testComposableReporting016() {
+        check("""
+            import androidx.compose.*;
+
+            <!WRONG_ANNOTATION_TARGET!>@Composable<!>
+            class Noise() {}
+        """)
+
+        check("""
+            import androidx.compose.*;
+
+            val adHoc = <!WRONG_ANNOTATION_TARGET!>@Composable()<!> object {
+                var x: Int = 0
+                var y: Int = 0
+            }
+        """)
+
+        check("""
+            import androidx.compose.*;
+
+            open class Noise() {}
+
+            val adHoc = <!WRONG_ANNOTATION_TARGET!>@Composable()<!> object : Noise() {
+                var x: Int = 0
+                var y: Int = 0
+            }
+        """)
+    }
+
+    fun testComposableReporting017() {
+        checkFail("""
+            import androidx.compose.*;
+            import android.widget.LinearLayout;
+            import android.widget.TextView;
+
+            @Composable
+            fun Foo(children: ()->Unit) {
+                children()
+            }
+
+            @Composable
+            fun main() {
+                Foo { TextView(text="Hello") }
+            }
+        """)
+        check("""
+            import androidx.compose.*;
+            import android.widget.LinearLayout;
+            import android.widget.TextView;
+
+            @Composable
+            fun Foo(children: ()->Unit) {
+                children()
+            }
+
+            @Composable
+            fun main() {
+                Foo { <!NONE_APPLICABLE!>TextView<!>(text="Hello") }
+            }
+        """)
+    }
+
+    fun testComposableReporting018() {
+        checkFail("""
+            import androidx.compose.*;
+            import android.widget.TextView;
+
+            fun foo() {
+                val myVariable: ()->Unit = @Composable { TextView(text="Hello World!") }
+                System.out.println(myVariable)
+            }
+        """)
+        check("""
+            import androidx.compose.*;
+            import android.widget.TextView;
+
+            fun foo() {
+                val myVariable: ()->Unit = <!TYPE_MISMATCH!>@Composable {
+                 TextView(text="Hello World!") }<!>
+                System.out.println(myVariable)
+            }
+        """)
+    }
+
+    fun testComposableReporting021() {
+        check("""
+            import androidx.compose.*;
+            import android.widget.TextView;
+
+            @Composable
+            fun foo() {
+                val myList = listOf(1,2,3,4,5)
+                myList.forEach @Composable { value: Int -> 
+                    TextView(text=value.toString())
+                    System.out.println(value)
+                }
+            }
+        """)
+    }
+
+    fun testComposableReporting022() {
+        checkFail("""
+            import androidx.compose.*;
+            import android.widget.TextView;
+
+            fun foo() {
+                val myList = listOf(1,2,3,4,5)
+                myList.forEach { value: Int -> 
+                    TextView(text=value.toString()) 
+                    System.out.println(value)
+                }
+            }
+        """)
+        check("""
+            import androidx.compose.*;
+            import android.widget.TextView;
+
+            fun <!COMPOSABLE_INVOCATION_IN_NON_COMPOSABLE!>foo<!>() {
+                val myList = listOf(1,2,3,4,5)
+                myList.forEach @Composable { value: Int -> 
+                    <!COMPOSABLE_INVOCATION_IN_NON_COMPOSABLE!>TextView<!>(text=value.toString())
+                    System.out.println(value) 
+                }
+            }
+        """)
+    }
+
+    fun testComposableReporting024() {
+        check("""
+            import androidx.compose.*;
+            import android.widget.TextView;
+            import android.widget.LinearLayout;
+
+            fun foo(ll: LinearLayout) {
+                ll.setViewContent({ TextView(text="Hello World!") })
+            }
+        """)
+    }
+
+    fun testComposableReporting025() {
+        check("""
+            import androidx.compose.*;
+            import android.widget.TextView;
+
+            @Composable
+            fun foo() {
+                listOf(1,2,3,4,5).forEach { TextView(text="Hello World!") }
+            }
+        """)
+    }
+
+    fun testComposableReporting026() {
+        check("""
+            import androidx.compose.*;
+            import android.widget.LinearLayout;
+            import android.widget.TextView;
+
+            @Composable
+            fun foo() {
+                LinearLayout {
+                    TextView(text="Hello Jim!")
+                }
+            }
+        """)
+    }
+
+    fun testComposableReporting027() {
+        check("""
+            import androidx.compose.*;
+            import android.widget.LinearLayout;
+            import android.widget.TextView;
+
+            @Composable
+            fun foo() {
+                LinearLayout {
+                    listOf(1,2,3).forEach {
+                        TextView(text="Hello Jim!")
+                    }
+                }
+            }
+        """)
+    }
+
+    fun testComposableReporting028() {
+        checkFail("""
+            import androidx.compose.*;
+            import android.widget.TextView;
+
+            fun foo(v: @Composable() ()->Unit) {
+                val myVariable: ()->Unit = v
+                myVariable()
+            }
+        """)
+        check("""
+            import androidx.compose.*;
+            import android.widget.TextView;
+
+            fun foo(v: @Composable() ()->Unit) {
+                val myVariable: ()->Unit = <!TYPE_MISMATCH!>v<!>
+                myVariable()
+            }
+        """)
+    }
+
+    fun testComposableReporting030() {
+        check("""
+            import androidx.compose.*;
+            import android.widget.TextView;
+
+            @Composable
+            fun foo() {
+                val myVariable: @Composable() ()->Unit = {};
+                myVariable()
+            }
+        """)
+    }
+
+    fun testComposableReporting031() {
+        check("""
+            import androidx.compose.*;
+            import android.widget.TextView;
+
+            fun foo() {
+                val myVariable: ()->Unit = { <!NONE_APPLICABLE!>TextView<!>(text="Hello") };
+                myVariable();
+            }
+        """)
+    }
+
+    fun testComposableReporting032() {
+        check("""
+            import androidx.compose.*;
+            import android.widget.TextView;
+
+            @Composable
+            fun MyComposable(children: @Composable() ()->Unit) { children() }
+
+            @Composable
+            fun foo() {
+                MyComposable { TextView(text="Hello") }
+            }
+        """)
+    }
+
+    fun testComposableReporting033() {
+        check("""
+            import androidx.compose.*;
+            import android.widget.TextView;
+
+            @Composable
+            fun MyComposable(children: @Composable() ()->Unit) { children() }
+
+            @Composable
+            fun foo() {
+                MyComposable(children={ TextView(text="Hello")})
+            }
+        """)
+    }
+
+    fun testComposableReporting034() {
+        checkFail("""
+            import androidx.compose.*;
+            import android.widget.TextView;
+
+            fun identity(f: ()->Unit): ()->Unit { return f; }
+
+            @Composable
+            fun test(f: @Composable() ()->Unit) {
+                val f2: @Composable() ()->Unit = identity(f);
+                f2()
+            }
+        """)
+        check("""
+            import androidx.compose.*;
+            import android.widget.TextView;
+
+            fun identity(f: ()->Unit): ()->Unit { return f; }
+
+            @Composable
+            fun test(f: @Composable() ()->Unit) {
+                val f2: @Composable() ()->Unit = <!TYPE_MISMATCH!>identity (<!TYPE_MISMATCH!>f<!>)<!>;
+                f2()
+            }
+        """)
+    }
+
+    fun testComposableReporting035() {
+        check("""
+            import androidx.compose.*
+
+            @Composable
+            fun Foo(x: String) {
+                @Composable operator fun String.invoke() {}
+                x()
+            }
+        """)
+    }
+
+    fun testComposableReporting036() {
+        checkFail("""
+            import androidx.compose.*
+            import android.widget.TextView;
+
+            fun Foo() {
+                repeat(5) {
+                    TextView(text="Hello World")
+                }
+            }
+
+            fun Bar() {
+                <!SVC_INVOCATION!>Foo<!>()
+            }
+        """)
+        checkFail("""
+            import androidx.compose.*
+            import android.widget.TextView;
+
+            fun Foo() {
+                repeat(5) {
+                    TextView(text="Hello World")
+                }
+            }
+
+            fun Bar() {
+                <!SVC_INVOCATION!>Foo<!>()
+            }
+        """)
+    }
+
+    fun testComposableReporting037() {
+        checkFail("""
+            import androidx.compose.*
+            import android.widget.TextView;
+
+            fun Foo() {
+                fun Noise() {
+                    TextView(text="Hello World")
+                }
+            }
+
+            fun Bar() {
+                Foo()
+            }
+        """)
+    }
+
+    fun testComposableReporting038() {
+        check("""
+            import androidx.compose.*
+            import android.widget.TextView;
+
+            // Function intentionally not inline
+            fun repeat(x: Int, l: ()->Unit) { for(i in 1..x) l() }
+
+            fun Foo() {
+                repeat(5) {
+                    <!NONE_APPLICABLE!>TextView<!>(text="Hello World")
+                }
+            }
+        """)
+    }
+
+    fun testComposableReporting039() {
+        check(
+            """
+            import androidx.compose.*
+            import android.widget.TextView;
+
+            fun composeInto(l: @Composable() ()->Unit) { System.out.println(l) }
+
+            fun Foo() {
+                composeInto {
+                    TextView(text="Hello World")
+                }
+            }
+
+            fun Bar() {
+                Foo()
+            }
+        """
+        )
+    }
+
+    fun testComposableReporting040() {
+        checkFail("""
+            import androidx.compose.*
+            import android.widget.TextView;
+
+            inline fun noise(l: ()->Unit) { l() }
+
+            fun Foo() {
+                noise {
+                    TextView(text="Hello World")
+                }
+            }
+
+            fun Bar() {
+                <!SVC_INVOCATION!>Foo<!>()
+            }
+        """)
+    }
+
+    fun testComposableReporting041() {
+        check("""
+            import androidx.compose.*
+            import android.widget.TextView;
+
+            typealias COMPOSABLE_UNIT_LAMBDA = @Composable() () -> Unit
+
+            @Composable
+            fun ComposeWrapperComposable(children: COMPOSABLE_UNIT_LAMBDA) {
+                MyComposeWrapper {
+                    children()
+                }
+            }
+
+            @Composable fun MyComposeWrapper(children: COMPOSABLE_UNIT_LAMBDA) {
+                print(children.hashCode())
+            }
+        """)
+    }
+
+    fun testComposableReporting043() {
+        check("""
+            import androidx.compose.*
+            import android.widget.TextView;
+
+            typealias UNIT_LAMBDA = () -> Unit
+
+            @Composable
+            fun FancyButton() {}
+
+            fun <!COMPOSABLE_INVOCATION_IN_NON_COMPOSABLE!>Noise<!>() {
+                <!COMPOSABLE_INVOCATION_IN_NON_COMPOSABLE,COMPOSABLE_INVOCATION_IN_NON_COMPOSABLE!>FancyButton<!>()
+            }
+        """)
+    }
+
+    fun testComposableReporting044() {
+        check("""
+            import androidx.compose.*
+            import android.widget.TextView;
+
+            typealias UNIT_LAMBDA = () -> Unit
+
+            @Composable
+            fun FancyButton() {}
+
+            @Composable
+            fun Noise() {
+                FancyButton()
+            }
+        """)
+    }
+
+    fun testComposableReporting045() {
+        check("""
+            import androidx.compose.*;
+
+            @Composable
+            fun foo() {
+                val bar = @Composable {}
+                bar()
+                System.out.println(bar)
+            }
+        """)
+    }
+
+    fun testComposableReporting047() {
+        check("""
+            import androidx.compose.*
+            import android.widget.LinearLayout;
+
+            @Composable
+            fun FancyButton() {}
+
+            @Composable
+            fun Foo() {
+                LinearLayout {
+                    FancyButton()
+                }
+            }
+        """)
+    }
+
+    fun testComposableReporting048() {
+        // Type inference for non-null @Composable lambdas
+        checkFail("""
+            import androidx.compose.*
+
+            val lambda: @Composable() (() -> Unit)? = null
+
+            @Composable
+            fun Foo() {
+	        // Should fail as null cannot be coerced to non-null
+                Bar(lambda)
+                Bar(null)
+                Bar {}
+            }
+
+            @Composable
+            fun Bar(child: @Composable() () -> Unit) {
+                child()
+            }
+        """)
+
+        // Type inference for nullable @Composable lambdas, with no default value
+        check("""
+            import androidx.compose.*
+
+            val lambda: @Composable() (() -> Unit)? = null
+
+            @Composable
+            fun Foo() {
+                Bar(lambda)
+                Bar(null)
+                Bar {}
+            }
+
+            @Composable
+            fun Bar(child: @Composable() (() -> Unit)?) {
+                child?.invoke()
+            }
+        """)
+
+        // Type inference for nullable @Composable lambdas, with a nullable default value
+        check("""
+            import androidx.compose.*
+
+            val lambda: @Composable() (() -> Unit)? = null
+
+            @Composable
+            fun Foo() {
+                Bar()
+                Bar(lambda)
+                Bar(null)
+                Bar {}
+            }
+
+            @Composable
+            fun Bar(child: @Composable() (() -> Unit)? = null) {
+                child?.invoke()
+            }
+        """)
+
+        // Type inference for nullable @Composable lambdas, with a non-null default value
+        check("""
+            import androidx.compose.*
+
+            val lambda: @Composable() (() -> Unit)? = null
+
+            @Composable
+            fun Foo() {
+                Bar()
+                Bar(lambda)
+                Bar(null)
+                Bar {}
+            }
+
+            @Composable
+            fun Bar(child: @Composable() (() -> Unit)? = {}) {
+                child?.invoke()
+            }
+        """)
+    }
+
+    fun testComposableReporting049() {
+        check("""
+            import androidx.compose.*
+            fun foo(<!WRONG_ANNOTATION_TARGET!>@Composable<!> bar: ()->Unit) {
+                println(bar)
+            }
+        """)
+    }
+
+    fun testComposableReporting050() {
+        checkFail("""
+            import androidx.compose.*;
+
+            @Composable val foo: Int = 123
+
+            fun App() {
+                foo
+            }
+        """)
+        check("""
+            import androidx.compose.*;
+
+            @Composable val foo: Int = 123
+
+            @Composable
+            fun App() {
+                println(foo)
+            }
+        """)
+    }
+
+    fun testComposableReporting051() {
+        checkFail("""
+            import androidx.compose.*;
+
+            class A {
+                @Composable val bar get() = 123
+            }
+
+            @Composable val A.bam get() = 123
+
+            fun App() {
+                val a = A()
+                a.bar
+            }
+        """)
+        checkFail("""
+            import androidx.compose.*;
+
+            class A {
+                @Composable val bar get() = 123
+            }
+
+            @Composable val A.bam get() = 123
+
+            fun App() {
+                val a = A()
+                a.bam
+            }
+        """)
+        check("""
+            import androidx.compose.*;
+
+            class A {
+                @Composable val bar get() = 123
+            }
+
+            @Composable val A.bam get() = 123
+
+            @Composable
+            fun App() {
+                val a = A()
+                a.bar
+                a.bam
+                with(a) {
+                    bar
+                    bam
+                }
+            }
+        """)
+    }
+
+    fun testComposableReporting052() {
+        checkFail("""
+            import androidx.compose.*;
+
+            @Composable fun Foo() {}
+
+            val bam: Int get() {
+                Foo()
+                return 123
+            }
+        """)
+
+        check("""
+            import androidx.compose.*;
+
+            @Composable fun Foo() {}
+
+            @Composable val bam: Int get() {
+                Foo()
+                return 123
+            }
+        """)
+    }
+
+    fun testComposableReporting053() {
+        check("""
+            import androidx.compose.*;
+
+            @Composable fun foo(): Int = 123
+
+            fun <!COMPOSABLE_INVOCATION_IN_NON_COMPOSABLE!>App<!>() {
+                val x = <!COMPOSABLE_INVOCATION_IN_NON_COMPOSABLE,COMPOSABLE_INVOCATION_IN_NON_COMPOSABLE!>foo<!>()
+                print(x)
+            }
+        """)
+    }
+
+    fun testComposableReporting054() {
+        check("""
+            import androidx.compose.*;
+
+            @Composable fun Foo() {}
+
+            val y: Any <!COMPOSABLE_INVOCATION_IN_NON_COMPOSABLE!>get() = 
+            <!COMPOSABLE_INVOCATION_IN_NON_COMPOSABLE,COMPOSABLE_INVOCATION_IN_NON_COMPOSABLE!>state<!> { 1 }<!>
+
+            fun App() {
+                val x = object {
+                  val a <!COMPOSABLE_INVOCATION_IN_NON_COMPOSABLE!>get() = 
+                  <!COMPOSABLE_INVOCATION_IN_NON_COMPOSABLE,COMPOSABLE_INVOCATION_IN_NON_COMPOSABLE!>state<!> { 2 }<!>
+                  @Composable val c get() = state { 4 }
+                  @Composable fun bar() { Foo() }
+                  fun <!COMPOSABLE_INVOCATION_IN_NON_COMPOSABLE!>foo<!>() {
+                    <!COMPOSABLE_INVOCATION_IN_NON_COMPOSABLE,COMPOSABLE_INVOCATION_IN_NON_COMPOSABLE!>Foo<!>() 
+                  }
+                }
+                class Bar {
+                  val b <!COMPOSABLE_INVOCATION_IN_NON_COMPOSABLE!>get() =
+                  <!COMPOSABLE_INVOCATION_IN_NON_COMPOSABLE,COMPOSABLE_INVOCATION_IN_NON_COMPOSABLE!>state<!> { 6 }<!>
+                  @Composable val c get() = state { 7 }
+                }
+                fun <!COMPOSABLE_INVOCATION_IN_NON_COMPOSABLE!>Bam<!>() {
+                    <!COMPOSABLE_INVOCATION_IN_NON_COMPOSABLE,COMPOSABLE_INVOCATION_IN_NON_COMPOSABLE!>Foo<!>()
+                }
+                @Composable fun Boo() {
+                    Foo()
+                }
+                print(x)
+            }
+        """)
+    }
+
+    fun testComposableReporting055() {
+        check("""
+            import androidx.compose.*;
+
+            @Composable fun Foo() {}
+
+            @Composable fun App() {
+                val x = object {
+                  val a <!COMPOSABLE_INVOCATION_IN_NON_COMPOSABLE!>get() = 
+                  <!COMPOSABLE_INVOCATION_IN_NON_COMPOSABLE!><!COMPOSABLE_INVOCATION_IN_NON_COMPOSABLE!>state<!><!> { 2 }<!>
+                  @Composable val c get() = state { 4 }
+                  fun <!COMPOSABLE_INVOCATION_IN_NON_COMPOSABLE!>foo<!>() { 
+                    <!COMPOSABLE_INVOCATION_IN_NON_COMPOSABLE,COMPOSABLE_INVOCATION_IN_NON_COMPOSABLE!>Foo<!>() 
+                  }
+                  @Composable fun bar() { Foo() }
+                }
+                class Bar {
+                  val b <!COMPOSABLE_INVOCATION_IN_NON_COMPOSABLE!>get() = 
+                  <!COMPOSABLE_INVOCATION_IN_NON_COMPOSABLE!><!COMPOSABLE_INVOCATION_IN_NON_COMPOSABLE!>state<!><!> { 6 }<!>
+                  @Composable val c get() = state { 7 }
+                }
+                fun <!COMPOSABLE_INVOCATION_IN_NON_COMPOSABLE!>Bam<!>() {
+                    <!COMPOSABLE_INVOCATION_IN_NON_COMPOSABLE,COMPOSABLE_INVOCATION_IN_NON_COMPOSABLE!>Foo<!>()
+                }
+                @Composable fun Boo() {
+                    Foo()
+                }
+                print(x)
+            }
+        """)
+    }
+
+    fun testComposableReporting057() {
+        // This tests composable calls in initialization expressions of object literals inside of
+        // composable functions. I don't see any reason why we shouldn't support this, but right now
+        // we catch it and prevent it. Enabling it is nontrivial so i'm writing the test to assert
+        // on the current behavior, and we can consider changing it at a later date.
+        check("""
+            import androidx.compose.*;
+
+            @Composable fun App() {
+                val x = object {
+                  <!COMPOSABLE_INVOCATION_IN_NON_COMPOSABLE!>val b = 
+                  <!COMPOSABLE_INVOCATION_IN_NON_COMPOSABLE!><!COMPOSABLE_INVOCATION_IN_NON_COMPOSABLE!>state<!><!> { 3 }<!>
+                }
+                class Bar {
+                  <!COMPOSABLE_INVOCATION_IN_NON_COMPOSABLE!>val a = 
+                  <!COMPOSABLE_INVOCATION_IN_NON_COMPOSABLE!><!COMPOSABLE_INVOCATION_IN_NON_COMPOSABLE!>state<!><!> { 5 }<!>
+                }
+                print(x)
+            }
+        """)
+    }
+}

plugins/compose-compiler-hosted/integration-tests/src/test/java/androidx/compose/plugins/kotlin/analysis/TryCatchComposableCheckerTests.kt

@@ -0,0 +1,92 @@
+/*
+ * Copyright 2019 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.compose.plugins.kotlin.analysis
+
+import androidx.compose.plugins.kotlin.AbstractComposeDiagnosticsTest
+
+/**
+ * We're strongly considering supporting try-catch-finally blocks in the future.
+ * If/when we do support them, these tests should be deleted.
+ */
+class TryCatchComposableCheckerTests : AbstractComposeDiagnosticsTest() {
+
+    fun testTryCatchReporting001() {
+        doTest(
+            """
+            import androidx.compose.*;
+
+            @Composable fun foo() { }
+
+            @Composable fun bar() {
+                <!ILLEGAL_TRY_CATCH_AROUND_COMPOSABLE!>try<!> {
+                    foo()
+                } catch(e: Exception) {
+                }
+            }
+        """)
+    }
+
+    fun testTryCatchReporting002() {
+        doTest(
+            """
+            import androidx.compose.*;
+
+            fun foo() { }
+
+            @Composable fun bar() {
+                try {
+                    foo()
+                } catch(e: Exception) {
+                }
+            }
+        """)
+    }
+
+    fun testTryCatchReporting003() {
+        doTest(
+            """
+            import androidx.compose.*;
+
+            @Composable fun foo() { }
+
+            @Composable fun bar() {
+                try {
+                } catch(e: Exception) {
+                    foo()
+                } finally {
+                    foo()
+                }
+            }
+        """)
+    }
+
+    fun testTryCatchReporting004() {
+        doTest(
+            """
+            import androidx.compose.*;
+
+            @Composable fun foo() { }
+
+            @Composable fun bar() {
+                <!ILLEGAL_TRY_CATCH_AROUND_COMPOSABLE!>try<!> {
+                    (1..10).forEach { foo() }
+                } catch(e: Exception) {
+                }
+            }
+        """)
+    }
+}

plugins/compose-compiler-hosted/integration-tests/src/test/java/androidx/compose/plugins/kotlin/analysis/UnionCheckerTests.kt

@@ -0,0 +1,88 @@
+/*
+ * Copyright 2018 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.compose.plugins.kotlin.analysis
+
+import androidx.compose.plugins.kotlin.AbstractComposeDiagnosticsTest
+
+class UnionCheckerTests : AbstractComposeDiagnosticsTest() {
+
+    fun testUnionTypeReporting001() {
+        doTest(
+            """
+            import androidx.compose.*;
+
+            @Composable fun foo(value: @UnionType(Int::class, String::class) Any) {
+                System.out.println(value)
+            }
+
+            @Composable
+            fun bar() {
+                foo(value=1)
+                foo(value="1")
+                foo(value=<!ILLEGAL_ASSIGN_TO_UNIONTYPE!>1f<!>)
+            }
+        """)
+    }
+
+    fun testUnionTypeReporting002() {
+        doTest(
+            """
+            import androidx.compose.*;
+
+            @Composable fun foo(value: @UnionType(Int::class, String::class) Any) {
+                System.out.println(value)
+            }
+
+            @Composable
+            fun bar(value: @UnionType(Int::class, String::class) Any) {
+                foo(value)
+            }
+        """)
+    }
+
+    fun testUnionTypeReporting003() {
+        doTest(
+            """
+            import androidx.compose.*;
+
+            @Composable fun foo(value: @UnionType(Int::class, String::class, Float::class) Any) {
+                System.out.println(value)
+            }
+
+            @Composable
+            fun bar(value: @UnionType(Int::class, String::class) Any) {
+                foo(value)
+            }
+        """)
+    }
+
+    fun testUnionTypeReporting004() {
+        doTest(
+            """
+            import androidx.compose.*;
+
+            @Composable fun foo(value: @UnionType(Int::class, String::class) Any) {
+                System.out.println(value)
+            }
+
+            @Composable
+            fun bar(value: @UnionType(Int::class, String::class, Float::class) Any) {
+                foo(<!ILLEGAL_ASSIGN_TO_UNIONTYPE!>value<!>)
+            }
+        """)
+    }
+}

plugins/compose-compiler-hosted/integration-tests/src/test/java/androidx/compose/plugins/kotlin/frames/FrameDiagnosticTests.kt

@@ -0,0 +1,72 @@
+package androidx.compose.plugins.kotlin.frames
+
+import androidx.compose.plugins.kotlin.AbstractComposeDiagnosticsTest
+
+class FrameDiagnosticTests : AbstractComposeDiagnosticsTest() {
+
+    // Ensure the simple case does not report an error
+    fun testModel_Accept_Simple() = doTest(
+        """
+        import androidx.compose.Model
+
+        @Model
+        class MyModel {
+          var strValue = "default"
+        }
+        """
+    )
+
+    // Ensure @Model is not used on an open class
+    fun testModel_Report_Open() = doTest(
+        """
+        import androidx.compose.Model
+
+        @Model
+        open class <!OPEN_MODEL!>MyModel<!> {
+          var strValue = "default"
+        }
+        """
+    )
+
+    // Ensure @Model is not used on an abstract class
+    fun testModel_Report_Abstract() = doTest(
+        """
+        import androidx.compose.Model
+
+        @Model
+        abstract class <!OPEN_MODEL!>MyModel<!> {
+          var strValue = "default"
+        }
+        """
+    )
+
+    // Ensure @Model supports inheriting from a non-model class
+    fun testModel_Report_Inheritance() = doTest(
+        """
+        import androidx.compose.Model
+
+        open class NonModel { }
+
+        @Model
+        class MyModel : NonModel() {
+          var strValue = "default"
+        }
+        """
+    )
+
+    // Ensure errors are reported when the class is nested.
+    fun testModel_Report_Nested_Inheritance() = doTest(
+        """
+        import androidx.compose.Model
+
+        open class NonModel { }
+
+        class Tests {
+            @Model
+            open class <!OPEN_MODEL!>MyModel<!> : NonModel() {
+              var strValue = "default"
+            }
+        }
+        """
+    )
+}

plugins/compose-compiler-hosted/integration-tests/src/test/java/androidx/compose/plugins/kotlin/frames/FrameTransformExtensionTests.kt

@@ -0,0 +1,437 @@
+package androidx.compose.plugins.kotlin.frames
+
+import org.jetbrains.kotlin.psi.KtFile
+import androidx.compose.plugins.kotlin.AbstractCodegenTest
+import org.junit.Before
+
+class FrameTransformExtensionTests : AbstractCodegenTest() {
+
+    @Before
+    fun before() {
+        setUp()
+    }
+
+    fun testTestUtilities() = testFile("""
+        class Foo {
+          val s = "This is a test"
+        }
+
+        class Test {
+          fun test() {
+            Foo().s.expectEqual("This is a test")
+          }
+        }
+    """)
+
+    fun testModel_Simple() = testFile("""
+        import androidx.compose.Model
+
+        @Model
+        class MyModel { }
+
+        class Test {
+          fun test() {
+            frame { MyModel() }
+          }
+        }
+    """)
+
+    fun testModel_OneField() = testFile("""
+        import androidx.compose.Model
+
+        @Model
+        class MyModel {
+          var value: String = "default"
+        }
+
+        class Test {
+          fun test() {
+            val instance = frame { MyModel() }
+            frame {
+              instance.value.expectEqual("default")
+              instance.value = "new value"
+              instance.value.expectEqual("new value")
+            }
+            frame {
+              instance.value.expectEqual("new value")
+            }
+          }
+        }
+    """)
+
+    fun testModel_OneField_Isolation() = testFile("""
+        import androidx.compose.Model
+
+        @Model
+        class MyModel {
+          var value: String = "default"
+        }
+
+        class Test {
+          fun test() {
+            val instance = frame { MyModel() }
+            val frame1 = suspended {
+              instance.value = "new value"
+            }
+            frame {
+              instance.value.expectEqual("default")
+            }
+            restored(frame1) {
+              instance.value.expectEqual("new value")
+            }
+            frame {
+              instance.value.expectEqual("new value")
+            }
+          }
+        }
+    """)
+
+    fun testModel_ThreeFields() = testFile("""
+        import androidx.compose.Model
+
+        @Model
+        class MyModel {
+          var strVal = "default"
+          var intVal = 1
+          var doubleVal = 27.2
+        }
+
+        class Test {
+          fun test() {
+            val instance = frame { MyModel() }
+            frame {
+              instance.strVal.expectEqual("default")
+              instance.intVal.expectEqual(1)
+              instance.doubleVal.expectEqual(27.2)
+            }
+            frame {
+              instance.strVal = "new value"
+            }
+            frame {
+              instance.strVal.expectEqual("new value")
+              instance.intVal.expectEqual(1)
+              instance.doubleVal.expectEqual(27.2)
+            }
+            frame {
+              instance.intVal = 2
+            }
+            frame {
+              instance.strVal.expectEqual("new value")
+              instance.intVal.expectEqual(2)
+              instance.doubleVal.expectEqual(27.2)
+            }
+          }
+        }
+    """)
+
+    fun testModel_ThreeFields_Isolation() = testFile("""
+        import androidx.compose.Model
+
+        @Model
+        class MyModel {
+          var strVal = "default"
+          var intVal = 1
+          var doubleVal = 27.2
+        }
+
+        class Test {
+          fun test() {
+            val instance = frame { MyModel() }
+            frame {
+              instance.strVal.expectEqual("default")
+              instance.intVal.expectEqual(1)
+              instance.doubleVal.expectEqual(27.2)
+            }
+            val frame1 = suspended {
+              instance.strVal = "new value"
+            }
+            frame {
+              instance.strVal.expectEqual("default")
+              instance.intVal.expectEqual(1)
+              instance.doubleVal.expectEqual(27.2)
+            }
+            restored(frame1) {
+              instance.intVal = 2
+            }
+            frame {
+              instance.strVal.expectEqual("new value")
+              instance.intVal.expectEqual(2)
+              instance.doubleVal.expectEqual(27.2)
+            }
+          }
+        }
+    """)
+
+    fun testModel_CustomSetter_Isolation() = testFile("""
+        import androidx.compose.Model
+
+        @Model
+        class MyModel {
+          var intVal = 0; set(value) { field = value; intVal2 = value + 10 }
+          var intVal2 = 10
+          var intVal3 = 0; get() = field + 7; set(value) { field = value - 7 }
+        }
+
+        class Test {
+          fun test() {
+            val instance = frame { MyModel() }
+            frame {
+              instance.intVal.expectEqual(0)
+              instance.intVal2.expectEqual(10)
+              instance.intVal3.expectEqual(7)
+
+              instance.intVal = 22
+              instance.intVal3 = 14
+
+              instance.intVal.expectEqual(22)
+              instance.intVal2.expectEqual(32)
+              instance.intVal3.expectEqual(14)
+            }
+            val frame1 = suspended {
+              instance.intVal = 32
+              instance.intVal3 = 21
+
+              instance.intVal.expectEqual(32)
+              instance.intVal2.expectEqual(42)
+              instance.intVal3.expectEqual(21)
+            }
+            frame {
+              instance.intVal.expectEqual(22)
+              instance.intVal2.expectEqual(32)
+              instance.intVal3.expectEqual(14)
+            }
+            restored(frame1) {
+              instance.intVal.expectEqual(32)
+              instance.intVal2.expectEqual(42)
+              instance.intVal3.expectEqual(21)
+            }
+            frame {
+              instance.intVal.expectEqual(32)
+              instance.intVal2.expectEqual(42)
+              instance.intVal3.expectEqual(21)
+            }
+          }
+        }
+    """)
+
+    fun testModel_PrivateFields_Isolation() = testFile("""
+        import androidx.compose.Model
+
+        @Model
+        class MyModel {
+          private var myIntVal = 1
+          private var myStrVal = "default"
+
+          var intVal get() = myIntVal; set(value) { myIntVal = value }
+          var strVal get() = myStrVal; set(value) { myStrVal = value }
+        }
+
+        class Test {
+          fun test() {
+            val instance = frame { MyModel() }
+            frame {
+              instance.strVal.expectEqual("default")
+              instance.intVal.expectEqual(1)
+            }
+            val frame1 = suspended {
+              instance.strVal = "new value"
+              instance.intVal = 2
+              instance.strVal.expectEqual("new value")
+              instance.intVal.expectEqual(2)
+            }
+            frame {
+              instance.strVal.expectEqual("default")
+              instance.intVal.expectEqual(1)
+            }
+            restored(frame1) {
+              instance.strVal.expectEqual("new value")
+              instance.intVal.expectEqual(2)
+            }
+            frame {
+              instance.strVal.expectEqual("new value")
+              instance.intVal.expectEqual(2)
+            }
+          }
+        }
+    """)
+
+    // regression: 144668818
+    fun testModel_InitSection() = testFile("""
+        import androidx.compose.Model
+
+        @Model
+        class Ints {
+            var values: IntArray
+
+            init {
+                values = intArrayOf(1, 2, 3)
+            }
+        }
+
+        class Test {
+          fun test() {
+            val instance = frame { Ints() }
+            frame {
+              instance.values.size.expectEqual(3)
+              instance.values[0].expectEqual(1)
+              instance.values[1].expectEqual(2)
+              instance.values[2].expectEqual(3)
+            }
+            frame {
+              instance.values = intArrayOf(1, 2, 3, 4)
+            }
+            frame {
+              instance.values.size.expectEqual(4)
+              instance.values[0].expectEqual(1)
+              instance.values[1].expectEqual(2)
+              instance.values[2].expectEqual(3)
+              instance.values[3].expectEqual(4)
+            }
+          }
+        }
+    """)
+
+    // regression: 144668818
+    fun testModel_Initializers() = testFile("""
+        import androidx.compose.Model
+
+        @Model
+        class ABC(var a: Int, b: Int = a) { val c = a }
+
+        class Test {
+          fun test() {
+            val instance = frame { ABC(1) }
+            frame {
+              instance.a.expectEqual(1)
+              instance.c.expectEqual(1)
+            }
+          }
+        }
+    """)
+
+    fun testModel_NestedClass() = testFile("""
+        import androidx.compose.Model
+
+        class Test {
+          @Model
+          class MyModel {
+            var value: String = "default"
+          }
+
+          fun test() {
+            val instance = frame { MyModel() }
+            val frame1 = suspended {
+              instance.value = "new value"
+            }
+            frame {
+              instance.value.expectEqual("default")
+            }
+            restored(frame1) {
+              instance.value.expectEqual("new value")
+            }
+            frame {
+              instance.value.expectEqual("new value")
+            }
+          }
+        }
+    """)
+
+    fun testModel_NonModelInheritance() = testFile("""
+        import androidx.compose.Model
+
+        open class NonModel(var nonTransactedValue: String = "default")
+
+        @Model
+        class MyModel : NonModel() {
+          var value: String = "default"
+        }
+
+        class Test {
+          fun test() {
+            val instance = frame { MyModel() }
+
+            frame {
+              // Ensure the fields have their default values
+              instance.nonTransactedValue.expectEqual("default")
+              instance.value.expectEqual("default")
+            }
+
+            val frame1 = suspended {
+              instance.nonTransactedValue = "modified in suspended transaction"
+              instance.value = "modified in suspended transaction"
+            }
+
+            frame {
+              // Transacted field should still be the default value.
+              instance.value.expectEqual("default")
+
+              // Inherited non-transacted field is not isolated and is expected to be modified
+              instance.nonTransactedValue.expectEqual("modified in suspended transaction")
+            }
+
+            restored(frame1) {
+              // Now both fields should appear modified.
+              instance.value.expectEqual("modified in suspended transaction")
+              instance.nonTransactedValue.expectEqual("modified in suspended transaction")
+            }
+
+            // Non-transacted values should be accessible outside a frame
+            instance.nonTransactedValue.expectEqual("modified in suspended transaction")
+
+            frame {
+              // New frames should see both fields as modified.
+              instance.value.expectEqual("modified in suspended transaction")
+              instance.nonTransactedValue.expectEqual("modified in suspended transaction")
+            }
+          }
+        }
+    """)
+
+    override fun helperFiles(): List<KtFile> = listOf(sourceFile("Helpers.kt",
+        HELPERS
+    ))
+}
+
+const val HELPERS = """
+    import androidx.compose.frames.open
+    import androidx.compose.frames.commit
+    import androidx.compose.frames.suspend
+    import androidx.compose.frames.restore
+    import androidx.compose.frames.Frame
+
+    inline fun <T> frame(crossinline block: ()->T): T {
+        open(false)
+        try {
+          return block()
+        } finally {
+          commit()
+        }
+    }
+
+    inline fun suspended(crossinline block: ()->Unit): Frame {
+      open(false)
+      block()
+      return suspend()
+    }
+
+    inline fun restored(frame: Frame, crossinline block: ()->Unit) {
+      restore(frame)
+      block()
+      commit()
+    }
+
+    inline fun continued(frame: Frame, crossinline block: ()->Unit): Frame {
+      restore(frame)
+      block()
+      return suspend()
+    }
+
+    fun Any.expectEqual(expected: Any) {
+      expect(expected, this)
+    }
+
+    fun expect(expected: Any, received: Any) {
+      if (expected != received) {
+        throw Exception("Expected \"${'$'}expected\" but received \"${'$'}received\"")
+      }
+    }"""

plugins/compose-compiler-hosted/src/main/java/androidx/compose/plugins/kotlin/ComposableAnnotationChecker.kt

@@ -0,0 +1,565 @@
+/*
+ * Copyright 2019 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.compose.plugins.kotlin
+
+import com.intellij.openapi.project.Project
+import com.intellij.psi.PsiElement
+import org.jetbrains.kotlin.container.StorageComponentContainer
+import org.jetbrains.kotlin.container.useInstance
+import org.jetbrains.kotlin.descriptors.ClassDescriptor
+import org.jetbrains.kotlin.descriptors.ConstructorDescriptor
+import org.jetbrains.kotlin.descriptors.DeclarationDescriptor
+import org.jetbrains.kotlin.descriptors.FunctionDescriptor
+import org.jetbrains.kotlin.descriptors.ModuleDescriptor
+import org.jetbrains.kotlin.descriptors.PropertyDescriptor
+import org.jetbrains.kotlin.descriptors.TypeAliasDescriptor
+import org.jetbrains.kotlin.descriptors.ValueParameterDescriptor
+import org.jetbrains.kotlin.descriptors.VariableDescriptor
+import org.jetbrains.kotlin.descriptors.annotations.KotlinTarget
+import org.jetbrains.kotlin.descriptors.impl.LocalVariableDescriptor
+import org.jetbrains.kotlin.diagnostics.Errors
+import org.jetbrains.kotlin.diagnostics.reportFromPlugin
+import org.jetbrains.kotlin.extensions.StorageComponentContainerContributor
+import org.jetbrains.kotlin.js.resolve.diagnostics.findPsi
+import org.jetbrains.kotlin.load.java.descriptors.JavaMethodDescriptor
+import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlin.psi.KtAnnotatedExpression
+import org.jetbrains.kotlin.psi.KtAnnotationEntry
+import org.jetbrains.kotlin.psi.KtCallExpression
+import org.jetbrains.kotlin.psi.KtClass
+import org.jetbrains.kotlin.psi.KtDeclaration
+import org.jetbrains.kotlin.psi.KtElement
+import org.jetbrains.kotlin.psi.KtExpression
+import org.jetbrains.kotlin.psi.KtFunction
+import org.jetbrains.kotlin.psi.KtLambdaExpression
+import org.jetbrains.kotlin.psi.KtNamedFunction
+import org.jetbrains.kotlin.psi.KtObjectLiteralExpression
+import org.jetbrains.kotlin.psi.KtParameter
+import org.jetbrains.kotlin.psi.KtProperty
+import org.jetbrains.kotlin.psi.KtTreeVisitorVoid
+import androidx.compose.plugins.kotlin.analysis.ComposeDefaultErrorMessages
+import androidx.compose.plugins.kotlin.analysis.ComposeErrors
+import androidx.compose.plugins.kotlin.analysis.ComposeWritableSlices.COMPOSABLE_ANALYSIS
+import androidx.compose.plugins.kotlin.analysis.ComposeWritableSlices.FCS_RESOLVEDCALL_COMPOSABLE
+import androidx.compose.plugins.kotlin.analysis.ComposeWritableSlices.INFERRED_COMPOSABLE_DESCRIPTOR
+import org.jetbrains.kotlin.descriptors.PropertyGetterDescriptor
+import org.jetbrains.kotlin.descriptors.impl.AnonymousFunctionDescriptor
+import org.jetbrains.kotlin.resolve.AdditionalAnnotationChecker
+import org.jetbrains.kotlin.resolve.BindingContext
+import org.jetbrains.kotlin.resolve.BindingTrace
+import org.jetbrains.kotlin.platform.TargetPlatform
+import org.jetbrains.kotlin.platform.jvm.isJvm
+import org.jetbrains.kotlin.psi.KtFunctionLiteral
+import org.jetbrains.kotlin.psi.KtLambdaArgument
+import org.jetbrains.kotlin.psi.KtPropertyAccessor
+import org.jetbrains.kotlin.psi.KtSimpleNameExpression
+import org.jetbrains.kotlin.resolve.calls.callUtil.getResolvedCall
+import org.jetbrains.kotlin.resolve.calls.checkers.AdditionalTypeChecker
+import org.jetbrains.kotlin.resolve.calls.checkers.CallChecker
+import org.jetbrains.kotlin.resolve.calls.checkers.CallCheckerContext
+import org.jetbrains.kotlin.resolve.calls.context.ResolutionContext
+import org.jetbrains.kotlin.resolve.calls.model.ResolvedCall
+import org.jetbrains.kotlin.resolve.calls.model.VariableAsFunctionResolvedCall
+import org.jetbrains.kotlin.resolve.checkers.DeclarationChecker
+import org.jetbrains.kotlin.resolve.checkers.DeclarationCheckerContext
+import org.jetbrains.kotlin.resolve.inline.InlineUtil.isInlinedArgument
+import org.jetbrains.kotlin.types.KotlinType
+import org.jetbrains.kotlin.types.TypeUtils
+import org.jetbrains.kotlin.types.lowerIfFlexible
+import org.jetbrains.kotlin.types.typeUtil.builtIns
+import org.jetbrains.kotlin.types.upperIfFlexible
+import org.jetbrains.kotlin.util.OperatorNameConventions
+import org.jetbrains.kotlin.utils.addToStdlib.cast
+
+open class ComposableAnnotationChecker : CallChecker, DeclarationChecker,
+    AdditionalTypeChecker, AdditionalAnnotationChecker, StorageComponentContainerContributor {
+
+    companion object {
+        fun get(project: Project): ComposableAnnotationChecker {
+            return StorageComponentContainerContributor.getInstances(project).single {
+                it is ComposableAnnotationChecker
+            } as ComposableAnnotationChecker
+        }
+    }
+
+    enum class Composability { NOT_COMPOSABLE, INFERRED, MARKED }
+
+    fun shouldInvokeAsTag(trace: BindingTrace, resolvedCall: ResolvedCall<*>): Boolean {
+        if (resolvedCall is VariableAsFunctionResolvedCall) {
+            if (resolvedCall.variableCall.candidateDescriptor.type.hasComposableAnnotation())
+                return true
+            if (resolvedCall.functionCall.resultingDescriptor.hasComposableAnnotation()) return true
+            return false
+        }
+        val candidateDescriptor = resolvedCall.candidateDescriptor
+        if (candidateDescriptor is FunctionDescriptor) {
+            if (candidateDescriptor.isOperator &&
+                candidateDescriptor.name == OperatorNameConventions.INVOKE) {
+                if (resolvedCall.dispatchReceiver?.type?.hasComposableAnnotation() == true) {
+                    return true
+                }
+            }
+        }
+        if (candidateDescriptor is FunctionDescriptor) {
+            when (analyze(trace, candidateDescriptor)) {
+                Composability.NOT_COMPOSABLE -> return false
+                Composability.INFERRED -> return true
+                Composability.MARKED -> return true
+            }
+        }
+        if (candidateDescriptor is ValueParameterDescriptor) {
+            return candidateDescriptor.type.hasComposableAnnotation()
+        }
+        if (candidateDescriptor is LocalVariableDescriptor) {
+            return candidateDescriptor.type.hasComposableAnnotation()
+        }
+        if (candidateDescriptor is PropertyDescriptor) {
+            return candidateDescriptor.hasComposableAnnotation()
+        }
+        return candidateDescriptor.hasComposableAnnotation()
+    }
+
+    fun analyze(trace: BindingTrace, descriptor: FunctionDescriptor): Composability {
+        val unwrappedDescriptor = when (descriptor) {
+            is ComposableFunctionDescriptor -> descriptor.underlyingDescriptor
+            else -> descriptor
+        }
+        val psi = unwrappedDescriptor.findPsi() as? KtElement
+        psi?.let { trace.bindingContext.get(COMPOSABLE_ANALYSIS, it)?.let { return it } }
+        if (unwrappedDescriptor.name == Name.identifier("compose") &&
+            unwrappedDescriptor.containingDeclaration is ClassDescriptor &&
+            ComposeUtils.isComposeComponent(unwrappedDescriptor.containingDeclaration)
+        ) return Composability.MARKED
+        var composability = Composability.NOT_COMPOSABLE
+        if (trace.bindingContext.get(
+                INFERRED_COMPOSABLE_DESCRIPTOR,
+                unwrappedDescriptor
+            ) ?: false) {
+            composability = Composability.MARKED
+        } else {
+            when (unwrappedDescriptor) {
+                is VariableDescriptor ->
+                    if (unwrappedDescriptor.hasComposableAnnotation() ||
+                        unwrappedDescriptor.type.hasComposableAnnotation()
+                    )
+                        composability =
+                            Composability.MARKED
+                is ConstructorDescriptor ->
+                    if (unwrappedDescriptor.hasComposableAnnotation()) composability =
+                        Composability.MARKED
+                is JavaMethodDescriptor ->
+                    if (unwrappedDescriptor.hasComposableAnnotation()) composability =
+                        Composability.MARKED
+                is AnonymousFunctionDescriptor -> {
+                    if (unwrappedDescriptor.hasComposableAnnotation()) composability =
+                        Composability.MARKED
+                    if (psi is KtFunctionLiteral && psi.isEmitInline(trace.bindingContext)) {
+                        composability = Composability.MARKED
+                    }
+                }
+                is PropertyGetterDescriptor ->
+                    if (unwrappedDescriptor.correspondingProperty.hasComposableAnnotation())
+                        composability = Composability.MARKED
+                else -> if (unwrappedDescriptor.hasComposableAnnotation()) composability =
+                    Composability.MARKED
+            }
+        }
+        (unwrappedDescriptor.findPsi() as? KtElement)?.let {
+                element -> composability = analyzeFunctionContents(trace, element, composability)
+        }
+        psi?.let { trace.record(COMPOSABLE_ANALYSIS, it, composability) }
+        return composability
+    }
+
+    private fun analyzeFunctionContents(
+        trace: BindingTrace,
+        element: KtElement,
+        signatureComposability: Composability
+    ): Composability {
+        var composability = signatureComposability
+        var localFcs = false
+        var isInlineLambda = false
+        element.accept(object : KtTreeVisitorVoid() {
+            override fun visitNamedFunction(function: KtNamedFunction) {
+                if (function == element) {
+                    super.visitNamedFunction(function)
+                }
+            }
+
+            override fun visitPropertyAccessor(accessor: KtPropertyAccessor) {
+                // this is basically a function, so unless it is the function we are analyzing, we
+                // stop here
+                if (accessor == element) {
+                    super.visitPropertyAccessor(accessor)
+                }
+            }
+
+            override fun visitClass(klass: KtClass) {
+                // never traverse a class
+            }
+
+            override fun visitLambdaExpression(lambdaExpression: KtLambdaExpression) {
+                val isInlineable = isInlinedArgument(
+                    lambdaExpression.functionLiteral,
+                    trace.bindingContext,
+                    true
+                )
+                if (isInlineable && lambdaExpression == element) isInlineLambda = true
+                if (isInlineable || lambdaExpression == element)
+                    super.visitLambdaExpression(lambdaExpression)
+            }
+
+            override fun visitSimpleNameExpression(expression: KtSimpleNameExpression) {
+                val resolvedCall = expression.getResolvedCall(trace.bindingContext)
+                if (resolvedCall?.candidateDescriptor is PropertyDescriptor) {
+                    checkResolvedCall(
+                        resolvedCall,
+                        trace.get(FCS_RESOLVEDCALL_COMPOSABLE, expression),
+                        expression
+                    )
+                }
+                super.visitSimpleNameExpression(expression)
+            }
+
+            override fun visitCallExpression(expression: KtCallExpression) {
+                val resolvedCall = expression.getResolvedCall(trace.bindingContext)
+                checkResolvedCall(
+                    resolvedCall,
+                    trace.get(FCS_RESOLVEDCALL_COMPOSABLE, expression),
+                    expression.calleeExpression ?: expression
+                )
+                super.visitCallExpression(expression)
+            }
+
+            private fun checkResolvedCall(
+                resolvedCall: ResolvedCall<*>?,
+                isCallComposable: Boolean?,
+                reportElement: KtExpression
+            ) {
+                when (resolvedCall?.candidateDescriptor) {
+                    is ComposableEmitDescriptor,
+                    is ComposablePropertyDescriptor,
+                    is ComposableFunctionDescriptor -> {
+                        localFcs = true
+                        if (!isInlineLambda && composability != Composability.MARKED) {
+                            // Report error on composable element to make it obvious which invocation is offensive
+                            trace.report(
+                                ComposeErrors.COMPOSABLE_INVOCATION_IN_NON_COMPOSABLE
+                                    .on(reportElement)
+                            )
+                        }
+                    }
+                }
+                // Can be null in cases where the call isn't resolvable (eg. due to a bug/typo in the user's code)
+                if (isCallComposable == true) {
+                    localFcs = true
+                    if (!isInlineLambda && composability != Composability.MARKED) {
+                        // Report error on composable element to make it obvious which invocation is offensive
+                        trace.report(
+                            ComposeErrors.COMPOSABLE_INVOCATION_IN_NON_COMPOSABLE
+                                .on(reportElement)
+                        )
+                    }
+                }
+            }
+        }, null)
+        if (
+            localFcs &&
+            !isInlineLambda && composability != Composability.MARKED
+        ) {
+            val reportElement = when (element) {
+                is KtNamedFunction -> element.nameIdentifier ?: element
+                else -> element
+            }
+            if (localFcs) {
+                trace.report(
+                    ComposeErrors.COMPOSABLE_INVOCATION_IN_NON_COMPOSABLE.on(reportElement)
+                )
+            }
+        }
+        if (localFcs && composability == Composability.NOT_COMPOSABLE)
+            composability =
+                Composability.INFERRED
+        return composability
+    }
+
+    /**
+     * Analyze a KtElement
+     *  - Determine if it is @Composable (eg. the element or inferred type has an @Composable annotation)
+     *  - Update the binding context to cache analysis results
+     *  - Report errors (eg. invocations of an @Composable, etc)
+     *  - Return true if element is @Composable, else false
+     */
+    fun analyze(trace: BindingTrace, element: KtElement, type: KotlinType?): Composability {
+        trace.bindingContext.get(COMPOSABLE_ANALYSIS, element)?.let { return it }
+
+        var composability =
+            Composability.NOT_COMPOSABLE
+
+        if (element is KtClass) {
+            val descriptor = trace.bindingContext.get(BindingContext.CLASS, element)
+                ?: error("Element class context not found")
+            val annotationEntry = element.annotationEntries.singleOrNull {
+                trace.bindingContext.get(BindingContext.ANNOTATION, it)?.isComposableAnnotation
+                    ?: false
+            }
+            if (annotationEntry != null && !ComposeUtils.isComposeComponent(descriptor)) {
+                trace.report(
+                    Errors.WRONG_ANNOTATION_TARGET.on(
+                        annotationEntry,
+                        "class which does not extend androidx.compose.Component"
+                    )
+                )
+            }
+            if (ComposeUtils.isComposeComponent(descriptor)) {
+                composability += Composability.MARKED
+            }
+        }
+        if (element is KtParameter) {
+            val composableAnnotation = element
+                .typeReference
+                ?.annotationEntries
+                ?.mapNotNull { trace.bindingContext.get(BindingContext.ANNOTATION, it) }
+                ?.singleOrNull { it.isComposableAnnotation }
+
+            if (composableAnnotation != null) {
+                composability += Composability.MARKED
+            }
+        }
+        if (element is KtParameter) {
+            val composableAnnotation = element
+                .typeReference
+                ?.annotationEntries
+                ?.mapNotNull { trace.bindingContext.get(BindingContext.ANNOTATION, it) }
+                ?.singleOrNull { it.isComposableAnnotation }
+
+            if (composableAnnotation != null) {
+                composability += Composability.MARKED
+            }
+        }
+
+        // if (candidateDescriptor.type.arguments.size != 1 || !candidateDescriptor.type.arguments[0].type.isUnit()) return false
+        if (
+            type != null &&
+            type !== TypeUtils.NO_EXPECTED_TYPE &&
+            type.hasComposableAnnotation()
+        ) {
+            composability += Composability.MARKED
+        }
+        val parent = element.parent
+        val annotations = when {
+            element is KtNamedFunction -> element.annotationEntries
+            parent is KtAnnotatedExpression -> parent.annotationEntries
+            element is KtProperty -> element.annotationEntries
+            element is KtParameter -> element.typeReference?.annotationEntries ?: emptyList()
+            else -> emptyList()
+        }
+
+        for (entry in annotations) {
+            val descriptor = trace.bindingContext.get(BindingContext.ANNOTATION, entry) ?: continue
+            if (descriptor.isComposableAnnotation) {
+                composability += Composability.MARKED
+            }
+        }
+
+        if (element is KtLambdaExpression || element is KtFunction) {
+            val associatedCall = parent?.parent as? KtCallExpression
+
+            if (associatedCall != null && parent is KtLambdaArgument) {
+                val resolvedCall = associatedCall.getResolvedCall(trace.bindingContext)
+                if (resolvedCall?.candidateDescriptor is ComposableEmitDescriptor) {
+                    composability += Composability.MARKED
+                }
+            }
+
+            composability = analyzeFunctionContents(trace, element, composability)
+        }
+
+        trace.record(COMPOSABLE_ANALYSIS, element, composability)
+        return composability
+    }
+
+    override fun registerModuleComponents(
+        container: StorageComponentContainer,
+        platform: TargetPlatform,
+        moduleDescriptor: ModuleDescriptor
+    ) {
+        if (!platform.isJvm()) return
+        container.useInstance(this)
+    }
+
+    override fun check(
+        declaration: KtDeclaration,
+        descriptor: DeclarationDescriptor,
+        context: DeclarationCheckerContext
+    ) {
+        when (descriptor) {
+            is ClassDescriptor -> {
+                val trace = context.trace
+                val element = descriptor.findPsi()
+                if (element is KtClass) {
+                    val classDescriptor =
+                        trace.bindingContext.get(
+                            BindingContext.CLASS,
+                            element
+                        ) ?: error("Element class context not found")
+                    val composableAnnotationEntry = element.annotationEntries.singleOrNull {
+                        trace.bindingContext.get(
+                            BindingContext.ANNOTATION,
+                            it
+                        )?.isComposableAnnotation ?: false
+                    }
+                    if (composableAnnotationEntry != null &&
+                        !ComposeUtils.isComposeComponent(classDescriptor)) {
+                        trace.report(
+                            Errors.WRONG_ANNOTATION_TARGET.on(
+                                composableAnnotationEntry,
+                                "class which does not extend androidx.compose.Component"
+                            )
+                        )
+                    }
+                }
+            }
+            is PropertyDescriptor -> {}
+            is LocalVariableDescriptor -> {}
+            is TypeAliasDescriptor -> {}
+            is FunctionDescriptor -> analyze(context.trace, descriptor)
+            else ->
+                throw Error("currently unsupported " + descriptor.javaClass)
+        }
+    }
+
+    override fun check(
+        resolvedCall: ResolvedCall<*>,
+        reportOn: PsiElement,
+        context: CallCheckerContext
+    ) {
+        val shouldBeTag = shouldInvokeAsTag(context.trace, resolvedCall)
+        context.trace.record(
+            FCS_RESOLVEDCALL_COMPOSABLE,
+            resolvedCall.call.callElement,
+            shouldBeTag
+        )
+    }
+
+    override fun checkType(
+        expression: KtExpression,
+        expressionType: KotlinType,
+        expressionTypeWithSmartCast: KotlinType,
+        c: ResolutionContext<*>
+    ) {
+        if (expression is KtLambdaExpression) {
+            val expectedType = c.expectedType
+            if (expectedType === TypeUtils.NO_EXPECTED_TYPE) return
+            val expectedComposable = expectedType.hasComposableAnnotation()
+            val composability = analyze(c.trace, expression, c.expectedType)
+            if ((expectedComposable && composability == Composability.NOT_COMPOSABLE) ||
+                (!expectedComposable && composability == Composability.MARKED)) {
+                val isInlineable =
+                    isInlinedArgument(
+                        expression.functionLiteral,
+                        c.trace.bindingContext,
+                        true
+                    )
+                if (isInlineable) return
+
+                if (expression.parent is KtLambdaArgument) {
+                    val callDescriptor = expression
+                        .parent
+                        ?.parent
+                        ?.cast<KtCallExpression>()
+                        ?.getResolvedCall(c.trace.bindingContext)
+                        ?.candidateDescriptor
+
+                    if (callDescriptor is ComposableEmitDescriptor) {
+                        return
+                    }
+                }
+
+                val reportOn =
+                    if (expression.parent is KtAnnotatedExpression)
+                        expression.parent as KtExpression
+                    else expression
+                c.trace.report(
+                    Errors.TYPE_MISMATCH.on(
+                        reportOn,
+                        expectedType,
+                        expressionTypeWithSmartCast
+                    )
+                )
+            }
+            return
+        } else {
+            val expectedType = c.expectedType
+
+            if (expectedType === TypeUtils.NO_EXPECTED_TYPE) return
+            if (expectedType === TypeUtils.UNIT_EXPECTED_TYPE) return
+
+            val nullableAnyType = expectedType.builtIns.nullableAnyType
+            val anyType = expectedType.builtIns.anyType
+
+            if (anyType == expectedType.lowerIfFlexible() &&
+                nullableAnyType == expectedType.upperIfFlexible()) return
+
+            val nullableNothingType = expectedType.builtIns.nullableNothingType
+
+            // Handle assigning null to a nullable composable type
+            if (expectedType.isMarkedNullable &&
+                expressionTypeWithSmartCast == nullableNothingType) return
+
+            val expectedComposable = expectedType.hasComposableAnnotation()
+            val isComposable = expressionType.hasComposableAnnotation()
+
+            if (expectedComposable != isComposable) {
+                val reportOn =
+                    if (expression.parent is KtAnnotatedExpression)
+                        expression.parent as KtExpression
+                    else expression
+                c.trace.report(
+                    Errors.TYPE_MISMATCH.on(
+                        reportOn,
+                        expectedType,
+                        expressionTypeWithSmartCast
+                    )
+                )
+            }
+            return
+        }
+    }
+
+    override fun checkEntries(
+        entries: List<KtAnnotationEntry>,
+        actualTargets: List<KotlinTarget>,
+        trace: BindingTrace
+    ) {
+        val entry = entries.singleOrNull {
+            trace.bindingContext.get(BindingContext.ANNOTATION, it)?.isComposableAnnotation ?: false
+        }
+        if ((entry?.parent as? KtAnnotatedExpression)?.baseExpression is
+                    KtObjectLiteralExpression) {
+            trace.report(
+                Errors.WRONG_ANNOTATION_TARGET.on(
+                    entry,
+                    "class which does not extend androidx.compose.Component"
+                )
+            )
+        }
+    }
+
+    operator fun Composability.plus(rhs: Composability): Composability =
+        if (this > rhs) this else rhs
+}

plugins/compose-compiler-hosted/src/main/java/androidx/compose/plugins/kotlin/ComposableEmitDescriptor.kt

@@ -0,0 +1,211 @@
+/*
+ * Copyright 2019 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.compose.plugins.kotlin
+
+import org.jetbrains.kotlin.builtins.DefaultBuiltIns
+import org.jetbrains.kotlin.descriptors.CallableDescriptor
+import org.jetbrains.kotlin.descriptors.CallableMemberDescriptor
+import org.jetbrains.kotlin.descriptors.DeclarationDescriptor
+import org.jetbrains.kotlin.descriptors.Modality
+import org.jetbrains.kotlin.descriptors.SimpleFunctionDescriptor
+import org.jetbrains.kotlin.descriptors.SourceElement
+import org.jetbrains.kotlin.descriptors.ValueParameterDescriptor
+import org.jetbrains.kotlin.descriptors.Visibilities
+import org.jetbrains.kotlin.descriptors.annotations.Annotations
+import org.jetbrains.kotlin.descriptors.impl.SimpleFunctionDescriptorImpl
+import org.jetbrains.kotlin.descriptors.impl.ValueParameterDescriptorImpl
+import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlin.resolve.calls.model.ResolvedCall
+import org.jetbrains.kotlin.types.KotlinType
+import org.jetbrains.kotlin.types.replace
+import org.jetbrains.kotlin.types.typeUtil.asTypeProjection
+
+interface ComposableEmitMetadata {
+    val composerMetadata: ComposerMetadata
+    val emitCall: ResolvedCall<*>
+    val hasChildren: Boolean
+    val pivotals: List<String>
+    val ctorCall: ResolvedCall<*>
+    val ctorParams: List<String>
+    val validations: List<ValidatedAssignment>
+}
+
+class ComposableEmitDescriptor(
+    override val composerMetadata: ComposerMetadata,
+    override val emitCall: ResolvedCall<*>,
+    override val hasChildren: Boolean,
+    override val pivotals: List<String>,
+    override val ctorCall: ResolvedCall<*>,
+    override val ctorParams: List<String>,
+    override val validations: List<ValidatedAssignment>,
+    containingDeclaration: DeclarationDescriptor,
+    original: SimpleFunctionDescriptor?,
+    annotations: Annotations,
+    name: Name,
+    kind: CallableMemberDescriptor.Kind,
+    source: SourceElement
+) : ComposableEmitMetadata, SimpleFunctionDescriptorImpl(
+    containingDeclaration,
+    original,
+    annotations,
+    name,
+    kind,
+    source
+) {
+
+    companion object {
+        fun build(
+            hasChildren: Boolean,
+            emitCall: ResolvedCall<*>,
+            pivotals: List<String>,
+            ctorCall: ResolvedCall<*>,
+            ctorParams: List<String>,
+            validations: List<ValidatedAssignment>,
+            composerMetadata: ComposerMetadata,
+            name: Name
+        ): ComposableEmitDescriptor {
+
+            val builtIns = DefaultBuiltIns.Instance
+
+            val resolvedCall = ctorCall
+
+            val original = resolvedCall.resultingDescriptor as? SimpleFunctionDescriptor
+
+            val descriptor = ComposableEmitDescriptor(
+                composerMetadata,
+                emitCall,
+                hasChildren,
+                pivotals,
+                ctorCall,
+                ctorParams,
+                validations,
+                emitCall.candidateDescriptor.containingDeclaration,
+                original,
+                Annotations.EMPTY,
+                name,
+                CallableMemberDescriptor.Kind.SYNTHESIZED,
+                SourceElement.NO_SOURCE
+            )
+
+            val valueArgs = mutableListOf<ValueParameterDescriptor>()
+            val paramSet = mutableSetOf<String>()
+
+            for (paramName in ctorParams) {
+                if (paramSet.contains(paramName)) continue
+                val param = resolvedCall.resultingDescriptor.valueParameters.find {
+                    it.name.identifier == paramName
+                } ?: continue
+
+                paramSet.add(paramName)
+                valueArgs.add(
+                    ValueParameterDescriptorImpl(
+                        descriptor, null, valueArgs.size,
+                        Annotations.EMPTY,
+                        param.name,
+                        param.type, false,
+                        false,
+                        false, null,
+                        SourceElement.NO_SOURCE
+                    )
+                )
+            }
+
+            for (validation in validations) {
+                if (paramSet.contains(validation.name)) continue
+                paramSet.add(validation.name)
+                valueArgs.add(
+                    ValueParameterDescriptorImpl(
+                        descriptor,
+                        null,
+                        valueArgs.size,
+                        Annotations.EMPTY,
+                        Name.identifier(validation.name),
+                        validation.type,
+                        false,
+                        false,
+                        false,
+                        null,
+                        SourceElement.NO_SOURCE
+                    )
+                )
+            }
+
+            val unitLambdaType = builtIns.getFunction(
+                0
+            ).defaultType.replace(
+                listOf(builtIns.unitType.asTypeProjection())
+            )
+            // NOTE(lmr): it's actually kind of important that this is *not* a composable lambda,
+            // so that the observe patcher doesn't insert an observe scope.
+            // In the future, we should reconsider how this is done since semantically a composable
+            // lambda is more correct here. I tried, but had trouble passing enough information to
+            // the observe patcher so it knew not to do this.
+            /*.makeComposable(scopeTower.module)*/
+            if (hasChildren) {
+                valueArgs.add(
+                    EmitChildrenValueParameterDescriptor(
+                        descriptor, null, valueArgs.size,
+                        Annotations.EMPTY,
+                        Name.identifier("\$CHILDREN"),
+                        unitLambdaType, false,
+                        false,
+                        false, null,
+                        SourceElement.NO_SOURCE
+                    )
+                )
+            }
+
+            descriptor.initialize(
+                null,
+                null,
+                mutableListOf(),
+                valueArgs,
+                builtIns.unitType,
+                Modality.FINAL,
+                Visibilities.DEFAULT_VISIBILITY
+            )
+
+            return descriptor
+        }
+    }
+}
+
+class EmitChildrenValueParameterDescriptor(
+    containingDeclaration: CallableDescriptor,
+    original: ValueParameterDescriptor?,
+    index: Int,
+    annotations: Annotations,
+    name: Name,
+    outType: KotlinType,
+    declaresDefaultValue: Boolean,
+    isCrossinline: Boolean,
+    isNoinline: Boolean,
+    varargElementType: KotlinType?,
+    source: SourceElement
+) : ValueParameterDescriptorImpl(
+    containingDeclaration,
+    original,
+    index,
+    annotations,
+    name,
+    outType,
+    declaresDefaultValue,
+    isCrossinline,
+    isNoinline,
+    varargElementType,
+    source
+)

plugins/compose-compiler-hosted/src/main/java/androidx/compose/plugins/kotlin/ComposableFunctionDescriptor.kt

@@ -0,0 +1,76 @@
+/*
+ * Copyright 2019 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.compose.plugins.kotlin
+
+import org.jetbrains.kotlin.descriptors.CallableDescriptor
+import org.jetbrains.kotlin.descriptors.FunctionDescriptor
+import org.jetbrains.kotlin.descriptors.PropertyDescriptor
+import org.jetbrains.kotlin.descriptors.SimpleFunctionDescriptor
+import org.jetbrains.kotlin.resolve.calls.model.ResolvedCall
+import org.jetbrains.kotlin.types.TypeSubstitutor
+
+interface ComposableCallableDescriptor : CallableDescriptor {
+    val underlyingDescriptor: CallableDescriptor
+}
+
+interface ComposableFunctionDescriptor : FunctionDescriptor, ComposableCallableDescriptor {
+    override val underlyingDescriptor: FunctionDescriptor
+}
+
+interface ComposablePropertyDescriptor : PropertyDescriptor, ComposableCallableDescriptor {
+    override val underlyingDescriptor: PropertyDescriptor
+}
+
+class ComposablePropertyDescriptorImpl(
+    override val underlyingDescriptor: PropertyDescriptor
+) : PropertyDescriptor by underlyingDescriptor, ComposablePropertyDescriptor {
+    override fun substitute(substitutor: TypeSubstitutor): PropertyDescriptor? {
+        return underlyingDescriptor.substitute(substitutor)?.let {
+            ComposablePropertyDescriptorImpl(it)
+        }
+    }
+}
+
+fun ComposableFunctionDescriptor(
+    underlyingDescriptor: FunctionDescriptor
+): ComposableFunctionDescriptor {
+    return if (underlyingDescriptor is SimpleFunctionDescriptor) {
+        ComposableSimpleFunctionDescriptorImpl(underlyingDescriptor)
+    } else {
+        ComposableFunctionDescriptorImpl(underlyingDescriptor)
+    }
+}
+
+class ComposableFunctionDescriptorImpl(
+    override val underlyingDescriptor: FunctionDescriptor
+) : FunctionDescriptor by underlyingDescriptor, ComposableFunctionDescriptor {
+    override fun substitute(substitutor: TypeSubstitutor): FunctionDescriptor? {
+        return underlyingDescriptor.substitute(substitutor)?.let {
+            ComposableFunctionDescriptor(it)
+        }
+    }
+}
+
+class ComposableSimpleFunctionDescriptorImpl(
+    override val underlyingDescriptor: SimpleFunctionDescriptor
+) : SimpleFunctionDescriptor by underlyingDescriptor, ComposableFunctionDescriptor {
+    override fun substitute(substitutor: TypeSubstitutor): FunctionDescriptor? {
+        return underlyingDescriptor.substitute(substitutor)?.let {
+            ComposableFunctionDescriptor(it)
+        }
+    }
+}

plugins/compose-compiler-hosted/src/main/java/androidx/compose/plugins/kotlin/ComposeCallResolutionInterceptorExtension.kt

@@ -0,0 +1,397 @@
+/*
+ * Copyright 2019 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.compose.plugins.kotlin
+
+import androidx.compose.plugins.kotlin.analysis.ComposeWritableSlices
+import com.intellij.psi.PsiElement
+import org.jetbrains.kotlin.descriptors.CallableDescriptor
+import org.jetbrains.kotlin.descriptors.ClassDescriptor
+import org.jetbrains.kotlin.descriptors.ConstructorDescriptor
+import org.jetbrains.kotlin.descriptors.FunctionDescriptor
+import org.jetbrains.kotlin.descriptors.Modality
+import org.jetbrains.kotlin.descriptors.PropertyDescriptor
+import org.jetbrains.kotlin.descriptors.findClassAcrossModuleDependencies
+import org.jetbrains.kotlin.extensions.internal.CallResolutionInterceptorExtension
+import org.jetbrains.kotlin.incremental.components.LookupLocation
+import org.jetbrains.kotlin.name.ClassId
+import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlin.psi.KtExpression
+import org.jetbrains.kotlin.psi.KtFunction
+import org.jetbrains.kotlin.psi.KtPsiFactory
+import org.jetbrains.kotlin.resolve.BindingContext
+import org.jetbrains.kotlin.resolve.calls.CallResolver
+import org.jetbrains.kotlin.resolve.calls.CandidateResolver
+import org.jetbrains.kotlin.resolve.calls.context.BasicCallResolutionContext
+import org.jetbrains.kotlin.resolve.calls.context.CheckArgumentTypesMode
+import org.jetbrains.kotlin.resolve.calls.context.TemporaryTraceAndCache
+import org.jetbrains.kotlin.resolve.calls.model.DataFlowInfoForArgumentsImpl
+import org.jetbrains.kotlin.resolve.calls.model.MutableResolvedCall
+import org.jetbrains.kotlin.resolve.calls.model.ResolvedCall
+import org.jetbrains.kotlin.resolve.calls.model.VariableAsFunctionMutableResolvedCall
+import org.jetbrains.kotlin.resolve.calls.model.VariableAsFunctionResolvedCall
+import org.jetbrains.kotlin.resolve.calls.model.VariableAsFunctionResolvedCallImpl
+import org.jetbrains.kotlin.resolve.calls.results.OverloadResolutionResults
+import org.jetbrains.kotlin.resolve.calls.tasks.TracingStrategy
+import org.jetbrains.kotlin.resolve.calls.tower.ImplicitScopeTower
+import org.jetbrains.kotlin.resolve.calls.tower.NewResolutionOldInference
+import org.jetbrains.kotlin.resolve.descriptorUtil.module
+import org.jetbrains.kotlin.resolve.inline.InlineUtil
+import org.jetbrains.kotlin.resolve.scopes.ResolutionScope
+import org.jetbrains.kotlin.types.KotlinType
+
+typealias Candidate = NewResolutionOldInference.MyCandidate
+
+fun ComposableCandidate(candidate: Candidate): Candidate {
+    val (eagerDiagnostics, resolvedCall) = candidate
+
+    if (resolvedCall !is VariableAsFunctionMutableResolvedCall) {
+        @Suppress("UNCHECKED_CAST")
+        return Candidate(
+            eagerDiagnostics = eagerDiagnostics,
+            resolvedCall = ComposableResolvedCall(resolvedCall),
+            finalDiagnosticsComputation = null
+        )
+    }
+
+    val functionCall = ComposableResolvedCall(resolvedCall.functionCall)
+
+    val variableCall = resolvedCall.variableCall
+
+    val newCall = VariableAsFunctionResolvedCallImpl(functionCall, variableCall)
+
+    return Candidate(
+        eagerDiagnostics = eagerDiagnostics,
+        resolvedCall = newCall,
+        finalDiagnosticsComputation = null
+    )
+}
+
+@Suppress("UNCHECKED_CAST")
+class ComposableResolvedCall<T : CallableDescriptor>(
+    private val underlying: MutableResolvedCall<T>
+) : MutableResolvedCall<T> by underlying {
+    private val composableCandidateDescriptor =
+        when (val descriptor = underlying.candidateDescriptor) {
+            is FunctionDescriptor -> ComposableFunctionDescriptor(descriptor)
+            is PropertyDescriptor -> ComposablePropertyDescriptorImpl(descriptor)
+            else -> error("Expected FunctionDescriptor or PropertyDescriptor, found $descriptor")
+        }
+    override fun getCandidateDescriptor(): T = composableCandidateDescriptor as T
+    override fun getResultingDescriptor(): T {
+        return when (val descriptor = underlying.resultingDescriptor) {
+            is FunctionDescriptor -> ComposableFunctionDescriptor(descriptor)
+            is PropertyDescriptor -> ComposablePropertyDescriptorImpl(descriptor)
+            else -> error("Expected FunctionDescriptor or PropertyDescriptor, found $descriptor")
+        } as T
+    }
+}
+
+@Suppress("INVISIBLE_REFERENCE", "EXPERIMENTAL_IS_NOT_ENABLED")
+@UseExperimental(org.jetbrains.kotlin.extensions.internal.InternalNonStableExtensionPoints::class)
+open class ComposeCallResolutionInterceptorExtension : CallResolutionInterceptorExtension {
+    override fun interceptCandidates(
+        candidates: Collection<Candidate>,
+        context: BasicCallResolutionContext,
+        candidateResolver: CandidateResolver,
+        callResolver: CallResolver?,
+        name: Name,
+        kind: NewResolutionOldInference.ResolutionKind,
+        tracing: TracingStrategy
+    ): Collection<Candidate> {
+        if (callResolver == null) throw IllegalArgumentException("Call resolver must be non-null")
+
+        if (candidates.isEmpty()) return candidates
+        val bindingContext = context.trace.bindingContext
+        val call = context.call
+        val shouldIgnore = bindingContext[
+                ComposeWritableSlices.IGNORE_COMPOSABLE_INTERCEPTION,
+                call
+        ] ?: false
+
+        if (shouldIgnore) return candidates
+
+        val composables = mutableListOf<Candidate>()
+        val nonComposablesNonConstructors = mutableListOf<Candidate>()
+        val alreadyInterceptedCandidates = mutableListOf<Candidate>()
+        val constructors = mutableListOf<Candidate>()
+
+        var needToLookupComposer = false
+
+        for (candidate in candidates) {
+            val resolvedCall = candidate.resolvedCall
+            val candidateDescriptor = resolvedCall.candidateDescriptor
+            when {
+                candidateDescriptor is ComposableFunctionDescriptor -> {
+                    alreadyInterceptedCandidates.add(candidate)
+                }
+                candidateDescriptor is ComposableEmitDescriptor -> {
+                    alreadyInterceptedCandidates.add(candidate)
+                }
+                resolvedCall is VariableAsFunctionResolvedCall &&
+                resolvedCall.variableCall
+                    .candidateDescriptor
+                    .type
+                    .hasComposableAnnotation() -> {
+                    needToLookupComposer = true
+                    composables.add(candidate)
+                }
+                resolvedCall.candidateDescriptor.hasComposableAnnotation() -> {
+                    needToLookupComposer = true
+                    composables.add(candidate)
+                }
+                resolvedCall.candidateDescriptor is ConstructorDescriptor -> {
+                    needToLookupComposer = true
+                    constructors.add(candidate)
+                }
+                else -> nonComposablesNonConstructors.add(candidate)
+            }
+        }
+
+        // If none of the candidates are composable or constructors, then it's unnecessary for us
+        // to do any work at all, since it will never be anything we intercept
+        if (!needToLookupComposer) return candidates
+
+        if (!isInComposableScope(context)) return candidates
+
+        // NOTE(lmr): I'm not implementing emittable interception here, since I believe it is not
+        // needed. So in this case we just pass constructors right through, untouched.
+        return nonComposablesNonConstructors +
+                constructors +
+                alreadyInterceptedCandidates +
+                composables.map { ComposableCandidate(it) }
+    }
+
+    override fun interceptCandidates(
+        candidates: Collection<FunctionDescriptor>,
+        scopeTower: ImplicitScopeTower,
+        resolutionContext: BasicCallResolutionContext,
+        resolutionScope: ResolutionScope,
+        callResolver: CallResolver?,
+        name: Name,
+        location: LookupLocation
+    ): Collection<FunctionDescriptor> {
+        val element = resolutionContext.call.callElement as KtExpression
+        val project = element.project
+        if (callResolver == null) throw IllegalArgumentException("Call resolver must be non-null")
+
+        if (candidates.isEmpty()) return candidates
+        val bindingContext = resolutionContext.trace.bindingContext
+        val call = resolutionContext.call
+        val shouldIgnore = bindingContext[
+                ComposeWritableSlices.IGNORE_COMPOSABLE_INTERCEPTION,
+                call
+        ] ?: false
+
+        if (shouldIgnore) return candidates
+
+        val composables = mutableListOf<FunctionDescriptor>()
+        val nonComposablesNonConstructors = mutableListOf<FunctionDescriptor>()
+        val constructors = mutableListOf<ConstructorDescriptor>()
+
+        var needToLookupComposer = false
+
+        for (candidate in candidates) {
+            when {
+                candidate.hasComposableAnnotation() -> {
+                    needToLookupComposer = true
+                    composables.add(candidate)
+                }
+                candidate is ConstructorDescriptor -> {
+                    needToLookupComposer = true
+                    constructors.add(candidate)
+                }
+                else -> nonComposablesNonConstructors.add(candidate)
+            }
+        }
+
+        // If none of the candidates are composable or constructors, then it's unnecessary for us
+        // to do any work at all, since it will never be anything we intercept
+        if (!needToLookupComposer) return candidates
+
+        // If there are no constructors, then all of the candidates are either composables or
+        // non-composable functions, and we follow normal resolution rules.
+        if (constructors.isEmpty()) {
+            // we wrap the composable descriptors into a ComposableFunctionDescriptor so we know
+            // to intercept it in the backend.
+            return nonComposablesNonConstructors + composables.map {
+                ComposableFunctionDescriptor(it)
+            }
+        }
+
+        if (!isInComposableScope(resolutionContext)) return candidates
+
+        val composerType = callResolver.findComposerCallAndDescriptor(resolutionContext)
+            ?: return nonComposablesNonConstructors + constructors
+
+        val psiFactory = KtPsiFactory(project, markGenerated = false)
+
+        // If we made it this far, we need to check and see if the constructors qualify as emit
+        // calls instead of constructor calls.  First, we need to look at the composer to see
+        // what kinds of "emittables" it accepts.
+        // We cache the metadata into a writeable slice based on the descriptor
+        val composerMetadata = ComposerMetadata.getOrBuild(
+            composerType,
+            callResolver,
+            psiFactory,
+            resolutionContext
+        )
+
+        val emittables = constructors.filter {
+            composerMetadata.isEmittable(it.returnType) && !it.returnType.isAbstract()
+        }
+        val hasEmittableCandidate = emittables.isNotEmpty()
+
+        // if none of the constructors are emittables, then all of the candidates are valid
+        if (!hasEmittableCandidate) {
+            return nonComposablesNonConstructors + constructors + composables.map {
+                ComposableFunctionDescriptor(it)
+            }
+        }
+
+        // since some of the constructors are emittables, we fall back to resolving using the
+        // emit resolver.
+        val emitResolver = ComposeEmitResolver(
+            callResolver,
+            project,
+            composerMetadata
+        )
+
+        val emitCandidates = emitResolver.resolveCandidates(
+            call,
+            emittables,
+            name,
+            resolutionContext
+        )
+
+        return nonComposablesNonConstructors +
+                composables.map {
+                    ComposableFunctionDescriptor(it)
+                } +
+                constructors.filter { !composerMetadata.isEmittable(it.returnType) } +
+                emitCandidates
+    }
+
+    private fun CallResolver.findComposerCallAndDescriptor(
+        context: BasicCallResolutionContext
+    ): KotlinType? {
+        // use the call resolver to find any variable that would resolve with "composer" in scope.
+        return resolveComposer(context)?.resultingDescriptor?.returnType
+            // If there is no composer in scope, then we decide to fall back to the ViewComposer
+            // as a default. When we are properly inferring the composer type, this step should no
+            // longer be needed. This provides a better developer experience currently though since
+            // developers won't be required to import the composer into scope.
+            ?: context
+                .scope
+                .ownerDescriptor
+                .module
+                .findClassAcrossModuleDependencies(ClassId.topLevel(ComposeFqNames.ViewComposer))
+                ?.defaultType
+    }
+
+    private fun isInComposableScope(resolutionContext: BasicCallResolutionContext): Boolean {
+        val call = resolutionContext.call
+        val temporaryTraceForComposeableCall =
+            TemporaryTraceAndCache.create(
+                resolutionContext,
+                "trace to resolve composable call", call.callElement as KtExpression
+            )
+
+        val composableAnnotationChecker =
+            ComposableAnnotationChecker.get(call.callElement.project)
+
+        // Ensure we are in a composable context
+        // TODO(lmr): there ought to be a better way to do this
+        var walker: PsiElement? = call.callElement
+        while (walker != null) {
+            val descriptor = try {
+                resolutionContext.trace[BindingContext.FUNCTION, walker]
+            } catch (e: Exception) {
+                null
+            }
+            if (descriptor != null) {
+                val composability = composableAnnotationChecker.analyze(
+                    temporaryTraceForComposeableCall.trace,
+                    descriptor
+                )
+                if (composability != ComposableAnnotationChecker.Composability.NOT_COMPOSABLE) {
+                    return true
+                }
+
+                // If the descriptor is for an inlined lambda, infer composability from the
+                // outer scope
+                if (!(walker is KtFunction) ||
+                    !InlineUtil.isInlinedArgument(
+                        walker,
+                        resolutionContext.trace.bindingContext,
+                        true
+                    )
+                )
+                    break
+            }
+            walker = try { walker.parent } catch (e: Throwable) { null }
+        }
+        return false
+    }
+
+    private fun CallResolver.resolveVar(
+        name: Name,
+        context: BasicCallResolutionContext
+    ): OverloadResolutionResults<CallableDescriptor> {
+        val temporaryForVariable = TemporaryTraceAndCache.create(
+            context,
+            "trace to resolve variable",
+            context.call.callElement as KtExpression
+        )
+        val call = makeCall(context.call.callElement)
+        val contextForVariable = BasicCallResolutionContext.create(
+            context.replaceTraceAndCache(temporaryForVariable),
+            call,
+            CheckArgumentTypesMode.CHECK_VALUE_ARGUMENTS,
+            DataFlowInfoForArgumentsImpl(context.dataFlowInfo, call)
+
+        )
+        contextForVariable.trace.record(
+            ComposeWritableSlices.IGNORE_COMPOSABLE_INTERCEPTION,
+            call,
+            true
+        )
+        return computeTasksAndResolveCall<CallableDescriptor>(
+            contextForVariable,
+            name,
+            TracingStrategy.EMPTY,
+            NewResolutionOldInference.ResolutionKind.Variable
+        )
+    }
+
+    private fun CallResolver.resolveComposer(context: BasicCallResolutionContext):
+            ResolvedCall<CallableDescriptor>? {
+
+        // The composer is currently resolved as whatever is currently in scope with the name "composer".
+        val resolvedComposer = resolveVar(KtxNameConventions.COMPOSER, context)
+
+        if (!resolvedComposer.isSuccess) {
+            return null
+        }
+
+        return resolvedComposer.resultingCall
+    }
+
+    private fun KotlinType.isAbstract(): Boolean {
+        val modality = (constructor.declarationDescriptor as? ClassDescriptor)?.modality
+        return modality == Modality.ABSTRACT || modality == Modality.SEALED
+    }
+}

plugins/compose-compiler-hosted/src/main/java/androidx/compose/plugins/kotlin/ComposeDiagnosticSuppressor.kt

@@ -0,0 +1,80 @@
+/*
+ * Copyright 2019 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.compose.plugins.kotlin
+
+import com.intellij.openapi.extensions.Extensions
+import com.intellij.openapi.project.Project
+import org.jetbrains.kotlin.diagnostics.Diagnostic
+import org.jetbrains.kotlin.diagnostics.Errors
+import org.jetbrains.kotlin.psi.KtAnnotatedExpression
+import org.jetbrains.kotlin.psi.KtCallExpression
+import org.jetbrains.kotlin.psi.KtExpression
+import org.jetbrains.kotlin.resolve.BindingContext
+import org.jetbrains.kotlin.resolve.BindingTraceContext
+import org.jetbrains.kotlin.resolve.TemporaryBindingTrace
+import org.jetbrains.kotlin.resolve.calls.callUtil.getCall
+import org.jetbrains.kotlin.resolve.calls.callUtil.getResolvedCall
+import org.jetbrains.kotlin.resolve.diagnostics.DiagnosticSuppressor
+
+open class ComposeDiagnosticSuppressor : DiagnosticSuppressor {
+
+    companion object {
+        fun registerExtension(
+            @Suppress("UNUSED_PARAMETER") project: Project,
+            extension: DiagnosticSuppressor
+        ) {
+            @Suppress("DEPRECATION")
+            Extensions.getRootArea().getExtensionPoint(DiagnosticSuppressor.EP_NAME)
+                .registerExtension(extension)
+        }
+    }
+
+    override fun isSuppressed(diagnostic: Diagnostic): Boolean {
+        return isSuppressed(diagnostic, null)
+    }
+
+    override fun isSuppressed(diagnostic: Diagnostic, bindingContext: BindingContext?): Boolean {
+        if (diagnostic.factory == Errors.NON_SOURCE_ANNOTATION_ON_INLINED_LAMBDA_EXPRESSION) {
+            for (entry in (
+                    diagnostic.psiElement.parent as KtAnnotatedExpression
+                    ).annotationEntries) {
+                if (bindingContext != null) {
+                    val annotation = bindingContext.get(BindingContext.ANNOTATION, entry)
+                    if (annotation != null && annotation.isComposableAnnotation) return true
+                }
+                // Best effort, maybe jetbrains can get rid of nullability.
+                else if (entry.shortName?.identifier == "Composable") return true
+            }
+        }
+        if (diagnostic.factory == Errors.NAMED_ARGUMENTS_NOT_ALLOWED) {
+            val functionCall = diagnostic.psiElement.parent.parent.parent.parent as KtExpression
+            if (bindingContext != null) {
+                val call = (diagnostic.psiElement.parent.parent.parent.parent as KtCallExpression)
+                    .getCall(bindingContext).getResolvedCall(bindingContext)
+                val temporaryTrace = TemporaryBindingTrace.create(
+                    BindingTraceContext.createTraceableBindingTrace(),
+                    "trace to resolve ktx call",
+                    functionCall
+                )
+                if (call != null) {
+                    return ComposableAnnotationChecker().shouldInvokeAsTag(temporaryTrace, call)
+                }
+            }
+        }
+        return false
+    }
+}

plugins/compose-compiler-hosted/src/main/java/androidx/compose/plugins/kotlin/ComposeFlags.kt

@@ -0,0 +1,22 @@
+/*
+ * Copyright 2019 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.compose.plugins.kotlin
+
+object ComposeFlags {
+    var FRAMED_COMPONENTS = false
+    var FRAMED_MODEL_CLASSES = true
+}

plugins/compose-compiler-hosted/src/main/java/androidx/compose/plugins/kotlin/ComposeFqNames.kt

@@ -0,0 +1,90 @@
+/*
+ * Copyright 2019 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.compose.plugins.kotlin
+
+import org.jetbrains.kotlin.descriptors.ModuleDescriptor
+import org.jetbrains.kotlin.descriptors.SourceElement
+import org.jetbrains.kotlin.descriptors.annotations.Annotated
+import org.jetbrains.kotlin.descriptors.annotations.AnnotationDescriptor
+import org.jetbrains.kotlin.descriptors.annotations.Annotations
+import org.jetbrains.kotlin.descriptors.findClassAcrossModuleDependencies
+import org.jetbrains.kotlin.name.ClassId
+import org.jetbrains.kotlin.name.FqName
+import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlin.resolve.constants.ConstantValue
+import org.jetbrains.kotlin.resolve.descriptorUtil.annotationClass
+import org.jetbrains.kotlin.types.KotlinType
+import org.jetbrains.kotlin.types.TypeUtils.NO_EXPECTED_TYPE
+import org.jetbrains.kotlin.types.TypeUtils.UNIT_EXPECTED_TYPE
+import org.jetbrains.kotlin.types.typeUtil.replaceAnnotations
+
+object ComposeFqNames {
+    val Composable = ComposeUtils.composeFqName("Composable")
+    val CurrentComposerIntrinsic = ComposeUtils.composeFqName("<get-currentComposer>")
+    val Pivotal = ComposeUtils.composeFqName("Pivotal")
+    val StableMarker = ComposeUtils.composeFqName("StableMarker")
+    val HiddenAttribute = ComposeUtils.composeFqName("HiddenAttribute")
+    val Composer = ComposeUtils.composeFqName("Composer")
+    val Untracked = ComposeUtils.composeFqName("Untracked")
+    val ViewComposer = ComposeUtils.composeFqName("ViewComposer")
+    val Package = FqName.fromSegments(listOf("androidx", "compose"))
+    val Function0 = FqName.fromSegments(listOf("kotlin", "jvm", "functions", "Function0"))
+    val Function1 = FqName.fromSegments(listOf("kotlin", "jvm", "functions", "Function1"))
+    fun makeComposableAnnotation(module: ModuleDescriptor): AnnotationDescriptor =
+        object : AnnotationDescriptor {
+            override val type: KotlinType
+                get() = module.findClassAcrossModuleDependencies(
+                    ClassId.topLevel(Composable)
+                )!!.defaultType
+            override val allValueArguments: Map<Name, ConstantValue<*>> get() = emptyMap()
+            override val source: SourceElement get() = SourceElement.NO_SOURCE
+            override fun toString() = "[@Composable]"
+        }
+}
+
+fun KotlinType.makeComposable(module: ModuleDescriptor): KotlinType {
+    if (hasComposableAnnotation()) return this
+    val annotation = ComposeFqNames.makeComposableAnnotation(module)
+    return replaceAnnotations(Annotations.create(annotations + annotation))
+}
+
+fun KotlinType.hasComposableAnnotation(): Boolean =
+    !isSpecialType && annotations.findAnnotation(ComposeFqNames.Composable) != null
+fun KotlinType.isMarkedStable(): Boolean =
+    !isSpecialType && (
+                    annotations.hasStableMarker() ||
+                    (constructor.declarationDescriptor?.annotations?.hasStableMarker() ?: false))
+fun Annotated.hasComposableAnnotation(): Boolean =
+    annotations.findAnnotation(ComposeFqNames.Composable) != null
+fun Annotated.hasUntrackedAnnotation(): Boolean =
+    annotations.findAnnotation(ComposeFqNames.Untracked) != null
+fun Annotated.hasPivotalAnnotation(): Boolean =
+    annotations.findAnnotation(ComposeFqNames.Pivotal) != null
+fun Annotated.hasHiddenAttributeAnnotation(): Boolean =
+    annotations.findAnnotation(ComposeFqNames.HiddenAttribute) != null
+
+internal val KotlinType.isSpecialType: Boolean get() =
+    this === NO_EXPECTED_TYPE || this === UNIT_EXPECTED_TYPE
+
+val AnnotationDescriptor.isComposableAnnotation: Boolean get() = fqName == ComposeFqNames.Composable
+
+fun Annotations.hasStableMarker(): Boolean = any(AnnotationDescriptor::isStableMarker)
+
+fun AnnotationDescriptor.isStableMarker(): Boolean {
+    val classDescriptor = annotationClass ?: return false
+    return classDescriptor.annotations.hasAnnotation(ComposeFqNames.StableMarker)
+}

plugins/compose-compiler-hosted/src/main/java/androidx/compose/plugins/kotlin/ComposeIrGenerationExtension.kt

@@ -0,0 +1,92 @@
+/*
+ * Copyright 2019 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.compose.plugins.kotlin
+
+import androidx.compose.plugins.kotlin.compiler.lower.ComposableCallTransformer
+import androidx.compose.plugins.kotlin.compiler.lower.ComposeObservePatcher
+import androidx.compose.plugins.kotlin.compiler.lower.ComposeResolutionMetadataTransformer
+import androidx.compose.plugins.kotlin.compiler.lower.ComposerIntrinsicTransformer
+import androidx.compose.plugins.kotlin.compiler.lower.ComposerLambdaMemoization
+import androidx.compose.plugins.kotlin.compiler.lower.ComposerParamTransformer
+import androidx.compose.plugins.kotlin.frames.FrameIrTransformer
+import org.jetbrains.kotlin.backend.common.BackendContext
+import org.jetbrains.kotlin.backend.common.extensions.IrGenerationExtension
+import org.jetbrains.kotlin.backend.common.extensions.IrPluginContext
+import org.jetbrains.kotlin.backend.jvm.JvmBackendContext
+import org.jetbrains.kotlin.ir.declarations.IrFile
+import org.jetbrains.kotlin.ir.declarations.IrModuleFragment
+import org.jetbrains.kotlin.ir.util.DeepCopySymbolRemapper
+import org.jetbrains.kotlin.resolve.BindingContext
+import org.jetbrains.kotlin.resolve.DelegatingBindingTrace
+
+class ComposeIrGenerationExtension : IrGenerationExtension {
+    override fun generate(moduleFragment: IrModuleFragment, pluginContext: IrPluginContext) {
+        TODO("Unimplemented in new API")
+    }
+
+    fun generate(
+        file: IrFile,
+        backendContext: BackendContext,
+        bindingContext: BindingContext
+    ) {
+        // TODO: refactor transformers to work with just BackendContext
+        val jvmContext = backendContext as JvmBackendContext
+        val module = jvmContext.ir.irModule
+        val bindingTrace = DelegatingBindingTrace(
+            bindingContext,
+            "trace in ComposeIrGenerationExtension"
+        )
+
+        // We transform the entire module all at once since we end up remapping symbols, we need to
+        // ensure that everything in the module points to the right symbol. There is no extension
+        // point that allows you to transform at the module level but we should communicate this
+        // need with JetBrains as it seems like the only reasonable way to update top-level symbols.
+        // If a module-based extension point gets added, we should refactor this to use it.
+        if (file == module.files.first()) {
+
+            // create a symbol remapper to be used across all transforms
+            val symbolRemapper = DeepCopySymbolRemapper()
+
+            // add metadata from the frontend onto IR Nodes so that the metadata will travel
+            // with the ir nodes as they transform and get copied
+            ComposeResolutionMetadataTransformer(jvmContext).lower(module)
+
+            // transform @Model classes
+            FrameIrTransformer(jvmContext).lower(module)
+
+            // Memoize normal lambdas and wrap composable lambdas
+            ComposerLambdaMemoization(jvmContext, symbolRemapper, bindingTrace).lower(module)
+
+            // transform all composable functions to have an extra synthetic composer
+            // parameter. this will also transform all types and calls to include the extra
+            // parameter.
+            ComposerParamTransformer(jvmContext, symbolRemapper, bindingTrace).lower(module)
+
+            // transform calls to the currentComposer to just use the local parameter from the
+            // previous transform
+            ComposerIntrinsicTransformer(jvmContext).lower(module)
+
+            // transform composable calls and emits into their corresponding calls appealing
+            // to the composer
+            ComposableCallTransformer(jvmContext, symbolRemapper, bindingTrace).lower(module)
+
+            // transform composable functions to have restart groups so that they can be
+            // recomposed
+            ComposeObservePatcher(jvmContext, symbolRemapper, bindingTrace).lower(module)
+        }
+    }
+}

plugins/compose-compiler-hosted/src/main/java/androidx/compose/plugins/kotlin/ComposePlugin.kt

@@ -0,0 +1,114 @@
+/*
+ * Copyright 2019 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.compose.plugins.kotlin
+
+import com.intellij.mock.MockProject
+import com.intellij.openapi.project.Project
+import org.jetbrains.kotlin.compiler.plugin.CliOption
+import org.jetbrains.kotlin.compiler.plugin.CliOptionProcessingException
+import org.jetbrains.kotlin.compiler.plugin.CommandLineProcessor
+import org.jetbrains.kotlin.compiler.plugin.ComponentRegistrar
+import org.jetbrains.kotlin.compiler.plugin.AbstractCliOption
+import org.jetbrains.kotlin.config.CompilerConfiguration
+import org.jetbrains.kotlin.extensions.StorageComponentContainerContributor
+import androidx.compose.plugins.kotlin.frames.analysis.FrameModelChecker
+import androidx.compose.plugins.kotlin.frames.analysis.FramePackageAnalysisHandlerExtension
+import org.jetbrains.kotlin.backend.common.extensions.IrGenerationExtension
+import org.jetbrains.kotlin.extensions.internal.CandidateInterceptor
+import org.jetbrains.kotlin.extensions.internal.TypeResolutionInterceptor
+import org.jetbrains.kotlin.resolve.jvm.extensions.AnalysisHandlerExtension
+
+class ComposeCommandLineProcessor : CommandLineProcessor {
+
+    companion object {
+        val PLUGIN_ID = "androidx.compose.plugins.kotlin"
+    }
+
+    override val pluginId =
+        PLUGIN_ID
+    override val pluginOptions = emptyList<CliOption>()
+
+    @Suppress("OverridingDeprecatedMember")
+    override fun processOption(
+        option: CliOption,
+        value: String,
+        configuration: CompilerConfiguration
+    ) =
+        throw CliOptionProcessingException("Unknown option: ${option.optionName}")
+
+    override fun processOption(
+        option: AbstractCliOption,
+        value: String,
+        configuration: CompilerConfiguration
+    ) = throw CliOptionProcessingException("Unknown option: ${option.optionName}")
+}
+
+class ComposeComponentRegistrar : ComponentRegistrar {
+    override fun registerProjectComponents(
+        project: MockProject,
+        configuration: CompilerConfiguration
+    ) {
+        registerProjectExtensions(
+            project as Project,
+            configuration
+        )
+    }
+
+    companion object {
+
+        @Suppress("UNUSED_PARAMETER")
+        fun registerProjectExtensions(
+            project: Project,
+            configuration: CompilerConfiguration
+        ) {
+            StorageComponentContainerContributor.registerExtension(
+                project,
+                ComposableAnnotationChecker()
+            )
+            StorageComponentContainerContributor.registerExtension(
+                project,
+                UnionAnnotationCheckerProvider()
+            )
+            StorageComponentContainerContributor.registerExtension(
+                project,
+                TryCatchComposableChecker()
+            )
+            ComposeDiagnosticSuppressor.registerExtension(
+                project,
+                ComposeDiagnosticSuppressor()
+            )
+            TypeResolutionInterceptor.registerExtension(
+                project,
+                ComposeTypeResolutionInterceptorExtension()
+            )
+            IrGenerationExtension.registerExtension(project,
+                ComposeIrGenerationExtension()
+            )
+            CandidateInterceptor.registerExtension(
+                project,
+                ComposeCallResolutionInterceptorExtension()
+            )
+            StorageComponentContainerContributor.registerExtension(project,
+                FrameModelChecker()
+            )
+            AnalysisHandlerExtension.registerExtension(
+                project,
+                FramePackageAnalysisHandlerExtension()
+            )
+        }
+    }
+}

plugins/compose-compiler-hosted/src/main/java/androidx/compose/plugins/kotlin/ComposeTypeResolutionInterceptorExtension.kt

@@ -0,0 +1,72 @@
+/*
+ * Copyright 2019 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.compose.plugins.kotlin
+
+import androidx.compose.plugins.kotlin.analysis.ComposeWritableSlices.INFERRED_COMPOSABLE_DESCRIPTOR
+import org.jetbrains.kotlin.descriptors.impl.AnonymousFunctionDescriptor
+import org.jetbrains.kotlin.extensions.StorageComponentContainerContributor
+import org.jetbrains.kotlin.extensions.internal.TypeResolutionInterceptorExtension
+import org.jetbrains.kotlin.psi.KtElement
+import org.jetbrains.kotlin.psi.KtLambdaExpression
+import org.jetbrains.kotlin.resolve.descriptorUtil.module
+import org.jetbrains.kotlin.types.KotlinType
+import org.jetbrains.kotlin.types.TypeUtils
+import org.jetbrains.kotlin.types.expressions.ExpressionTypingContext
+
+/**
+ * If a lambda is marked as `@Composable`, then the inferred type should become `@Composable`
+ */
+@Suppress("INVISIBLE_REFERENCE", "EXPERIMENTAL_IS_NOT_ENABLED")
+@UseExperimental(org.jetbrains.kotlin.extensions.internal.InternalNonStableExtensionPoints::class)
+open class ComposeTypeResolutionInterceptorExtension : TypeResolutionInterceptorExtension {
+
+    override fun interceptFunctionLiteralDescriptor(
+        expression: KtLambdaExpression,
+        context: ExpressionTypingContext,
+        descriptor: AnonymousFunctionDescriptor
+    ): AnonymousFunctionDescriptor {
+        if (context.expectedType.hasComposableAnnotation()) {
+            // If the expected type has an @Composable annotation then the literal function
+            // expression should infer a an @Composable annotation
+            context.trace.record(INFERRED_COMPOSABLE_DESCRIPTOR, descriptor, true)
+        }
+        return descriptor
+    }
+
+    override fun interceptType(
+        element: KtElement,
+        context: ExpressionTypingContext,
+        resultType: KotlinType
+    ): KotlinType {
+        if (resultType === TypeUtils.NO_EXPECTED_TYPE) return resultType
+        if (element !is KtLambdaExpression) return resultType
+        val module = context.scope.ownerDescriptor.module
+        val checker =
+            StorageComponentContainerContributor.getInstances(element.project).single {
+                it is ComposableAnnotationChecker
+            } as ComposableAnnotationChecker
+        if ((context.expectedType.hasComposableAnnotation() || checker.analyze(
+                context.trace,
+                element,
+                resultType
+            ) != ComposableAnnotationChecker.Composability.NOT_COMPOSABLE)
+        ) {
+            return resultType.makeComposable(module)
+        }
+        return resultType
+    }
+}

plugins/compose-compiler-hosted/src/main/java/androidx/compose/plugins/kotlin/ComposeUtils.kt

@@ -0,0 +1,73 @@
+/*
+ * Copyright 2019 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.compose.plugins.kotlin
+
+import org.jetbrains.kotlin.descriptors.ClassDescriptor
+import org.jetbrains.kotlin.descriptors.DeclarationDescriptor
+import org.jetbrains.kotlin.descriptors.findClassAcrossModuleDependencies
+import org.jetbrains.kotlin.name.ClassId
+import org.jetbrains.kotlin.name.FqName
+import org.jetbrains.kotlin.psi.KtCallExpression
+import org.jetbrains.kotlin.psi.KtFunction
+import org.jetbrains.kotlin.psi.KtFunctionLiteral
+import org.jetbrains.kotlin.psi.KtLambdaArgument
+import org.jetbrains.kotlin.resolve.BindingContext
+import org.jetbrains.kotlin.resolve.calls.callUtil.getResolvedCall
+import org.jetbrains.kotlin.resolve.descriptorUtil.isSubclassOf
+import org.jetbrains.kotlin.resolve.descriptorUtil.module
+
+object ComposeUtils {
+
+    fun generateComposePackageName() = "androidx.compose"
+
+    fun composeFqName(cname: String) = FqName("${generateComposePackageName()}.$cname")
+
+    fun setterMethodFromPropertyName(name: String): String {
+        return "set${name[0].toUpperCase()}${name.slice(1 until name.length)}"
+    }
+
+    fun propertyNameFromSetterMethod(name: String): String {
+        return if (name.startsWith("set")) "${
+            name[3].toLowerCase()
+        }${name.slice(4 until name.length)}" else name
+    }
+
+    fun isSetterMethodName(name: String): Boolean {
+        // use !lower to capture non-alpha chars
+        return name.startsWith("set") && name.length > 3 && !name[3].isLowerCase()
+    }
+
+    fun isComposeComponent(descriptor: DeclarationDescriptor): Boolean {
+        if (descriptor !is ClassDescriptor) return false
+        val baseComponentDescriptor =
+            descriptor.module.findClassAcrossModuleDependencies(
+                ClassId.topLevel(
+                    FqName(ComposeUtils.generateComposePackageName() + ".Component")
+                )
+            ) ?: return false
+        return descriptor.isSubclassOf(baseComponentDescriptor)
+    }
+}
+
+fun KtFunction.isEmitInline(bindingContext: BindingContext): Boolean {
+    if (this !is KtFunctionLiteral) return false
+    if (parent?.parent !is KtLambdaArgument) return false
+    val call = parent?.parent?.parent as? KtCallExpression
+    val resolvedCall = call?.getResolvedCall(bindingContext)
+    return resolvedCall != null &&
+            resolvedCall.candidateDescriptor is ComposableEmitDescriptor
+}

plugins/compose-compiler-hosted/src/main/java/androidx/compose/plugins/kotlin/ComposerMetadata.kt

@@ -0,0 +1,212 @@
+/*
+ * Copyright 2019 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.compose.plugins.kotlin
+
+import androidx.compose.plugins.kotlin.analysis.ComposeWritableSlices
+import org.jetbrains.kotlin.builtins.getReturnTypeFromFunctionType
+import org.jetbrains.kotlin.builtins.getValueParameterTypesFromFunctionType
+import org.jetbrains.kotlin.builtins.isFunctionTypeOrSubtype
+import org.jetbrains.kotlin.descriptors.CallableDescriptor
+import org.jetbrains.kotlin.descriptors.SimpleFunctionDescriptor
+import org.jetbrains.kotlin.descriptors.ValueParameterDescriptor
+import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlin.psi.KtPsiFactory
+import org.jetbrains.kotlin.resolve.calls.CallResolver
+import org.jetbrains.kotlin.resolve.calls.context.BasicCallResolutionContext
+import org.jetbrains.kotlin.resolve.calls.context.CheckArgumentTypesMode
+import org.jetbrains.kotlin.resolve.calls.model.DataFlowInfoForArgumentsImpl
+import org.jetbrains.kotlin.resolve.calls.model.ResolvedCall
+import org.jetbrains.kotlin.resolve.scopes.receivers.TransientReceiver
+import org.jetbrains.kotlin.types.KotlinType
+import org.jetbrains.kotlin.types.isError
+import org.jetbrains.kotlin.types.typeUtil.isNothingOrNullableNothing
+import org.jetbrains.kotlin.types.typeUtil.isSubtypeOf
+
+class ComposerMetadata(
+    val type: KotlinType,
+    // Set of valid upper bound types that were defined on the composer that can't have children
+    // For android, this should be [View]
+    private val emitSimpleUpperBoundTypes: Set<KotlinType>,
+    // Set of valid upper bound types that were defined on the composer that can have children.
+    // For android, this would be [ViewGroup]
+    private val emitCompoundUpperBoundTypes: Set<KotlinType>,
+    // The specification for `emit` on a composer allows for the `ctor` parameter to be a function type
+    // with any number of parameters. We allow for these parameters to be used as parameters in the
+    // Constructors that are emitted with a KTX tag. These parameters can be overridden with attributes
+    // in the KTX tag, but if there are required parameters with a type that matches one declared in the
+    // ctor parameter, we will resolve it automatically with the value passed in the `ctor` lambda.
+    //
+    // In order to do this resolution, we store a list of pairs of "upper bounds" to parameter types. For example,
+    // the following emit call:
+    //
+    //      fun <T : View> emit(key: Any, ctor: (context: Context) -> T, update: U<T>.() -> Unit)
+    //
+    // would produce a Pair of [View] to [Context]
+    private val emittableTypeToImplicitCtorTypes: List<Pair<List<KotlinType>, Set<KotlinType>>>
+) {
+
+    companion object {
+        private fun resolveComposerMethodCandidates(
+            name: Name,
+            context: BasicCallResolutionContext,
+            composerType: KotlinType,
+            callResolver: CallResolver,
+            psiFactory: KtPsiFactory
+        ): Collection<ResolvedCall<*>> {
+            val calleeExpression = psiFactory.createSimpleName(name.asString())
+
+            val methodCall = makeCall(
+                callElement = context.call.callElement,
+                calleeExpression = calleeExpression,
+                receiver = TransientReceiver(
+                    composerType
+                )
+            )
+
+            val contextForVariable =
+                BasicCallResolutionContext.create(
+                    context,
+                    methodCall,
+                    CheckArgumentTypesMode.CHECK_VALUE_ARGUMENTS,
+                    DataFlowInfoForArgumentsImpl(
+                        context.dataFlowInfo,
+                        methodCall
+                    )
+                )
+
+            val results = callResolver.resolveCallWithGivenName(
+                // it's important that we use "collectAllCandidates" so that extension functions get included
+                contextForVariable.replaceCollectAllCandidates(true),
+                methodCall,
+                calleeExpression,
+                name
+            )
+
+            return results.allCandidates ?: emptyList()
+        }
+
+        fun build(
+            composerType: KotlinType,
+            callResolver: CallResolver,
+            psiFactory: KtPsiFactory,
+            resolutionContext: BasicCallResolutionContext
+        ): ComposerMetadata {
+            val emitSimpleUpperBoundTypes = mutableSetOf<KotlinType>()
+            val emitCompoundUpperBoundTypes = mutableSetOf<KotlinType>()
+            val emittableTypeToImplicitCtorTypes =
+                mutableListOf<Pair<List<KotlinType>, Set<KotlinType>>>()
+
+            val emitCandidates = resolveComposerMethodCandidates(
+                KtxNameConventions.EMIT,
+                resolutionContext,
+                composerType,
+                callResolver,
+                psiFactory
+            )
+
+            for (candidate in emitCandidates.map { it.candidateDescriptor }) {
+                if (candidate.name != KtxNameConventions.EMIT) continue
+                if (candidate !is SimpleFunctionDescriptor) continue
+                val params = candidate.valueParameters
+                // NOTE(lmr): we could report diagnostics on some of these? it seems strange to emit diagnostics about a function
+                // that is not necessarily being used though. I think it's probably better to just ignore them here.
+
+                // the signature of emit that we are looking for has 3 or 4 parameters
+                if (params.size < 3 || params.size > 4) continue
+                val ctorParam = params.find {
+                    it.name == KtxNameConventions.EMIT_CTOR_PARAMETER
+                } ?: continue
+                if (!ctorParam.type.isFunctionTypeOrSubtype) continue
+
+                // the return type from the ctor param is the "upper bound" of the node type. It will often be a generic type with constraints.
+                val upperBounds = ctorParam.type.getReturnTypeFromFunctionType().upperBounds()
+
+                // the ctor param can have parameters itself, which we interpret as implicit parameter types that the composer knows how to
+                // automatically provide to the component. In the case of Android Views, this is how we automatically provide Context.
+                val implicitParamTypes =
+                    ctorParam.type.getValueParameterTypesFromFunctionType().map {
+                        it.type
+                    }
+
+                for (implicitType in implicitParamTypes) {
+                    emittableTypeToImplicitCtorTypes.add(upperBounds to implicitParamTypes.toSet())
+                }
+
+                emitSimpleUpperBoundTypes.addAll(upperBounds)
+
+                if (params.any { it.name == KtxNameConventions.EMIT_CHILDREN_PARAMETER }) {
+                    emitCompoundUpperBoundTypes.addAll(upperBounds)
+                }
+            }
+
+            return ComposerMetadata(
+                composerType,
+                emitSimpleUpperBoundTypes,
+                emitCompoundUpperBoundTypes,
+                emittableTypeToImplicitCtorTypes
+            )
+        }
+
+        fun getOrBuild(
+            composerType: KotlinType,
+            callResolver: CallResolver,
+            psiFactory: KtPsiFactory,
+            resolutionContext: BasicCallResolutionContext
+        ): ComposerMetadata {
+            val meta = resolutionContext.trace.bindingContext[
+                    ComposeWritableSlices.COMPOSER_METADATA,
+                    composerType
+            ]
+            return if (meta == null) {
+                val built = build(composerType, callResolver, psiFactory, resolutionContext)
+                resolutionContext.trace.record(
+                    ComposeWritableSlices.COMPOSER_METADATA,
+                    composerType,
+                    built
+                )
+                built
+            } else {
+                meta
+            }
+        }
+    }
+
+    fun isEmittable(type: KotlinType) =
+        !type.isError && !type.isNothingOrNullableNothing() && emitSimpleUpperBoundTypes.any {
+            type.isSubtypeOf(it)
+        }
+
+    fun isCompoundEmittable(type: KotlinType) = !type.isError &&
+            !type.isNothingOrNullableNothing() &&
+            emitCompoundUpperBoundTypes.any {
+                type.isSubtypeOf(it)
+            }
+
+    fun isImplicitConstructorParam(
+        param: ValueParameterDescriptor,
+        fn: CallableDescriptor
+    ): Boolean {
+        val returnType = fn.returnType ?: return false
+        val paramType = param.type
+        for ((upperBounds, implicitTypes) in emittableTypeToImplicitCtorTypes) {
+            if (!implicitTypes.any { it.isSubtypeOf(paramType) }) continue
+            if (!returnType.satisfiesConstraintsOf(upperBounds)) continue
+            return true
+        }
+        return false
+    }
+}

plugins/compose-compiler-hosted/src/main/java/androidx/compose/plugins/kotlin/KtxNameConventions.kt

@@ -0,0 +1,26 @@
+package androidx.compose.plugins.kotlin
+
+import org.jetbrains.kotlin.name.Name
+
+object KtxNameConventions {
+    val COMPOSER = Name.identifier("composer")
+    val COMPOSER_PARAMETER = Name.identifier("\$composer")
+    val EMIT = Name.identifier("emit")
+    val CALL = Name.identifier("call")
+    val START_EXPR = Name.identifier("startExpr")
+    val END_EXPR = Name.identifier("endExpr")
+    val JOINKEY = Name.identifier("joinKey")
+    val STARTRESTARTGROUP = Name.identifier("startRestartGroup")
+    val ENDRESTARTGROUP = Name.identifier("endRestartGroup")
+    val UPDATE_SCOPE = Name.identifier("updateScope")
+
+    val EMIT_KEY_PARAMETER = Name.identifier("key")
+    val EMIT_CTOR_PARAMETER = Name.identifier("ctor")
+    val EMIT_UPDATER_PARAMETER = Name.identifier("update")
+    val EMIT_CHILDREN_PARAMETER = Name.identifier("children")
+
+    val CALL_KEY_PARAMETER = Name.identifier("key")
+    val CALL_CTOR_PARAMETER = Name.identifier("ctor")
+    val CALL_INVALID_PARAMETER = Name.identifier("invalid")
+    val CALL_BLOCK_PARAMETER = Name.identifier("block")
+}

plugins/compose-compiler-hosted/src/main/java/androidx/compose/plugins/kotlin/TryCatchComposableChecker.kt

@@ -0,0 +1,70 @@
+/*
+ * Copyright 2019 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.compose.plugins.kotlin
+
+import androidx.compose.plugins.kotlin.analysis.ComposeDefaultErrorMessages
+import androidx.compose.plugins.kotlin.analysis.ComposeErrors
+import com.intellij.psi.PsiElement
+import org.jetbrains.kotlin.container.StorageComponentContainer
+import org.jetbrains.kotlin.container.useInstance
+import org.jetbrains.kotlin.descriptors.ModuleDescriptor
+import org.jetbrains.kotlin.diagnostics.reportFromPlugin
+import org.jetbrains.kotlin.extensions.StorageComponentContainerContributor
+import org.jetbrains.kotlin.platform.TargetPlatform
+import org.jetbrains.kotlin.platform.jvm.isJvm
+import org.jetbrains.kotlin.psi.KtTryExpression
+import org.jetbrains.kotlin.resolve.calls.checkers.CallChecker
+import org.jetbrains.kotlin.resolve.calls.checkers.CallCheckerContext
+import org.jetbrains.kotlin.resolve.calls.model.ResolvedCall
+
+open class TryCatchComposableChecker : CallChecker, StorageComponentContainerContributor {
+
+    override fun registerModuleComponents(
+        container: StorageComponentContainer,
+        platform: TargetPlatform,
+        moduleDescriptor: ModuleDescriptor
+    ) {
+        if (!platform.isJvm()) return
+        container.useInstance(this)
+    }
+
+    override fun check(
+        resolvedCall: ResolvedCall<*>,
+        reportOn: PsiElement,
+        context: CallCheckerContext
+    ) {
+        val trace = context.trace
+        val call = resolvedCall.call.callElement
+        val shouldBeTag =
+            ComposableAnnotationChecker.get(call.project).shouldInvokeAsTag(trace, resolvedCall)
+        if (shouldBeTag) {
+            var walker: PsiElement? = call
+            while (walker != null) {
+                val parent = walker.parent
+                if (parent is KtTryExpression) {
+                    if (walker == parent.tryBlock)
+                        trace.report(
+                            ComposeErrors.ILLEGAL_TRY_CATCH_AROUND_COMPOSABLE.on(
+                                parent.tryKeyword!!
+                            )
+                        )
+                }
+                walker = try { walker.parent } catch (e: Throwable) { null }
+            }
+        }
+    }
+}

plugins/compose-compiler-hosted/src/main/java/androidx/compose/plugins/kotlin/UnionAnnotationChecker.kt

@@ -0,0 +1,95 @@
+/*
+ * Copyright 2019 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.compose.plugins.kotlin
+
+import org.jetbrains.kotlin.types.TypeUtils
+import org.jetbrains.kotlin.container.StorageComponentContainer
+import org.jetbrains.kotlin.container.useInstance
+import org.jetbrains.kotlin.descriptors.ModuleDescriptor
+import org.jetbrains.kotlin.extensions.StorageComponentContainerContributor
+import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlin.psi.KtExpression
+import androidx.compose.plugins.kotlin.analysis.ComposeErrors
+import org.jetbrains.kotlin.platform.TargetPlatform
+import org.jetbrains.kotlin.platform.jvm.isJvm
+import org.jetbrains.kotlin.resolve.calls.checkers.AdditionalTypeChecker
+import org.jetbrains.kotlin.resolve.calls.context.ResolutionContext
+import org.jetbrains.kotlin.resolve.constants.ArrayValue
+import org.jetbrains.kotlin.types.KotlinType
+import org.jetbrains.kotlin.types.checker.KotlinTypeChecker
+
+open class UnionAnnotationCheckerProvider() : StorageComponentContainerContributor {
+    override fun registerModuleComponents(
+        container: StorageComponentContainer,
+        platform: TargetPlatform,
+        moduleDescriptor: ModuleDescriptor
+    ) {
+        if (!platform.isJvm()) return
+        container.useInstance(
+            UnionAnnotationChecker(
+                moduleDescriptor
+            )
+        )
+    }
+}
+
+open class UnionAnnotationChecker(val moduleDescriptor: ModuleDescriptor) : AdditionalTypeChecker {
+    companion object {
+        val UNIONTYPE_ANNOTATION_NAME =
+            ComposeUtils.composeFqName("UnionType")
+    }
+
+    override fun checkType(
+        expression: KtExpression,
+        expressionType: KotlinType,
+        expressionTypeWithSmartCast: KotlinType,
+        c: ResolutionContext<*>
+    ) {
+        val expectedType = c.expectedType
+        if (TypeUtils.noExpectedType(expectedType)) return
+
+        if (!expectedType.annotations.hasAnnotation(UNIONTYPE_ANNOTATION_NAME) &&
+            !expressionTypeWithSmartCast.annotations.hasAnnotation(UNIONTYPE_ANNOTATION_NAME)) {
+            return
+        }
+
+        val expressionTypes = getUnionTypes(expressionTypeWithSmartCast)
+        val permittedTypes = getUnionTypes(expectedType)
+
+        outer@ for (potentialExpressionType in expressionTypes) {
+            for (permittedType in permittedTypes) {
+                if (KotlinTypeChecker.DEFAULT.isSubtypeOf(potentialExpressionType, permittedType))
+                    continue@outer
+            }
+            c.trace.report(
+                ComposeErrors.ILLEGAL_ASSIGN_TO_UNIONTYPE.on(
+                    expression,
+                    listOf(potentialExpressionType),
+                    permittedTypes
+                )
+            )
+            return
+        }
+    }
+
+    private fun getUnionTypes(type: KotlinType): List<KotlinType> {
+        val annotation =
+            type.annotations.findAnnotation(UNIONTYPE_ANNOTATION_NAME) ?: return listOf(type)
+        val types = annotation.allValueArguments.get(Name.identifier("types")) as ArrayValue
+        return types.value.map { it.getType(moduleDescriptor).arguments.single().type }
+    }
+}

plugins/compose-compiler-hosted/src/main/java/androidx/compose/plugins/kotlin/ValidatedAssignment.kt

@@ -0,0 +1,39 @@
+/*
+ * Copyright 2019 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.compose.plugins.kotlin
+
+import org.jetbrains.kotlin.descriptors.DeclarationDescriptor
+import org.jetbrains.kotlin.descriptors.FunctionDescriptor
+import org.jetbrains.kotlin.resolve.calls.model.ResolvedCall
+import org.jetbrains.kotlin.types.KotlinType
+
+class ValidatedAssignment(
+    val validationType: ValidationType,
+    val validationCall: ResolvedCall<*>?,
+    val uncheckedValidationCall: ResolvedCall<*>?,
+    val assignment: ResolvedCall<*>?,
+    val assignmentLambda: FunctionDescriptor?, // needed?
+    val type: KotlinType,
+    val name: String,
+    val descriptor: DeclarationDescriptor
+)
+
+enum class ValidationType {
+    CHANGED,
+    SET,
+    UPDATE
+}

plugins/compose-compiler-hosted/src/main/java/androidx/compose/plugins/kotlin/WeakBindingTrace.kt

@@ -0,0 +1,59 @@
+/*
+ * Copyright 2019 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.compose.plugins.kotlin
+
+import com.intellij.util.keyFMap.KeyFMap
+import org.jetbrains.kotlin.backend.common.BackendContext
+import org.jetbrains.kotlin.codegen.state.GenerationState
+import org.jetbrains.kotlin.ir.declarations.IrAttributeContainer
+import org.jetbrains.kotlin.psi2ir.generators.GeneratorContext
+import org.jetbrains.kotlin.util.slicedMap.ReadOnlySlice
+import org.jetbrains.kotlin.util.slicedMap.WritableSlice
+import java.util.WeakHashMap
+
+/**
+ * This class is meant to have the shape of a BindingTrace object that could exist and flow
+ * through the Psi2Ir -> Ir phase, but doesn't currently exist. Ideally, this gets replaced in
+ * the future by a trace that handles this use case in upstream. For now, we are okay with this
+ * because the combination of IrAttributeContainer and WeakHashMap makes this relatively safe.
+ */
+class WeakBindingTrace {
+    private val map = WeakHashMap<Any, KeyFMap>()
+
+    fun <K : IrAttributeContainer, V> record(slice: WritableSlice<K, V>, key: K, value: V) {
+        var holder = map[key.attributeOwnerId] ?: KeyFMap.EMPTY_MAP
+        val prev = holder.get(slice.key)
+        if (prev != null) {
+            holder = holder.minus(slice.key)
+        }
+        holder = holder.plus(slice.key, value)
+        map[key.attributeOwnerId] = holder
+    }
+
+    operator fun <K : IrAttributeContainer, V> get(slice: ReadOnlySlice<K, V>, key: K): V? {
+        return map[key.attributeOwnerId]?.get(slice.key)
+    }
+}
+
+private val ComposeTemporaryGlobalBindingTrace = WeakBindingTrace()
+
+@Suppress("unused")
+val GeneratorContext.irTrace: WeakBindingTrace get() = ComposeTemporaryGlobalBindingTrace
+@Suppress("unused")
+val GenerationState.irTrace: WeakBindingTrace get() = ComposeTemporaryGlobalBindingTrace
+@Suppress("unused")
+val BackendContext.irTrace: WeakBindingTrace get() = ComposeTemporaryGlobalBindingTrace

plugins/compose-compiler-hosted/src/main/java/androidx/compose/plugins/kotlin/analysis/ComposeDefaultErrorMessages.kt

@@ -0,0 +1,60 @@
+package androidx.compose.plugins.kotlin.analysis
+
+import org.jetbrains.kotlin.diagnostics.rendering.DefaultErrorMessages
+import org.jetbrains.kotlin.diagnostics.rendering.DiagnosticFactoryToRendererMap
+import org.jetbrains.kotlin.diagnostics.rendering.DiagnosticParameterRenderer
+import org.jetbrains.kotlin.diagnostics.rendering.Renderers
+import org.jetbrains.kotlin.diagnostics.rendering.Renderers.RENDER_COLLECTION_OF_TYPES
+import org.jetbrains.kotlin.diagnostics.rendering.RenderingContext
+
+object ComposeDefaultErrorMessages : DefaultErrorMessages.Extension {
+    private val MAP = DiagnosticFactoryToRendererMap("Compose")
+    override fun getMap() = MAP
+
+    val OUR_STRING_RENDERER = object : DiagnosticParameterRenderer<String> {
+        override fun render(obj: String, renderingContext: RenderingContext): String {
+            return obj
+        }
+    }
+
+    init {
+        MAP.put(
+            ComposeErrors.NO_COMPOSER_FOUND,
+            "Couldn't find a valid composer."
+        )
+        MAP.put(
+            ComposeErrors.OPEN_MODEL,
+            "Model objects cannot be open or abstract"
+        )
+        MAP.put(
+            ComposeErrors.INVALID_COMPOSER_IMPLEMENTATION,
+            "Composer of type ''{0}'' was found to be an invalid Composer implementation. " +
+                    "Reason: {1}",
+            Renderers.RENDER_TYPE,
+            OUR_STRING_RENDERER
+        )
+        MAP.put(
+            ComposeErrors.SUSPEND_FUNCTION_USED_AS_SFC,
+            "Suspend functions are not allowed to be used as Components"
+        )
+        MAP.put(
+            ComposeErrors.INVALID_TYPE_SIGNATURE_SFC,
+            "Only Unit-returning functions are allowed to be used as Components"
+        )
+        MAP.put(
+            ComposeErrors.COMPOSABLE_INVOCATION_IN_NON_COMPOSABLE,
+            "Functions which invoke @Composable functions must be marked with the @Composable " +
+                    "annotation"
+        )
+        MAP.put(
+            ComposeErrors.ILLEGAL_ASSIGN_TO_UNIONTYPE,
+            "Value of type {0} can't be assigned to union type {1}.",
+            RENDER_COLLECTION_OF_TYPES,
+            RENDER_COLLECTION_OF_TYPES
+        )
+        MAP.put(
+            ComposeErrors.ILLEGAL_TRY_CATCH_AROUND_COMPOSABLE,
+            "Try catch is not supported around composable function invocations."
+        )
+    }
+}

plugins/compose-compiler-hosted/src/main/java/androidx/compose/plugins/kotlin/analysis/ComposeErrors.java

@@ -0,0 +1,58 @@
+/*
+ * Copyright 2018 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.compose.plugins.kotlin.analysis;
+
+import static org.jetbrains.kotlin.diagnostics.Severity.ERROR;
+
+import com.intellij.psi.PsiElement;
+
+import org.jetbrains.kotlin.diagnostics.DiagnosticFactory0;
+import org.jetbrains.kotlin.diagnostics.DiagnosticFactory2;
+import org.jetbrains.kotlin.diagnostics.Errors;
+import org.jetbrains.kotlin.psi.KtElement;
+import org.jetbrains.kotlin.psi.KtExpression;
+import org.jetbrains.kotlin.types.KotlinType;
+
+import java.util.Collection;
+
+/**
+ * Error messages
+ */
+public interface ComposeErrorsLegacy {
+    DiagnosticFactory0<PsiElement> OPEN_MODEL = DiagnosticFactory0.create(ERROR);
+    DiagnosticFactory0<KtElement>
+            SUSPEND_FUNCTION_USED_AS_SFC = DiagnosticFactory0.create(ERROR);
+    DiagnosticFactory0<PsiElement>
+            COMPOSABLE_INVOCATION_IN_NON_COMPOSABLE = DiagnosticFactory0.create(ERROR);
+    DiagnosticFactory0<KtElement>
+            INVALID_TYPE_SIGNATURE_SFC = DiagnosticFactory0.create(ERROR);
+    DiagnosticFactory0<KtElement>
+            NO_COMPOSER_FOUND = DiagnosticFactory0.create(ERROR);
+    DiagnosticFactory2<KtElement, KotlinType, String>
+            INVALID_COMPOSER_IMPLEMENTATION = DiagnosticFactory2.create(ERROR);
+    DiagnosticFactory2<KtExpression, Collection<KotlinType>, Collection<KotlinType>>
+            ILLEGAL_ASSIGN_TO_UNIONTYPE = DiagnosticFactory2.create(ERROR);
+    DiagnosticFactory0<PsiElement>
+            ILLEGAL_TRY_CATCH_AROUND_COMPOSABLE = DiagnosticFactory0.create(ERROR);
+
+    @SuppressWarnings("UnusedDeclaration")
+    Object INITIALIZER = new Object() {
+        {
+            Errors.Initializer.initializeFactoryNames(ComposeErrors.class);
+        }
+    };
+
+}

plugins/compose-compiler-hosted/src/main/java/androidx/compose/plugins/kotlin/analysis/ComposeErrors.kt

@@ -0,0 +1,56 @@
+/*
+ * Copyright 2010-2020 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package androidx.compose.plugins.kotlin.analysis
+
+import com.intellij.psi.PsiElement
+import org.jetbrains.kotlin.diagnostics.DiagnosticFactory0
+import org.jetbrains.kotlin.diagnostics.DiagnosticFactory2
+import org.jetbrains.kotlin.diagnostics.Errors
+import org.jetbrains.kotlin.diagnostics.Severity
+import org.jetbrains.kotlin.psi.KtElement
+import org.jetbrains.kotlin.psi.KtExpression
+import org.jetbrains.kotlin.types.KotlinType
+
+/**
+ * Error messages
+ */
+interface ComposeErrors {
+    companion object {
+        val OPEN_MODEL =
+            DiagnosticFactory0
+                .create<PsiElement>(Severity.ERROR)
+        val SUSPEND_FUNCTION_USED_AS_SFC =
+            DiagnosticFactory0
+                .create<KtElement>(Severity.ERROR)
+        val COMPOSABLE_INVOCATION_IN_NON_COMPOSABLE =
+            DiagnosticFactory0
+                .create<PsiElement>(Severity.ERROR)
+        val INVALID_TYPE_SIGNATURE_SFC =
+            DiagnosticFactory0
+                .create<KtElement>(Severity.ERROR)
+        val NO_COMPOSER_FOUND =
+            DiagnosticFactory0
+                .create<KtElement>(Severity.ERROR)
+        val INVALID_COMPOSER_IMPLEMENTATION =
+            DiagnosticFactory2
+                .create<KtElement, KotlinType, String>(
+                    Severity.ERROR
+                )
+        val ILLEGAL_ASSIGN_TO_UNIONTYPE =
+            DiagnosticFactory2
+                .create<KtExpression, Collection<KotlinType>, Collection<KotlinType>>(
+                    Severity.ERROR
+                )
+        val ILLEGAL_TRY_CATCH_AROUND_COMPOSABLE =
+            DiagnosticFactory0
+                .create<PsiElement>(Severity.ERROR)
+        val INITIALIZER: Any = object : Any() {
+            init {
+                Errors.Initializer.initializeFactoryNames(ComposeErrors::class.java)
+            }
+        }
+    }
+}

plugins/compose-compiler-hosted/src/main/java/androidx/compose/plugins/kotlin/analysis/ComposeWritableSlices.kt

@@ -0,0 +1,37 @@
+package androidx.compose.plugins.kotlin.analysis
+
+import androidx.compose.plugins.kotlin.ComposableAnnotationChecker
+import androidx.compose.plugins.kotlin.ComposableEmitMetadata
+import androidx.compose.plugins.kotlin.ComposerMetadata
+import org.jetbrains.kotlin.descriptors.FunctionDescriptor
+import org.jetbrains.kotlin.ir.declarations.IrAttributeContainer
+import org.jetbrains.kotlin.ir.expressions.IrFunctionAccessExpression
+import org.jetbrains.kotlin.psi.Call
+import org.jetbrains.kotlin.psi.KtElement
+import org.jetbrains.kotlin.util.slicedMap.BasicWritableSlice
+import org.jetbrains.kotlin.util.slicedMap.RewritePolicy
+import org.jetbrains.kotlin.util.slicedMap.WritableSlice
+import org.jetbrains.kotlin.types.KotlinType
+
+object ComposeWritableSlices {
+    val COMPOSABLE_ANALYSIS: WritableSlice<KtElement, ComposableAnnotationChecker.Composability> =
+        BasicWritableSlice(RewritePolicy.DO_NOTHING)
+    val FCS_RESOLVEDCALL_COMPOSABLE: WritableSlice<KtElement, Boolean> =
+        BasicWritableSlice(RewritePolicy.DO_NOTHING)
+    val INFERRED_COMPOSABLE_DESCRIPTOR: WritableSlice<FunctionDescriptor, Boolean> =
+        BasicWritableSlice(RewritePolicy.DO_NOTHING)
+    val STABLE_TYPE: WritableSlice<KotlinType, Boolean?> =
+        BasicWritableSlice(RewritePolicy.DO_NOTHING)
+    val COMPOSER_METADATA: WritableSlice<KotlinType, ComposerMetadata> =
+        BasicWritableSlice(RewritePolicy.DO_NOTHING)
+    val IGNORE_COMPOSABLE_INTERCEPTION: WritableSlice<Call, Boolean> =
+        BasicWritableSlice(RewritePolicy.DO_NOTHING)
+    val COMPOSABLE_EMIT_METADATA: WritableSlice<IrAttributeContainer, ComposableEmitMetadata> =
+        BasicWritableSlice(RewritePolicy.DO_NOTHING)
+    val IS_COMPOSABLE_CALL: WritableSlice<IrAttributeContainer, Boolean> =
+        BasicWritableSlice(RewritePolicy.DO_NOTHING)
+    val IS_INLINE_COMPOSABLE_CALL: WritableSlice<IrAttributeContainer, Boolean> =
+        BasicWritableSlice(RewritePolicy.DO_NOTHING)
+    val IS_SYNTHETIC_COMPOSABLE_CALL: WritableSlice<IrFunctionAccessExpression, Boolean> =
+        BasicWritableSlice(RewritePolicy.DO_NOTHING)
+}

plugins/compose-compiler-hosted/src/main/java/androidx/compose/plugins/kotlin/compiler/lower/AbstractComposeLowering.kt

@@ -0,0 +1,353 @@
+/*
+ * Copyright 2020 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.compose.plugins.kotlin.compiler.lower
+
+import androidx.compose.plugins.kotlin.ComposableAnnotationChecker
+import androidx.compose.plugins.kotlin.ComposeFqNames
+import androidx.compose.plugins.kotlin.KtxNameConventions
+import androidx.compose.plugins.kotlin.analysis.ComposeWritableSlices
+import androidx.compose.plugins.kotlin.irTrace
+import androidx.compose.plugins.kotlin.isMarkedStable
+import androidx.compose.plugins.kotlin.isSpecialType
+import org.jetbrains.kotlin.backend.common.descriptors.isFunctionOrKFunctionType
+import org.jetbrains.kotlin.backend.common.lower.DeclarationIrBuilder
+import org.jetbrains.kotlin.backend.jvm.JvmBackendContext
+import org.jetbrains.kotlin.builtins.KotlinBuiltIns
+import org.jetbrains.kotlin.builtins.extractParameterNameFromFunctionTypeArgument
+import org.jetbrains.kotlin.builtins.getReceiverTypeFromFunctionType
+import org.jetbrains.kotlin.builtins.getReturnTypeFromFunctionType
+import org.jetbrains.kotlin.builtins.getValueParameterTypesFromFunctionType
+import org.jetbrains.kotlin.descriptors.CallableDescriptor
+import org.jetbrains.kotlin.descriptors.CallableMemberDescriptor
+import org.jetbrains.kotlin.descriptors.ClassConstructorDescriptor
+import org.jetbrains.kotlin.descriptors.ClassDescriptor
+import org.jetbrains.kotlin.descriptors.ClassKind
+import org.jetbrains.kotlin.descriptors.DeclarationDescriptor
+import org.jetbrains.kotlin.descriptors.FunctionDescriptor
+import org.jetbrains.kotlin.descriptors.Modality
+import org.jetbrains.kotlin.descriptors.ParameterDescriptor
+import org.jetbrains.kotlin.descriptors.SimpleFunctionDescriptor
+import org.jetbrains.kotlin.descriptors.SourceElement
+import org.jetbrains.kotlin.descriptors.TypeParameterDescriptor
+import org.jetbrains.kotlin.descriptors.ValueParameterDescriptor
+import org.jetbrains.kotlin.descriptors.Visibilities
+import org.jetbrains.kotlin.descriptors.annotations.Annotations
+import org.jetbrains.kotlin.descriptors.findClassAcrossModuleDependencies
+import org.jetbrains.kotlin.descriptors.impl.AnonymousFunctionDescriptor
+import org.jetbrains.kotlin.descriptors.impl.ValueParameterDescriptorImpl
+import org.jetbrains.kotlin.incremental.components.NoLookupLocation
+import org.jetbrains.kotlin.ir.UNDEFINED_OFFSET
+import org.jetbrains.kotlin.ir.builders.*
+import org.jetbrains.kotlin.ir.declarations.IrAnnotationContainer
+import org.jetbrains.kotlin.ir.declarations.IrDeclarationOrigin
+import org.jetbrains.kotlin.ir.declarations.IrFunction
+import org.jetbrains.kotlin.ir.declarations.IrValueParameter
+import org.jetbrains.kotlin.ir.declarations.impl.IrFunctionImpl
+import org.jetbrains.kotlin.ir.declarations.impl.IrTypeParameterImpl
+import org.jetbrains.kotlin.ir.declarations.impl.IrValueParameterImpl
+import org.jetbrains.kotlin.ir.expressions.IrCall
+import org.jetbrains.kotlin.ir.expressions.IrExpression
+import org.jetbrains.kotlin.ir.expressions.IrFunctionExpression
+import org.jetbrains.kotlin.ir.expressions.IrStatementOrigin
+import org.jetbrains.kotlin.ir.expressions.impl.IrFunctionReferenceImpl
+import org.jetbrains.kotlin.ir.expressions.typeParametersCount
+import org.jetbrains.kotlin.ir.symbols.IrClassSymbol
+import org.jetbrains.kotlin.ir.symbols.IrConstructorSymbol
+import org.jetbrains.kotlin.ir.symbols.IrFunctionSymbol
+import org.jetbrains.kotlin.ir.symbols.IrSimpleFunctionSymbol
+import org.jetbrains.kotlin.ir.symbols.impl.IrSimpleFunctionSymbolImpl
+import org.jetbrains.kotlin.ir.symbols.impl.IrTypeParameterSymbolImpl
+import org.jetbrains.kotlin.ir.types.IrType
+import org.jetbrains.kotlin.ir.util.ConstantValueGenerator
+import org.jetbrains.kotlin.ir.util.DeepCopySymbolRemapper
+import org.jetbrains.kotlin.ir.util.TypeTranslator
+import org.jetbrains.kotlin.ir.util.endOffset
+import org.jetbrains.kotlin.ir.util.hasAnnotation
+import org.jetbrains.kotlin.ir.util.referenceFunction
+import org.jetbrains.kotlin.ir.util.startOffset
+import org.jetbrains.kotlin.ir.visitors.IrElementTransformerVoid
+import org.jetbrains.kotlin.name.ClassId
+import org.jetbrains.kotlin.name.FqName
+import org.jetbrains.kotlin.resolve.BindingTrace
+import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlin.resolve.DescriptorFactory
+import org.jetbrains.kotlin.resolve.descriptorUtil.fqNameSafe
+import org.jetbrains.kotlin.resolve.isInlineClassType
+import org.jetbrains.kotlin.resolve.unsubstitutedUnderlyingType
+import org.jetbrains.kotlin.types.KotlinType
+import org.jetbrains.kotlin.types.isError
+import org.jetbrains.kotlin.types.isNullable
+import org.jetbrains.kotlin.types.typeUtil.isTypeParameter
+import org.jetbrains.kotlin.types.typeUtil.makeNotNullable
+
+abstract class AbstractComposeLowering(
+    val context: JvmBackendContext,
+    val symbolRemapper: DeepCopySymbolRemapper,
+    val bindingTrace: BindingTrace
+) : IrElementTransformerVoid() {
+
+    protected val typeTranslator =
+        TypeTranslator(
+            context.ir.symbols.externalSymbolTable,
+            context.state.languageVersionSettings,
+            context.builtIns
+        ).apply {
+            constantValueGenerator = ConstantValueGenerator(
+                context.state.module,
+                context.ir.symbols.externalSymbolTable
+            )
+            constantValueGenerator.typeTranslator = this
+        }
+
+    protected val builtIns = context.irBuiltIns
+
+    protected val composerTypeDescriptor = context.state.module.findClassAcrossModuleDependencies(
+        ClassId.topLevel(ComposeFqNames.Composer)
+    ) ?: error("Cannot find the Composer class")
+
+    private val symbolTable get() = context.ir.symbols.externalSymbolTable
+
+    fun referenceFunction(descriptor: CallableDescriptor): IrFunctionSymbol {
+        return symbolRemapper.getReferencedFunction(symbolTable.referenceFunction(descriptor))
+    }
+
+    fun referenceSimpleFunction(descriptor: SimpleFunctionDescriptor): IrSimpleFunctionSymbol {
+        return symbolRemapper.getReferencedSimpleFunction(
+            symbolTable.referenceSimpleFunction(descriptor)
+        )
+    }
+
+    fun referenceConstructor(descriptor: ClassConstructorDescriptor): IrConstructorSymbol {
+        return symbolRemapper.getReferencedConstructor(symbolTable.referenceConstructor(descriptor))
+    }
+
+    fun getTopLevelClass(fqName: FqName): IrClassSymbol {
+        val descriptor = context.state.module.getPackage(fqName.parent()).memberScope
+            .getContributedClassifier(
+                fqName.shortName(), NoLookupLocation.FROM_BACKEND
+            ) as ClassDescriptor? ?: error("Class is not found: $fqName")
+        return symbolTable.referenceClass(descriptor)
+    }
+
+    fun getTopLevelFunction(fqName: FqName): IrFunctionSymbol {
+        val descriptor = context.state.module.getPackage(fqName.parent()).memberScope
+            .getContributedFunctions(
+                fqName.shortName(), NoLookupLocation.FROM_BACKEND
+            ).singleOrNull() ?: error("Function not found $fqName")
+        return symbolTable.referenceSimpleFunction(descriptor)
+    }
+
+    fun KotlinType.toIrType(): IrType = typeTranslator.translateType(this)
+
+    fun IrValueParameter.isComposerParam(): Boolean =
+        (descriptor as? ValueParameterDescriptor)?.isComposerParam() ?: false
+
+    fun ValueParameterDescriptor.isComposerParam(): Boolean =
+        name == KtxNameConventions.COMPOSER_PARAMETER &&
+                type.constructor.declarationDescriptor?.fqNameSafe == ComposeFqNames.Composer
+
+    fun IrAnnotationContainer.hasComposableAnnotation(): Boolean {
+        return annotations.hasAnnotation(ComposeFqNames.Composable)
+    }
+
+    fun IrCall.isTransformedComposableCall(): Boolean {
+        return context.irTrace[ComposeWritableSlices.IS_COMPOSABLE_CALL, this] ?: false
+    }
+
+    fun IrCall.isSyntheticComposableCall(): Boolean {
+        return context.irTrace[ComposeWritableSlices.IS_SYNTHETIC_COMPOSABLE_CALL, this] == true
+    }
+
+    private val composableChecker = ComposableAnnotationChecker()
+
+    protected val KotlinType.isEnum
+        get() =
+            (constructor.declarationDescriptor as? ClassDescriptor)?.kind == ClassKind.ENUM_CLASS
+
+    protected fun KotlinType?.isStable(): Boolean {
+        if (this == null) return false
+
+        val trace = bindingTrace
+        val calculated = trace.get(ComposeWritableSlices.STABLE_TYPE, this)
+        return if (calculated == null) {
+            val isStable = !isError &&
+                    !isTypeParameter() &&
+                    !isSpecialType &&
+                    (
+                            KotlinBuiltIns.isPrimitiveType(this) ||
+                                    isFunctionOrKFunctionType ||
+                                    isEnum ||
+                                    KotlinBuiltIns.isString(this) ||
+                                    isMarkedStable() ||
+                                    (
+                                            isNullable() &&
+                                                    makeNotNullable().isStable()
+                                            ) ||
+                                    (
+                                            isInlineClassType() &&
+                                                    unsubstitutedUnderlyingType().isStable()
+                                            )
+                            )
+            trace.record(ComposeWritableSlices.STABLE_TYPE, this, isStable)
+            isStable
+        } else calculated
+    }
+
+    fun FunctionDescriptor.isComposable(): Boolean {
+        val composability = composableChecker.analyze(bindingTrace, this)
+        return when (composability) {
+            ComposableAnnotationChecker.Composability.NOT_COMPOSABLE -> false
+            ComposableAnnotationChecker.Composability.MARKED -> true
+            ComposableAnnotationChecker.Composability.INFERRED -> true
+        }
+    }
+
+    fun IrFunction.isComposable(): Boolean = descriptor.isComposable()
+    fun IrFunctionExpression.isComposable(): Boolean = function.isComposable()
+
+    private fun IrFunction.createParameterDeclarations() {
+        fun ParameterDescriptor.irValueParameter() = IrValueParameterImpl(
+            this.startOffset ?: UNDEFINED_OFFSET,
+            this.endOffset ?: UNDEFINED_OFFSET,
+            IrDeclarationOrigin.DEFINED,
+            this,
+            type.toIrType(),
+            (this as? ValueParameterDescriptor)?.varargElementType?.toIrType()
+        ).also {
+            it.parent = this@createParameterDeclarations
+        }
+
+        fun TypeParameterDescriptor.irTypeParameter() = IrTypeParameterImpl(
+            this.startOffset ?: UNDEFINED_OFFSET,
+            this.endOffset ?: UNDEFINED_OFFSET,
+            IrDeclarationOrigin.DEFINED,
+            IrTypeParameterSymbolImpl(this)
+        ).also {
+            it.parent = this@createParameterDeclarations
+        }
+
+        dispatchReceiverParameter = descriptor.dispatchReceiverParameter?.irValueParameter()
+        extensionReceiverParameter = descriptor.extensionReceiverParameter?.irValueParameter()
+
+        assert(valueParameters.isEmpty())
+        valueParameters = descriptor.valueParameters.map { it.irValueParameter() }
+
+        assert(typeParameters.isEmpty())
+        typeParameters = descriptor.typeParameters.map { it.irTypeParameter() }
+    }
+
+    protected fun IrBuilderWithScope.irLambdaExpression(
+        descriptor: FunctionDescriptor,
+        type: IrType,
+        body: IrBlockBodyBuilder.(IrFunction) -> Unit
+    ) = irLambdaExpression(this.startOffset, this.endOffset, descriptor, type, body)
+
+    protected fun IrBuilderWithScope.irLambdaExpression(
+        startOffset: Int,
+        endOffset: Int,
+        descriptor: FunctionDescriptor,
+        type: IrType,
+        body: IrBlockBodyBuilder.(IrFunction) -> Unit
+    ): IrExpression {
+        val symbol = IrSimpleFunctionSymbolImpl(descriptor)
+
+        val returnType = descriptor.returnType!!.toIrType()
+
+        val lambda = IrFunctionImpl(
+            startOffset, endOffset,
+            IrDeclarationOrigin.LOCAL_FUNCTION_FOR_LAMBDA,
+            symbol,
+            returnType
+        ).also {
+            it.parent = scope.getLocalDeclarationParent()
+            it.createParameterDeclarations()
+            it.body = DeclarationIrBuilder(this@AbstractComposeLowering.context as IrGeneratorContext, symbol)
+                .irBlockBody { body(it) }
+        }
+
+        return irBlock(
+            startOffset = startOffset,
+            endOffset = endOffset,
+            origin = IrStatementOrigin.LAMBDA,
+            resultType = type
+        ) {
+            +lambda
+            +IrFunctionReferenceImpl(
+                startOffset = startOffset,
+                endOffset = endOffset,
+                type = type,
+                symbol = symbol,
+                typeArgumentsCount = descriptor.typeParametersCount,
+                reflectionTarget = symbol,
+                origin = IrStatementOrigin.LAMBDA
+            )
+        }
+    }
+
+    protected fun IrBuilderWithScope.createFunctionDescriptor(
+        type: KotlinType,
+        owner: DeclarationDescriptor = scope.scopeOwner
+    ): FunctionDescriptor {
+        return AnonymousFunctionDescriptor(
+            owner,
+            Annotations.EMPTY,
+            CallableMemberDescriptor.Kind.SYNTHESIZED,
+            SourceElement.NO_SOURCE,
+            false
+        ).apply {
+            initialize(
+                type.getReceiverTypeFromFunctionType()?.let {
+                    DescriptorFactory.createExtensionReceiverParameterForCallable(
+                        this,
+                        it,
+                        Annotations.EMPTY
+                    )
+                },
+                null,
+                emptyList(),
+                type.getValueParameterTypesFromFunctionType().mapIndexed { i, t ->
+                    ValueParameterDescriptorImpl(
+                        containingDeclaration = this,
+                        original = null,
+                        index = i,
+                        annotations = Annotations.EMPTY,
+                        name = t.type.extractParameterNameFromFunctionTypeArgument()
+                            ?: Name.identifier("p$i"),
+                        outType = t.type,
+                        declaresDefaultValue = false,
+                        isCrossinline = false,
+                        isNoinline = false,
+                        varargElementType = null,
+                        source = SourceElement.NO_SOURCE
+                    )
+                },
+                type.getReturnTypeFromFunctionType(),
+                Modality.FINAL,
+                Visibilities.LOCAL,
+                null
+            )
+            isOperator = false
+            isInfix = false
+            isExternal = false
+            isInline = false
+            isTailrec = false
+            isSuspend = false
+            isExpect = false
+            isActual = false
+        }
+    }
+}

plugins/compose-compiler-hosted/src/main/java/androidx/compose/plugins/kotlin/compiler/lower/ComposableCallTransformer.kt

@@ -0,0 +1,994 @@
+package androidx.compose.plugins.kotlin.compiler.lower
+
+import androidx.compose.plugins.kotlin.ComposableEmitMetadata
+import androidx.compose.plugins.kotlin.ComposeFqNames
+import androidx.compose.plugins.kotlin.EmitChildrenValueParameterDescriptor
+import androidx.compose.plugins.kotlin.KtxNameConventions
+import androidx.compose.plugins.kotlin.ValidatedAssignment
+import androidx.compose.plugins.kotlin.ValidationType
+import androidx.compose.plugins.kotlin.analysis.ComposeWritableSlices
+import androidx.compose.plugins.kotlin.hasPivotalAnnotation
+import androidx.compose.plugins.kotlin.irTrace
+import org.jetbrains.kotlin.backend.common.FileLoweringPass
+import org.jetbrains.kotlin.backend.common.deepCopyWithVariables
+import org.jetbrains.kotlin.backend.common.ir.copyTo
+import org.jetbrains.kotlin.backend.common.ir.createImplicitParameterDeclarationWithWrappedDescriptor
+import org.jetbrains.kotlin.backend.common.lower.DeclarationIrBuilder
+import org.jetbrains.kotlin.backend.common.pop
+import org.jetbrains.kotlin.backend.common.push
+import org.jetbrains.kotlin.backend.jvm.JvmBackendContext
+import org.jetbrains.kotlin.backend.jvm.JvmLoweredDeclarationOrigin
+import org.jetbrains.kotlin.builtins.getReceiverTypeFromFunctionType
+import org.jetbrains.kotlin.descriptors.ClassConstructorDescriptor
+import org.jetbrains.kotlin.descriptors.FunctionDescriptor
+import org.jetbrains.kotlin.descriptors.Modality
+import org.jetbrains.kotlin.descriptors.PropertyDescriptor
+import org.jetbrains.kotlin.descriptors.ValueParameterDescriptor
+import org.jetbrains.kotlin.descriptors.Visibilities
+import org.jetbrains.kotlin.ir.IrStatement
+import org.jetbrains.kotlin.ir.builders.*
+import org.jetbrains.kotlin.ir.builders.declarations.addConstructor
+import org.jetbrains.kotlin.ir.builders.declarations.addFunction
+import org.jetbrains.kotlin.ir.builders.declarations.buildClass
+import org.jetbrains.kotlin.ir.declarations.IrConstructor
+import org.jetbrains.kotlin.ir.declarations.IrDeclaration
+import org.jetbrains.kotlin.ir.declarations.IrFile
+import org.jetbrains.kotlin.ir.declarations.IrFunction
+import org.jetbrains.kotlin.ir.declarations.IrModuleFragment
+import org.jetbrains.kotlin.ir.declarations.IrSimpleFunction
+import org.jetbrains.kotlin.ir.declarations.IrSymbolOwner
+import org.jetbrains.kotlin.ir.declarations.IrValueParameter
+import org.jetbrains.kotlin.ir.declarations.getIrValueParameter
+import org.jetbrains.kotlin.ir.expressions.IrCall
+import org.jetbrains.kotlin.ir.expressions.IrConst
+import org.jetbrains.kotlin.ir.expressions.IrExpression
+import org.jetbrains.kotlin.ir.expressions.IrFunctionExpression
+import org.jetbrains.kotlin.ir.expressions.IrGetValue
+import org.jetbrains.kotlin.ir.expressions.IrReturn
+import org.jetbrains.kotlin.ir.expressions.IrStatementOrigin
+import org.jetbrains.kotlin.ir.expressions.copyTypeArgumentsFrom
+import org.jetbrains.kotlin.ir.expressions.getValueArgument
+import org.jetbrains.kotlin.ir.expressions.impl.IrFunctionReferenceImpl
+import org.jetbrains.kotlin.ir.expressions.impl.IrInstanceInitializerCallImpl
+import org.jetbrains.kotlin.ir.expressions.putTypeArguments
+import org.jetbrains.kotlin.ir.expressions.putValueArgument
+import org.jetbrains.kotlin.ir.expressions.typeParametersCount
+import org.jetbrains.kotlin.ir.symbols.impl.IrSimpleFunctionSymbolImpl
+import org.jetbrains.kotlin.ir.types.IrSimpleType
+import org.jetbrains.kotlin.ir.types.IrTypeProjection
+import org.jetbrains.kotlin.ir.types.getClass
+import org.jetbrains.kotlin.ir.types.isUnit
+import org.jetbrains.kotlin.ir.types.toKotlinType
+import org.jetbrains.kotlin.ir.types.typeWith
+import org.jetbrains.kotlin.ir.util.DeepCopySymbolRemapper
+import org.jetbrains.kotlin.ir.util.constructors
+import org.jetbrains.kotlin.ir.util.defaultType
+import org.jetbrains.kotlin.ir.util.explicitParameters
+import org.jetbrains.kotlin.ir.util.findAnnotation
+import org.jetbrains.kotlin.ir.util.fqNameWhenAvailable
+import org.jetbrains.kotlin.ir.util.functions
+import org.jetbrains.kotlin.ir.util.getArgumentsWithIr
+import org.jetbrains.kotlin.ir.util.isSuspend
+import org.jetbrains.kotlin.ir.util.parentAsClass
+import org.jetbrains.kotlin.ir.util.statements
+import org.jetbrains.kotlin.ir.util.substitute
+import org.jetbrains.kotlin.ir.util.typeSubstitutionMap
+import org.jetbrains.kotlin.ir.visitors.IrElementTransformerVoid
+import org.jetbrains.kotlin.ir.visitors.transformChildrenVoid
+import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlin.psi2ir.findFirstFunction
+import org.jetbrains.kotlin.resolve.BindingTrace
+import org.jetbrains.kotlin.resolve.calls.tasks.ExplicitReceiverKind
+import org.jetbrains.kotlin.resolve.descriptorUtil.fqNameSafe
+import org.jetbrains.kotlin.types.typeUtil.isUnit
+
+class ComposableCallTransformer(
+    context: JvmBackendContext,
+    symbolRemapper: DeepCopySymbolRemapper,
+    bindingTrace: BindingTrace
+) :
+    AbstractComposeLowering(context, symbolRemapper, bindingTrace),
+    FileLoweringPass,
+    ModuleLoweringPass {
+
+    override fun lower(module: IrModuleFragment) {
+        module.transformChildrenVoid(this)
+    }
+
+    private val orFunctionDescriptor = builtIns.builtIns.booleanType.memberScope
+        .findFirstFunction("or") { it is FunctionDescriptor && it.isInfix }
+
+    override fun lower(irFile: IrFile) {
+        irFile.transformChildrenVoid(this)
+    }
+
+    private val declarationStack = mutableListOf<IrSymbolOwner>()
+
+    override fun visitDeclaration(declaration: IrDeclaration): IrStatement {
+        if (declaration !is IrSymbolOwner) return super.visitDeclaration(declaration)
+        try {
+            declarationStack.push(declaration)
+            return super.visitDeclaration(declaration)
+        } finally {
+            declarationStack.pop()
+        }
+    }
+
+    override fun visitCall(expression: IrCall): IrExpression {
+        if (expression.isTransformedComposableCall() || expression.isSyntheticComposableCall()) {
+            val descriptor = expression.symbol.descriptor
+            val returnType = descriptor.returnType
+            val isUnit = returnType == null || returnType.isUnit() || expression.type.isUnit()
+            val isInline = descriptor.isInline || context.irTrace[
+                    ComposeWritableSlices.IS_INLINE_COMPOSABLE_CALL,
+                    expression
+            ] == true
+            return if (isUnit && !isInline) {
+                DeclarationIrBuilder(context as IrGeneratorContext, declarationStack.last().symbol).irBlock {
+                    +irComposableCall(expression.transformChildren())
+                }
+            } else {
+                val call = if (isInline)
+                    expression.transformChildrenWithoutConvertingLambdas()
+                else
+                    expression.transformChildren()
+                DeclarationIrBuilder(context, declarationStack.last().symbol)
+                    .irComposableExpr(call)
+            }
+        }
+
+        val emitMetadata = context.irTrace[
+                ComposeWritableSlices.COMPOSABLE_EMIT_METADATA,
+                expression
+        ]
+        if (emitMetadata != null) {
+            return DeclarationIrBuilder(context, declarationStack.last().symbol).irBlock {
+                +irComposableEmit(expression.transformChildren(), emitMetadata)
+            }
+        }
+        return super.visitCall(expression)
+    }
+
+    private fun IrCall.transformChildrenWithoutConvertingLambdas(): IrCall {
+        dispatchReceiver = dispatchReceiver?.transform(this@ComposableCallTransformer, null)
+        extensionReceiver = extensionReceiver?.transform(this@ComposableCallTransformer, null)
+        for (i in 0 until valueArgumentsCount) {
+            val arg = getValueArgument(i) ?: continue
+            if (arg is IrFunctionExpression) {
+                // we convert function expressions into their lowered lambda equivalents, but we
+                // want to avoid doing this for inlined lambda calls.
+                putValueArgument(i, super.visitFunctionExpression(arg))
+            } else {
+                putValueArgument(i, arg.transform(this@ComposableCallTransformer, null))
+            }
+        }
+        return this
+    }
+
+    override fun visitFunctionExpression(expression: IrFunctionExpression): IrExpression {
+        if (expression.origin == IrStatementOrigin.LAMBDA) {
+            if (expression.function.valueParameters.lastOrNull()?.isComposerParam() == true) {
+                return DeclarationIrBuilder(context, declarationStack.last().symbol).irBlock {
+                    +covertLambdaIfNecessary(expression.transformChildren())
+                }
+            }
+        }
+        return super.visitFunctionExpression(expression)
+    }
+
+    private fun IrBlockBuilder.irComposableCall(
+        original: IrCall
+    ): IrExpression {
+        val composerArg = original.getValueArgument(original.valueArgumentsCount - 1)!!
+        // TODO(lmr): we may want to rewrite this in a way that doesn't do a deepCopy...
+        val getComposer = { composerArg.deepCopyWithVariables() }
+        return irComposableCallBase(
+            original,
+            getComposer
+        )
+    }
+
+    private fun IrBlockBuilder.irComposableCallBase(
+        original: IrCall,
+        getComposer: () -> IrExpression
+    ): IrExpression {
+
+        /*
+
+        Foo(text="foo")
+
+        // transforms into
+
+        val attr_text = "foo"
+        composer.call(
+            key = 123,
+            invalid = { changed(attr_text) },
+            block = { Foo(attr_text) }
+        )
+         */
+        // TODO(lmr): the way we grab temporaries here feels wrong. We should investigate the right
+        // way to do this. Additionally, we are creating temporary vars for variables which is
+        // causing larger stack space than needed in our generated code.
+
+        val irGetArguments = original
+            .symbol
+            .descriptor
+            .valueParameters
+            .map {
+                val arg = original.getValueArgument(it)
+                it to getParameterExpression(it, arg)
+            }
+
+        val tmpDispatchReceiver = original.dispatchReceiver?.let { irTemporary(it) }
+        val tmpExtensionReceiver = original.extensionReceiver?.let { irTemporary(it) }
+
+        val callDescriptor = composerTypeDescriptor
+            .unsubstitutedMemberScope
+            .findFirstFunction(KtxNameConventions.CALL.identifier) {
+                it.valueParameters.size == 3
+            }
+
+        val joinKeyDescriptor = composerTypeDescriptor
+            .unsubstitutedMemberScope
+            .findFirstFunction(KtxNameConventions.JOINKEY.identifier) {
+                it.valueParameters.size == 2
+            }
+
+        val callParameters = callDescriptor.valueParameters
+            .map { it.name to it }
+            .toMap()
+
+        fun getCallParameter(name: Name) = callParameters[name]
+            ?: error("Expected $name parameter to exist")
+
+        return irCall(
+            callee = referenceFunction(callDescriptor),
+            type = builtIns.unitType // TODO(lmr): refactor call(...) to return a type
+        ).apply {
+            dispatchReceiver = getComposer()
+
+            putValueArgument(
+                getCallParameter(KtxNameConventions.CALL_KEY_PARAMETER),
+                irGroupKey(
+                    original = original,
+                    getComposer = getComposer,
+                    joinKey = joinKeyDescriptor,
+                    pivotals = irGetArguments.mapNotNull { (param, getExpr) ->
+                        if (!param.hasPivotalAnnotation()) null
+                        else getExpr()
+                    }
+                )
+            )
+
+            val invalidParameter = getCallParameter(KtxNameConventions.CALL_INVALID_PARAMETER)
+
+            val validatorType = invalidParameter.type.getReceiverTypeFromFunctionType()
+                ?: error("Expected validator type to be on receiver of the invalid lambda")
+
+            val changedDescriptor = validatorType
+                .memberScope
+                .findFirstFunction("changed") { it.typeParametersCount == 1 }
+
+            val validatedArguments = irGetArguments
+                .take(irGetArguments.size - 1)
+                .mapNotNull { (_, getExpr) -> getExpr() } +
+                    listOfNotNull(
+                        tmpDispatchReceiver?.let { irGet(it) },
+                        tmpExtensionReceiver?.let { irGet(it) }
+                    )
+
+            val isSkippable = validatedArguments.all { it.type.toKotlinType().isStable() }
+
+            putValueArgument(
+                invalidParameter,
+                irLambdaExpression(
+                    original.startOffset,
+                    original.endOffset,
+                    descriptor = createFunctionDescriptor(
+                        type = invalidParameter.type,
+                        owner = symbol.descriptor.containingDeclaration
+                    ),
+                    type = invalidParameter.type.toIrType()
+                ) { fn ->
+                    if (!isSkippable) {
+                        // if it's not skippable, we don't validate any arguments.
+                        +irReturn(irTrue())
+                    } else {
+                        val validationCalls = validatedArguments
+                            .map {
+                                irChangedCall(
+                                    changedDescriptor = changedDescriptor,
+                                    receiver = irGet(fn.extensionReceiverParameter!!),
+                                    attributeValue = it
+                                )
+                            }
+
+                        // all as one expression: a or b or c ... or z
+                        +irReturn(when (validationCalls.size) {
+                            0 -> irFalse()
+                            1 -> validationCalls.single()
+                            else -> validationCalls.reduce { accumulator, value ->
+                                when {
+                                    // if it is a constant, the value is `false`
+                                    accumulator is IrConst<*> -> value
+                                    value is IrConst<*> -> accumulator
+                                    else -> irOr(accumulator, value)
+                                }
+                            }
+                        })
+                    }
+                }
+            )
+
+            val blockParameter = getCallParameter(KtxNameConventions.CALL_BLOCK_PARAMETER)
+
+            putValueArgument(
+                blockParameter,
+                irLambdaExpression(
+                    original.startOffset,
+                    original.endOffset,
+                    descriptor = createFunctionDescriptor(
+                        type = blockParameter.type,
+                        owner = symbol.descriptor.containingDeclaration
+                    ),
+                    type = blockParameter.type.toIrType()
+                ) {
+                    +irCall(
+                        callee = IrSimpleFunctionSymbolImpl(original.symbol.descriptor).also {
+                            it.bind(original.symbol.owner as IrSimpleFunction)
+                        },
+                        type = original.type
+                    ).apply {
+                        copyTypeArgumentsFrom(original)
+
+                        dispatchReceiver = tmpDispatchReceiver?.let { irGet(it) }
+                        extensionReceiver = tmpExtensionReceiver?.let { irGet(it) }
+
+                        irGetArguments.forEach { (param, getExpr) ->
+                            putValueArgument(param, getExpr())
+                        }
+                    }
+                }
+            )
+        }
+    }
+
+    private fun DeclarationIrBuilder.irComposableExpr(
+        original: IrCall
+    ): IrExpression {
+        return irBlock(resultType = original.type) {
+            val composerParam = nearestComposer()
+            val getComposer = { irGet(composerParam) }
+            irComposableExprBase(
+                original,
+                getComposer
+            )
+        }
+    }
+
+    private fun IrBlockBuilder.irComposableExprBase(
+        original: IrCall,
+        getComposer: () -> IrExpression
+    ) {
+        /*
+
+        Foo(text="foo")
+
+        // transforms into
+
+        composer.startExpr(123)
+        val result = Foo(text="foo")
+        composer.endExpr()
+        result
+         */
+
+        // TODO(lmr): the way we grab temporaries here feels wrong. We should investigate the right
+        // way to do this. Additionally, we are creating temporary vars for variables which is
+        // causing larger stack space than needed in our generated code.
+
+        // for composableExpr, we only need to create temporaries if there are any pivotals
+        val hasPivotals = original
+            .symbol
+            .descriptor
+            .valueParameters
+            .any { it.hasPivotalAnnotation() }
+
+        // if we don't have any pivotal parameters, we don't use the parameters more than once,
+        // so we can just use the original call itself.
+        val irGetArguments = original
+            .symbol
+            .descriptor
+            .valueParameters
+            .map {
+                val arg = original.getValueArgument(it)
+                val expr = if (hasPivotals)
+                    getParameterExpression(it, arg)
+                else
+                    ({ arg })
+                it to expr
+            }
+
+        val startExpr = composerTypeDescriptor
+            .unsubstitutedMemberScope
+            .findFirstFunction(KtxNameConventions.START_EXPR.identifier) {
+                it.valueParameters.size == 1
+            }
+
+        val endExpr = composerTypeDescriptor
+            .unsubstitutedMemberScope
+            .findFirstFunction(KtxNameConventions.END_EXPR.identifier) {
+                it.valueParameters.size == 0
+            }
+
+        val joinKeyDescriptor = composerTypeDescriptor
+            .unsubstitutedMemberScope
+            .findFirstFunction(KtxNameConventions.JOINKEY.identifier) {
+                it.valueParameters.size == 2
+            }
+
+        val startCall = irCall(
+            callee = referenceFunction(startExpr),
+            type = builtIns.unitType
+        ).apply {
+            dispatchReceiver = getComposer()
+            putValueArgument(
+                startExpr.valueParameters.first(),
+                irGroupKey(
+                    original = original,
+                    getComposer = getComposer,
+                    joinKey = joinKeyDescriptor,
+                    pivotals = irGetArguments.mapNotNull { (param, getExpr) ->
+                        if (!param.hasPivotalAnnotation()) null
+                        else getExpr()
+                    }
+                )
+            )
+        }
+
+        val newCall = if (hasPivotals) irCall(
+            callee = IrSimpleFunctionSymbolImpl(original.symbol.descriptor).also {
+                it.bind(original.symbol.owner as IrSimpleFunction)
+            },
+            type = original.type
+        ).apply {
+            copyTypeArgumentsFrom(original)
+
+            dispatchReceiver = original.dispatchReceiver
+            extensionReceiver = original.extensionReceiver
+
+            irGetArguments.forEach { (param, getExpr) ->
+                putValueArgument(param, getExpr())
+            }
+        } else original
+
+        val endCall = irCall(
+            callee = referenceFunction(endExpr),
+            type = builtIns.unitType
+        ).apply {
+            dispatchReceiver = getComposer()
+        }
+
+        val hasResult = !original.type.isUnit()
+
+        if (hasResult) {
+            +startCall
+            val tmpResult = irTemporary(newCall, irType = original.type)
+            +endCall
+            +irGet(tmpResult)
+        } else {
+            +startCall
+            +newCall
+            +endCall
+        }
+    }
+
+    private fun isChildrenParameter(desc: ValueParameterDescriptor, expr: IrExpression): Boolean {
+        return expr is IrFunctionExpression &&
+                expr.origin == IrStatementOrigin.LAMBDA &&
+                desc is EmitChildrenValueParameterDescriptor
+    }
+
+    private fun IrBlockBuilder.getParameterExpression(
+        desc: ValueParameterDescriptor,
+        expr: IrExpression?
+    ): () -> IrExpression? {
+        if (expr == null)
+            return { null }
+        return when {
+            expr is IrConst<*> ->
+                ({ expr.copy() })
+            isChildrenParameter(desc, expr) ->
+                ({ expr })
+            else -> {
+                val temp = irTemporary(
+                    covertLambdaIfNecessary(expr),
+                    irType = expr.type
+                )
+                ({ irGet(temp) })
+            }
+        }
+    }
+
+    private fun nearestComposer(): IrValueParameter {
+        for (fn in declarationStack.asReversed().asSequence()) {
+            if (fn is IrFunction) {
+                val param = fn.valueParameters.lastOrNull()
+                if (param != null && param.isComposerParam()) {
+                    return param
+                }
+            }
+        }
+        error("Couldn't find composer parameter")
+    }
+
+    private fun IrBlockBuilder.irComposableEmit(
+        original: IrCall,
+        emitMetadata: ComposableEmitMetadata
+    ): IrExpression {
+        val composerParam = nearestComposer()
+        return irComposableEmitBase(
+            original,
+            { irGet(composerParam) },
+            emitMetadata
+        )
+    }
+
+    private fun IrBlockBuilder.irComposableEmitBase(
+        original: IrCall,
+        getComposer: () -> IrExpression,
+        emitMetadata: ComposableEmitMetadata
+    ): IrExpression {
+        /*
+
+        TextView(text="foo")
+
+        // transforms into
+
+        val attr_text = "foo"
+        composer.emit(
+            key = 123,
+            ctor = { context -> TextView(context) },
+            update = { set(attr_text) { text -> this.text = text } }
+        )
+         */
+        val parametersByName = original
+            .symbol
+            .descriptor
+            .valueParameters
+            .mapNotNull { desc ->
+                original.getValueArgument(desc)?.let { desc to it }
+            }
+            .map { (desc, expr) ->
+                desc.name.asString() to getParameterExpression(desc, expr)
+            }
+            .toMap()
+
+        val emitCall = emitMetadata.emitCall
+        val emitFunctionDescriptor = emitCall.candidateDescriptor
+
+        val emitParameters = emitFunctionDescriptor.valueParameters
+            .map { it.name to it }
+            .toMap()
+
+        fun getEmitParameter(name: Name) = emitParameters[name]
+            ?: error("Expected $name parameter to exist")
+
+        val emitFunctionSymbol = referenceFunction(emitFunctionDescriptor)
+
+        val joinKeyDescriptor = composerTypeDescriptor
+            .unsubstitutedMemberScope
+            .findFirstFunction(KtxNameConventions.JOINKEY.identifier) {
+                it.valueParameters.size == 2
+            }
+
+        fun irGetParameter(name: String): IrExpression = parametersByName[name]?.invoke()
+            ?: error("No parameter found with name $name")
+
+        return irCall(
+            callee = emitFunctionSymbol,
+            type = builtIns.unitType
+        ).apply {
+            dispatchReceiver = getComposer()
+            // TODO(lmr): extensionReceiver.
+            // We would want to do this to enable "emit" and "call" implementations that are
+            // extensions on the composer
+
+            putTypeArguments(emitCall.typeArguments) { it.toIrType() }
+
+            putValueArgument(
+                getEmitParameter(KtxNameConventions.EMIT_KEY_PARAMETER),
+                irGroupKey(
+                    original = original,
+                    getComposer = getComposer,
+                    joinKey = joinKeyDescriptor,
+                    pivotals = emitMetadata.pivotals.map { irGetParameter(it) }
+                )
+            )
+
+            val ctorParam = getEmitParameter(KtxNameConventions.EMIT_CTOR_PARAMETER)
+
+            val ctorLambdaDescriptor = createFunctionDescriptor(ctorParam.type)
+
+            putValueArgument(
+                ctorParam,
+                irLambdaExpression(
+                    original.startOffset,
+                    original.endOffset,
+                    descriptor = ctorLambdaDescriptor,
+                    type = ctorParam.type.toIrType()
+                ) { fn ->
+
+                    val ctorCall = emitMetadata.ctorCall
+
+                    val ctorCallSymbol = referenceConstructor(
+                        ctorCall.candidateDescriptor as ClassConstructorDescriptor
+                    )
+
+                    +irReturn(irCall(ctorCallSymbol).apply {
+                        putTypeArguments(ctorCall.typeArguments) { it.toIrType() }
+                        ctorLambdaDescriptor.valueParameters.zip(
+                            ctorCall
+                                .candidateDescriptor!!
+                                .valueParameters
+                        ) { a, b ->
+                            putValueArgument(
+                                b,
+                                irGet(fn.getIrValueParameter(a))
+                            )
+                        }
+                        emitMetadata.ctorParams.forEach { name ->
+                            val param = ctorCall
+                                .candidateDescriptor
+                                .valueParameters
+                                .firstOrNull { it.name.identifier == name }
+                            if (param != null) {
+                                putValueArgument(
+                                    param,
+                                    irGetParameter(name)
+                                )
+                            }
+                        }
+                    })
+                }
+            )
+
+            val updateParam = getEmitParameter(
+                KtxNameConventions
+                    .EMIT_UPDATER_PARAMETER
+            )
+
+            val updateLambdaDescriptor = createFunctionDescriptor(updateParam.type)
+
+            putValueArgument(
+                updateParam,
+                irLambdaExpression(
+                    original.startOffset,
+                    original.endOffset,
+                    descriptor = updateLambdaDescriptor,
+                    type = updateParam.type.toIrType()
+                ) { fn ->
+                    emitMetadata.validations.forEach {
+                        // set(attr_text) { text -> this.text = text }
+                        val arg = irGetParameter(it.name)
+                        +irValidatedAssignment(
+                            arg.startOffset,
+                            arg.endOffset,
+                            memoizing = true,
+                            validation = it,
+                            receiver = irGet(fn.extensionReceiverParameter!!),
+                            attributeValue = arg
+                        )
+                    }
+                    +irReturnUnit()
+                }
+            )
+
+            if (emitMetadata.hasChildren) {
+                val bodyParam = getEmitParameter(KtxNameConventions.EMIT_CHILDREN_PARAMETER)
+
+                val childrenExpr = irGetParameter("\$CHILDREN")
+
+                putValueArgument(
+                    bodyParam,
+                    childrenExpr
+                )
+            }
+        }
+    }
+
+    private fun IrBuilderWithScope.irGroupKey(
+        original: IrCall,
+        joinKey: FunctionDescriptor,
+        getComposer: () -> IrExpression,
+        pivotals: List<IrExpression>
+    ): IrExpression {
+        val keyValueExpression = irInt(original.sourceLocationHash())
+        return if (pivotals.isEmpty()) keyValueExpression
+        else (listOf(keyValueExpression) + pivotals).reduce { accumulator, value ->
+            irCall(
+                callee = referenceFunction(joinKey),
+                type = joinKey.returnType!!.toIrType()
+            ).apply {
+                dispatchReceiver = getComposer()
+                putValueArgument(0, accumulator)
+                putValueArgument(1, value)
+            }
+        }
+    }
+
+    private fun IrCall.sourceLocationHash(): Int {
+        return symbol.descriptor.fqNameSafe.toString().hashCode() xor startOffset
+    }
+
+    private fun IrBuilderWithScope.irChangedCall(
+        changedDescriptor: FunctionDescriptor,
+        receiver: IrExpression,
+        attributeValue: IrExpression
+    ): IrExpression {
+        // TODO(lmr): make it so we can use the "changed" overloads with primitive types.
+        // Right now this is causing a lot of boxing/unboxing for primitives
+        return if (attributeValue is IrConst<*>) irFalse()
+        else irCall(
+            callee = referenceFunction(changedDescriptor),
+            type = changedDescriptor.returnType?.toIrType()!!
+        ).apply {
+            putTypeArgument(0, attributeValue.type)
+            dispatchReceiver = receiver
+            putValueArgument(0, attributeValue)
+        }
+    }
+
+    private fun IrBuilderWithScope.irOr(
+        left: IrExpression,
+        right: IrExpression
+    ): IrExpression {
+        return irCall(
+            callee = referenceFunction(orFunctionDescriptor),
+            type = builtIns.booleanType
+        ).apply {
+            dispatchReceiver = left
+            putValueArgument(0, right)
+        }
+    }
+
+    private fun IrBuilderWithScope.irValidatedAssignment(
+        startOffset: Int,
+        endOffset: Int,
+        memoizing: Boolean,
+        validation: ValidatedAssignment,
+        receiver: IrExpression,
+        attributeValue: IrExpression
+    ): IrExpression {
+        // for emit, fnDescriptor is Validator.(Value) -> Unit    or Validator.(Value, Element.(Value) -> Unit) -> Unit
+        // for call, fnDescriptor is Validator.(Value) -> Boolean or Validator.(Value, (Value) -> Unit) -> Boolean
+
+        // in emit, the element is passed through an extension parameter
+        // in call, the element is passed through a capture scope
+        val validationCall =
+            if (memoizing) validation.validationCall
+            else validation.uncheckedValidationCall
+
+        if (validationCall == null) error("Expected validationCall to be non-null")
+
+        val validationCallDescriptor = validationCall.candidateDescriptor as FunctionDescriptor
+
+        return irCall(
+            callee = referenceFunction(validationCallDescriptor),
+            type = validationCallDescriptor.returnType?.toIrType()!!
+        ).apply {
+
+            dispatchReceiver = receiver
+            // TODO(lmr): extensionReceiver.
+            // This might be something we want to be able to do in the cases where we want to
+            // build extension `changed(...)`  or `set(..) { ... }` methods
+
+            putTypeArguments(validationCall.typeArguments) { it.toIrType() }
+
+            putValueArgument(0, attributeValue)
+            val assignment = validation.assignment
+            if (assignment != null && validation.validationType != ValidationType.CHANGED) {
+                val assignmentLambdaDescriptor = validation.assignmentLambda
+                    ?: error("Expected assignmentLambda to be non-null")
+                val assignmentDescriptor = assignment.candidateDescriptor.original
+
+                val assignmentSymbol = when (assignmentDescriptor) {
+                    is PropertyDescriptor -> referenceFunction(
+                        assignmentDescriptor.setter!!
+                    )
+                    else -> referenceFunction(assignmentDescriptor)
+                }
+                val assignmentValueParameterDescriptor = assignmentLambdaDescriptor
+                    .valueParameters[0]
+
+                putValueArgument(
+                    1,
+                    irLambdaExpression(
+                        startOffset,
+                        endOffset,
+                        assignmentLambdaDescriptor,
+                        validationCallDescriptor.valueParameters[1].type.toIrType()
+                    ) { fn ->
+                        +irReturn(
+                            irCall(
+                                callee = assignmentSymbol,
+                                type = builtIns.unitType
+                            ).apply {
+                                putTypeArguments(assignment.typeArguments) { it.toIrType() }
+                                when (assignment.explicitReceiverKind) {
+                                    ExplicitReceiverKind.DISPATCH_RECEIVER -> {
+                                        dispatchReceiver = irGet(fn.extensionReceiverParameter!!)
+                                    }
+                                    ExplicitReceiverKind.EXTENSION_RECEIVER -> {
+                                        extensionReceiver = irGet(fn.extensionReceiverParameter!!)
+                                    }
+                                    ExplicitReceiverKind.BOTH_RECEIVERS -> {
+                                        // NOTE(lmr): This should not be possible. This would have
+                                        // to be an extension method on the ComposerUpdater class
+                                        // itself for the emittable type.
+                                        error(
+                                            "Extension instance methods are not allowed for " +
+                                                    "assignments"
+                                        )
+                                    }
+                                    ExplicitReceiverKind.NO_EXPLICIT_RECEIVER -> {
+                                        // NOTE(lmr): This is not possible
+                                        error("Static methods are invalid for assignments")
+                                    }
+                                }
+                                putValueArgument(
+                                    0,
+                                    irGet(
+                                        fn.getIrValueParameter(assignmentValueParameterDescriptor)
+                                    )
+                                )
+                            }
+                        )
+                    }
+                )
+            }
+        }
+    }
+
+    /**
+     * Convert a function-reference into a inner class constructor call.
+     *
+     * This is a transformed copy of the work done in CallableReferenceLowering to allow the
+     * [ComposeObservePatcher] access to the this parameter.
+     */
+    private fun IrBlockBuilder.covertLambdaIfNecessary(expression: IrExpression): IrExpression {
+        val functionExpression = expression as? IrFunctionExpression ?: return expression
+
+        val function = functionExpression.function
+
+        if (!function.isComposable()) return expression
+
+        // A temporary node created so the code matches more closely to the
+        // CallableReferenceLowering code that was copied.
+        val functionReference = IrFunctionReferenceImpl(
+            -1,
+            -1,
+            expression.type,
+            function.symbol,
+            function.descriptor,
+            0,
+            expression.origin
+        )
+
+        val context = this@ComposableCallTransformer.context
+        val superType =
+            FakeJvmSymbols(context.state.module, context.irBuiltIns).lambdaClass.typeWith()
+        val parameterTypes = (functionExpression.type as IrSimpleType).arguments.map {
+            (it as IrTypeProjection).type
+        }
+        val functionSuperClass = FakeJvmSymbols(context.state.module, context.irBuiltIns)
+            .getJvmFunctionClass(
+            parameterTypes.size - 1
+        )
+        val jvmClass = functionSuperClass.typeWith(parameterTypes)
+        val boundReceiver = functionReference.getArgumentsWithIr().singleOrNull()
+        val typeArgumentsMap = functionReference.typeSubstitutionMap
+        val callee = functionReference.symbol.owner
+        var constructor: IrConstructor? = null
+        val irClass = buildClass {
+            setSourceRange(functionReference)
+            visibility = Visibilities.LOCAL
+            origin = JvmLoweredDeclarationOrigin.LAMBDA_IMPL
+            name = Name.special("<function reference to ${callee.fqNameWhenAvailable}>")
+        }.apply {
+            parent = scope.getLocalDeclarationParent()
+            superTypes += superType
+            superTypes += jvmClass
+            createImplicitParameterDeclarationWithWrappedDescriptor()
+        }.also { irClass ->
+            // Add constructor
+            val superConstructor = superType.getClass()!!.constructors.single {
+                it.valueParameters.size == if (boundReceiver != null) 2 else 1
+            }
+            constructor = irClass.addConstructor {
+                setSourceRange(functionReference)
+                origin = JvmLoweredDeclarationOrigin.FUNCTION_REFERENCE_IMPL
+                returnType = irClass.defaultType
+                isPrimary = true
+            }.apply {
+                boundReceiver?.first?.let { param ->
+                    valueParameters += param.copyTo(
+                        irFunction = this,
+                        index = 0,
+                        type = param.type.substitute(typeArgumentsMap)
+                    )
+                }
+                body = DeclarationIrBuilder(context, symbol).irBlockBody(startOffset, endOffset) {
+                    +irDelegatingConstructorCall(superConstructor).apply {
+                        putValueArgument(0, irInt(parameterTypes.size - 1))
+                        if (boundReceiver != null)
+                            putValueArgument(1, irGet(valueParameters.first()))
+                    }
+                    +IrInstanceInitializerCallImpl(
+                        startOffset,
+                        endOffset,
+                        irClass.symbol,
+                        context.irBuiltIns.unitType
+                    )
+                }
+            }
+
+            // Add the invoke method
+            val superMethod = functionSuperClass.functions.single {
+                it.owner.modality == Modality.ABSTRACT
+            }
+            irClass.addFunction {
+                name = superMethod.owner.name
+                returnType = callee.returnType
+                isSuspend = callee.isSuspend
+            }.apply {
+                overriddenSymbols += superMethod
+                dispatchReceiverParameter = parentAsClass.thisReceiver!!.copyTo(this)
+                annotations += callee.annotations
+                if (annotations.findAnnotation(ComposeFqNames.Composable) == null) {
+                    expression.type.annotations.findAnnotation(ComposeFqNames.Composable)?.let {
+                        annotations += it
+                    }
+                }
+                val valueParameterMap =
+                    callee.explicitParameters.withIndex().associate { (index, param) ->
+                        param to param.copyTo(this, index = index)
+                    }
+                valueParameters += valueParameterMap.values
+                body = DeclarationIrBuilder(context, symbol).irBlockBody(startOffset, endOffset) {
+                    callee.body?.statements?.forEach { statement ->
+                        +statement.transform(object : IrElementTransformerVoid() {
+                            override fun visitGetValue(expression: IrGetValue): IrExpression {
+                                val replacement = valueParameterMap[expression.symbol.owner]
+                                    ?: return super.visitGetValue(expression)
+
+                                at(expression.startOffset, expression.endOffset)
+                                return irGet(replacement)
+                            }
+
+                            override fun visitReturn(expression: IrReturn): IrExpression =
+                                if (expression.returnTargetSymbol != callee.symbol) {
+                                    super.visitReturn(expression)
+                                } else {
+                                    at(expression.startOffset, expression.endOffset)
+                                    irReturn(expression.value.transform(this, null))
+                                }
+
+                            override fun visitDeclaration(declaration: IrDeclaration): IrStatement {
+                                if (declaration.parent == callee)
+                                    declaration.parent = this@apply
+                                return super.visitDeclaration(declaration)
+                            }
+                        }, null)
+                    }
+                }
+            }
+        }
+
+        return irBlock {
+            +irClass
+            +irCall(constructor!!.symbol).apply {
+                if (valueArgumentsCount > 0) putValueArgument(0, boundReceiver!!.second)
+            }
+        }
+    }
+}

plugins/compose-compiler-hosted/src/main/java/androidx/compose/plugins/kotlin/compiler/lower/ComposableTypeRemapper.kt

@@ -0,0 +1,380 @@
+/*
+ * Copyright 2020 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.compose.plugins.kotlin.compiler.lower
+
+import androidx.compose.plugins.kotlin.ComposeFqNames
+import org.jetbrains.kotlin.backend.common.pop
+import org.jetbrains.kotlin.backend.jvm.JvmBackendContext
+import org.jetbrains.kotlin.builtins.isFunctionType
+import org.jetbrains.kotlin.descriptors.ClassDescriptor
+import org.jetbrains.kotlin.ir.IrElement
+import org.jetbrains.kotlin.ir.IrStatement
+import org.jetbrains.kotlin.ir.builders.declarations.addValueParameter
+import org.jetbrains.kotlin.ir.declarations.IrClass
+import org.jetbrains.kotlin.ir.declarations.IrConstructor
+import org.jetbrains.kotlin.ir.declarations.IrDeclarationOrigin
+import org.jetbrains.kotlin.ir.declarations.IrField
+import org.jetbrains.kotlin.ir.declarations.IrFile
+import org.jetbrains.kotlin.ir.declarations.IrFunction
+import org.jetbrains.kotlin.ir.declarations.IrMetadataSourceOwner
+import org.jetbrains.kotlin.ir.declarations.IrProperty
+import org.jetbrains.kotlin.ir.declarations.IrSimpleFunction
+import org.jetbrains.kotlin.ir.declarations.IrTypeParametersContainer
+import org.jetbrains.kotlin.ir.declarations.copyAttributes
+import org.jetbrains.kotlin.ir.declarations.impl.IrClassImpl
+import org.jetbrains.kotlin.ir.declarations.impl.IrFieldImpl
+import org.jetbrains.kotlin.ir.declarations.impl.IrFileImpl
+import org.jetbrains.kotlin.ir.declarations.impl.IrFunctionBase
+import org.jetbrains.kotlin.ir.declarations.impl.IrFunctionImpl
+import org.jetbrains.kotlin.ir.declarations.impl.IrPropertyImpl
+import org.jetbrains.kotlin.ir.expressions.IrCall
+import org.jetbrains.kotlin.ir.expressions.IrConstructorCall
+import org.jetbrains.kotlin.ir.expressions.IrMemberAccessExpression
+import org.jetbrains.kotlin.ir.expressions.IrStatementOrigin
+import org.jetbrains.kotlin.ir.expressions.impl.IrCallImpl
+import org.jetbrains.kotlin.ir.expressions.impl.IrConstructorCallImpl
+import org.jetbrains.kotlin.ir.symbols.IrSimpleFunctionSymbol
+import org.jetbrains.kotlin.ir.types.IrSimpleType
+import org.jetbrains.kotlin.ir.types.IrType
+import org.jetbrains.kotlin.ir.types.IrTypeAbbreviation
+import org.jetbrains.kotlin.ir.types.IrTypeArgument
+import org.jetbrains.kotlin.ir.types.IrTypeProjection
+import org.jetbrains.kotlin.ir.types.impl.IrSimpleTypeImpl
+import org.jetbrains.kotlin.ir.types.impl.IrTypeAbbreviationImpl
+import org.jetbrains.kotlin.ir.types.impl.makeTypeProjection
+import org.jetbrains.kotlin.ir.types.toKotlinType
+import org.jetbrains.kotlin.ir.types.withHasQuestionMark
+import org.jetbrains.kotlin.ir.util.DeepCopyIrTreeWithSymbols
+import org.jetbrains.kotlin.ir.util.DeepCopySymbolRemapper
+import org.jetbrains.kotlin.ir.util.SymbolRemapper
+import org.jetbrains.kotlin.ir.util.SymbolRenamer
+import org.jetbrains.kotlin.ir.util.TypeRemapper
+import org.jetbrains.kotlin.ir.util.TypeTranslator
+import org.jetbrains.kotlin.ir.util.hasAnnotation
+import org.jetbrains.kotlin.ir.util.isFunction
+import org.jetbrains.kotlin.ir.util.patchDeclarationParents
+import org.jetbrains.kotlin.ir.visitors.IrElementTransformerVoid
+import org.jetbrains.kotlin.psi2ir.PsiSourceManager
+import org.jetbrains.kotlin.psi2ir.findFirstFunction
+import org.jetbrains.kotlin.types.KotlinType
+import org.jetbrains.kotlin.types.TypeProjectionImpl
+import org.jetbrains.kotlin.types.replace
+
+class DeepCopyIrTreeWithSymbolsPreservingMetadata(
+    val context: JvmBackendContext,
+    val symbolRemapper: DeepCopySymbolRemapper,
+    val typeRemapper: TypeRemapper,
+    val typeTranslator: TypeTranslator,
+    symbolRenamer: SymbolRenamer = SymbolRenamer.DEFAULT
+) : DeepCopyIrTreeWithSymbols(symbolRemapper, typeRemapper, symbolRenamer) {
+
+    override fun visitClass(declaration: IrClass): IrClass {
+        return super.visitClass(declaration).also { it.copyMetadataFrom(declaration) }
+    }
+
+    override fun visitFunction(declaration: IrFunction): IrStatement {
+        return super.visitFunction(declaration).also { it.copyMetadataFrom(declaration) }
+    }
+
+    override fun visitSimpleFunction(declaration: IrSimpleFunction): IrSimpleFunction {
+        return super.visitSimpleFunction(declaration).also {
+            it.correspondingPropertySymbol = declaration.correspondingPropertySymbol
+            it.copyMetadataFrom(declaration)
+        }
+    }
+
+    override fun visitField(declaration: IrField): IrField {
+        return super.visitField(declaration).also {
+            (it as IrFieldImpl).metadata = declaration.metadata
+        }
+    }
+
+    override fun visitProperty(declaration: IrProperty): IrProperty {
+        return super.visitProperty(declaration).also { it.copyMetadataFrom(declaration) }
+    }
+
+    override fun visitFile(declaration: IrFile): IrFile {
+        val srcManager = context.psiSourceManager
+        val fileEntry = srcManager.getFileEntry(declaration) as? PsiSourceManager.PsiFileEntry
+        return super.visitFile(declaration).also {
+            if (fileEntry != null) {
+                srcManager.putFileEntry(it, fileEntry)
+            }
+            if (it is IrFileImpl) {
+                it.metadata = declaration.metadata
+            }
+        }
+    }
+
+    override fun visitConstructorCall(expression: IrConstructorCall): IrConstructorCall {
+        val ownerFn = expression.symbol.owner as? IrConstructor
+        // If we are calling an external constructor, we want to "remap" the types of its signature
+        // as well, since if it they are @Composable it will have its unmodified signature. These
+        // types won't be traversed by default by the DeepCopyIrTreeWithSymbols so we have to
+        // do it ourself here.
+        if (
+            ownerFn != null &&
+            ownerFn.origin == IrDeclarationOrigin.IR_EXTERNAL_DECLARATION_STUB
+        ) {
+            symbolRemapper.visitConstructor(ownerFn)
+            val newFn = super.visitConstructor(ownerFn).also {
+                it.parent = ownerFn.parent
+                it.patchDeclarationParents(it.parent)
+            }
+            val newCallee = symbolRemapper.getReferencedConstructor(newFn.symbol)
+
+            return IrConstructorCallImpl(
+                expression.startOffset, expression.endOffset,
+                expression.type.remapType(),
+                newCallee,
+                newCallee.descriptor,
+                expression.typeArgumentsCount,
+                expression.constructorTypeArgumentsCount,
+                expression.valueArgumentsCount,
+                mapStatementOrigin(expression.origin)
+            ).apply {
+                copyRemappedTypeArgumentsFrom(expression)
+                transformValueArguments(expression)
+            }.copyAttributes(expression)
+        }
+        return super.visitConstructorCall(expression)
+    }
+
+    override fun visitCall(expression: IrCall): IrCall {
+        val ownerFn = expression.symbol.owner as? IrSimpleFunction
+        val containingClass = expression.symbol.descriptor.containingDeclaration as? ClassDescriptor
+
+        // Any virtual calls on composable functions we want to make sure we update the call to
+        // the right function base class (of n+1 arity). The most often virtual call to make on
+        // a function instance is `invoke`, which we *already* do in the ComposeParamTransformer.
+        // There are others that can happen though as well, such as `equals` and `hashCode`. In this
+        // case, we want to update those calls as well.
+        if (
+            ownerFn != null &&
+            ownerFn.origin == IrDeclarationOrigin.FAKE_OVERRIDE &&
+            containingClass != null &&
+            containingClass.defaultType.isFunctionType &&
+            expression.dispatchReceiver?.type?.isComposable() == true
+        ) {
+            val typeArguments = containingClass.defaultType.arguments
+            val newFnClass = context.irIntrinsics.symbols.externalSymbolTable
+                .referenceClass(context.builtIns.getFunction(typeArguments.size))
+            val newDescriptor = newFnClass
+                .descriptor
+                .unsubstitutedMemberScope
+                .findFirstFunction(ownerFn.name.identifier) { true }
+
+            var newFn: IrSimpleFunction = IrFunctionImpl(
+                ownerFn.startOffset,
+                ownerFn.endOffset,
+                ownerFn.origin,
+                newDescriptor,
+                expression.type
+            )
+            symbolRemapper.visitSimpleFunction(newFn)
+            newFn = super.visitSimpleFunction(newFn).also { fn ->
+                fn.parent = newFnClass.owner
+                ownerFn.overriddenSymbols.mapTo(fn.overriddenSymbols) { it }
+                fn.dispatchReceiverParameter = ownerFn.dispatchReceiverParameter
+                fn.extensionReceiverParameter = ownerFn.extensionReceiverParameter
+                newDescriptor.valueParameters.forEach { p ->
+                    fn.addValueParameter(p.name.identifier, p.type.toIrType())
+                }
+                fn.patchDeclarationParents(fn.parent)
+                assert(fn.body == null) { "expected body to be null" }
+            }
+
+            val newCallee = symbolRemapper.getReferencedSimpleFunction(newFn.symbol)
+            return shallowCopyCall(expression, newCallee).apply {
+                copyRemappedTypeArgumentsFrom(expression)
+                transformValueArguments(expression)
+            }
+        }
+
+        // If we are calling an external function, we want to "remap" the types of its signature
+        // as well, since if it is @Composable it will have its unmodified signature. These
+        // functions won't be traversed by default by the DeepCopyIrTreeWithSymbols so we have to
+        // do it ourself here.
+        if (
+            ownerFn != null &&
+            ownerFn.origin == IrDeclarationOrigin.IR_EXTERNAL_DECLARATION_STUB
+        ) {
+            symbolRemapper.visitSimpleFunction(ownerFn)
+            val newFn = super.visitSimpleFunction(ownerFn).also {
+                it.parent = ownerFn.parent
+                it.correspondingPropertySymbol = ownerFn.correspondingPropertySymbol
+                it.patchDeclarationParents(it.parent)
+            }
+            val newCallee = symbolRemapper.getReferencedSimpleFunction(newFn.symbol)
+            return shallowCopyCall(expression, newCallee).apply {
+                copyRemappedTypeArgumentsFrom(expression)
+                transformValueArguments(expression)
+            }
+        }
+        return super.visitCall(expression)
+    }
+
+    /* copied verbatim from DeepCopyIrTreeWithSymbols, except with newCallee as a parameter */
+    private fun shallowCopyCall(expression: IrCall, newCallee: IrSimpleFunctionSymbol): IrCall {
+        return IrCallImpl(
+            expression.startOffset, expression.endOffset,
+            expression.type.remapType(),
+            newCallee,
+            newCallee.descriptor,
+            expression.typeArgumentsCount,
+            expression.valueArgumentsCount,
+            mapStatementOrigin(expression.origin),
+            symbolRemapper.getReferencedClassOrNull(expression.superQualifierSymbol)
+        ).apply {
+            copyRemappedTypeArgumentsFrom(expression)
+        }.copyAttributes(expression)
+    }
+
+    /* copied verbatim from DeepCopyIrTreeWithSymbols */
+    private fun IrMemberAccessExpression.copyRemappedTypeArgumentsFrom(
+        other: IrMemberAccessExpression
+    ) {
+        assert(typeArgumentsCount == other.typeArgumentsCount) {
+            "Mismatching type arguments: $typeArgumentsCount vs ${other.typeArgumentsCount} "
+        }
+        for (i in 0 until typeArgumentsCount) {
+            putTypeArgument(i, other.getTypeArgument(i)?.remapType())
+        }
+    }
+
+    /* copied verbatim from DeepCopyIrTreeWithSymbols */
+    private fun <T : IrMemberAccessExpression> T.transformValueArguments(original: T) {
+        transformReceiverArguments(original)
+        for (i in 0 until original.valueArgumentsCount) {
+            putValueArgument(i, original.getValueArgument(i)?.transform())
+        }
+    }
+
+    /* copied verbatim from DeepCopyIrTreeWithSymbols */
+    private fun <T : IrMemberAccessExpression> T.transformReceiverArguments(original: T): T =
+        apply {
+            dispatchReceiver = original.dispatchReceiver?.transform()
+            extensionReceiver = original.extensionReceiver?.transform()
+        }
+
+    /* copied verbatim from DeepCopyIrTreeWithSymbols */
+    private inline fun <reified T : IrElement> T.transform() =
+        transform(this@DeepCopyIrTreeWithSymbolsPreservingMetadata, null) as T
+
+    /* copied verbatim from DeepCopyIrTreeWithSymbols */
+    private fun IrType.remapType() = typeRemapper.remapType(this)
+
+    /* copied verbatim from DeepCopyIrTreeWithSymbols */
+    private fun mapStatementOrigin(origin: IrStatementOrigin?) = origin
+
+    private fun IrElement.copyMetadataFrom(owner: IrMetadataSourceOwner) {
+        when (this) {
+            is IrPropertyImpl -> metadata = owner.metadata
+            is IrFunctionBase -> metadata = owner.metadata
+            is IrClassImpl -> metadata = owner.metadata
+        }
+    }
+
+    private fun IrType.isComposable(): Boolean {
+        return annotations.hasAnnotation(ComposeFqNames.Composable)
+    }
+
+    private fun KotlinType.toIrType(): IrType = typeTranslator.translateType(this)
+}
+
+class ComposerTypeRemapper(
+    private val context: JvmBackendContext,
+    private val symbolRemapper: SymbolRemapper,
+    private val typeTranslator: TypeTranslator,
+    private val composerTypeDescriptor: ClassDescriptor
+) : TypeRemapper {
+
+    lateinit var deepCopy: IrElementTransformerVoid
+
+    private val scopeStack = mutableListOf<IrTypeParametersContainer>()
+
+    private val shouldTransform: Boolean get() {
+        // we don't want to remap the types of composable decoys. they are there specifically for
+        // their types to be unaltered!
+        return scopeStack.isEmpty() || scopeStack.last().origin != COMPOSABLE_DECOY_IMPL
+    }
+
+    override fun enterScope(irTypeParametersContainer: IrTypeParametersContainer) {
+        scopeStack.add(irTypeParametersContainer)
+    }
+
+    override fun leaveScope() {
+        scopeStack.pop()
+    }
+
+    private fun IrType.isComposable(): Boolean {
+        return annotations.hasAnnotation(ComposeFqNames.Composable)
+    }
+
+    private fun KotlinType.toIrType(): IrType = typeTranslator.translateType(this)
+
+    override fun remapType(type: IrType): IrType {
+        // TODO(lmr):
+        // This is basically creating the KotlinType and then converting to an IrType. Consider
+        // rewriting to just create the IrType directly, which would probably be more efficient.
+        if (type !is IrSimpleType) return type
+        if (!type.isFunction()) return underlyingRemapType(type)
+        if (!type.isComposable()) return underlyingRemapType(type)
+        if (!shouldTransform) return underlyingRemapType(type)
+        val oldArguments = type.toKotlinType().arguments
+        val newArguments =
+            oldArguments.subList(0, oldArguments.size - 1) +
+                    TypeProjectionImpl(composerTypeDescriptor.defaultType) +
+                    oldArguments.last()
+
+        val transformedComposableType = context
+            .irBuiltIns
+            .builtIns
+            .getFunction(oldArguments.size) // return type is an argument, so this is n + 1
+            .defaultType
+            .replace(newArguments)
+            .toIrType()
+            .withHasQuestionMark(type.hasQuestionMark) as IrSimpleType
+
+        return underlyingRemapType(transformedComposableType)
+    }
+
+    private fun underlyingRemapType(type: IrSimpleType): IrType {
+        return IrSimpleTypeImpl(
+            null,
+            symbolRemapper.getReferencedClassifier(type.classifier),
+            type.hasQuestionMark,
+            type.arguments.map { remapTypeArgument(it) },
+            type.annotations.map { it.transform(deepCopy, null) as IrConstructorCall },
+            type.abbreviation?.remapTypeAbbreviation()
+        )
+    }
+
+    private fun remapTypeArgument(typeArgument: IrTypeArgument): IrTypeArgument =
+        if (typeArgument is IrTypeProjection)
+            makeTypeProjection(this.remapType(typeArgument.type), typeArgument.variance)
+        else
+            typeArgument
+
+    private fun IrTypeAbbreviation.remapTypeAbbreviation() =
+        IrTypeAbbreviationImpl(
+            symbolRemapper.getReferencedTypeAlias(typeAlias),
+            hasQuestionMark,
+            arguments.map { remapTypeArgument(it) },
+            annotations
+        )
+}

plugins/compose-compiler-hosted/src/main/java/androidx/compose/plugins/kotlin/compiler/lower/ComposeObservePatcher.kt

@@ -0,0 +1,507 @@
+/*
+ * Copyright 2018 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.compose.plugins.kotlin.compiler.lower
+
+import androidx.compose.plugins.kotlin.ComposeFqNames
+import androidx.compose.plugins.kotlin.KtxNameConventions
+import androidx.compose.plugins.kotlin.KtxNameConventions.UPDATE_SCOPE
+import androidx.compose.plugins.kotlin.isEmitInline
+import org.jetbrains.kotlin.backend.common.FileLoweringPass
+import org.jetbrains.kotlin.backend.common.lower.DeclarationIrBuilder
+import org.jetbrains.kotlin.backend.common.lower.irIfThen
+import org.jetbrains.kotlin.backend.common.lower.irNot
+import org.jetbrains.kotlin.backend.jvm.JvmBackendContext
+import org.jetbrains.kotlin.descriptors.CallableMemberDescriptor
+import org.jetbrains.kotlin.descriptors.DeclarationDescriptor
+import org.jetbrains.kotlin.descriptors.FunctionDescriptor
+import org.jetbrains.kotlin.descriptors.Modality
+import org.jetbrains.kotlin.descriptors.SimpleFunctionDescriptor
+import org.jetbrains.kotlin.descriptors.SourceElement
+import org.jetbrains.kotlin.descriptors.Visibilities
+import org.jetbrains.kotlin.descriptors.annotations.Annotations
+import org.jetbrains.kotlin.descriptors.impl.AnonymousFunctionDescriptor
+import org.jetbrains.kotlin.descriptors.impl.ValueParameterDescriptorImpl
+import org.jetbrains.kotlin.incremental.components.NoLookupLocation
+import org.jetbrains.kotlin.ir.IrElement
+import org.jetbrains.kotlin.ir.IrStatement
+import org.jetbrains.kotlin.ir.UNDEFINED_OFFSET
+import org.jetbrains.kotlin.ir.builders.IrBlockBuilder
+import org.jetbrains.kotlin.ir.builders.declarations.addValueParameter
+import org.jetbrains.kotlin.ir.builders.irBlock
+import org.jetbrains.kotlin.ir.builders.irBlockBody
+import org.jetbrains.kotlin.ir.builders.irCall
+import org.jetbrains.kotlin.ir.builders.irEqeqeq
+import org.jetbrains.kotlin.ir.builders.irGet
+import org.jetbrains.kotlin.ir.builders.irNull
+import org.jetbrains.kotlin.ir.builders.irReturn
+import org.jetbrains.kotlin.ir.builders.irTemporary
+import org.jetbrains.kotlin.ir.declarations.IrDeclaration
+import org.jetbrains.kotlin.ir.declarations.IrDeclarationOrigin
+import org.jetbrains.kotlin.ir.declarations.IrFile
+import org.jetbrains.kotlin.ir.declarations.IrFunction
+import org.jetbrains.kotlin.ir.declarations.IrModuleFragment
+import org.jetbrains.kotlin.ir.declarations.impl.IrFunctionImpl
+import org.jetbrains.kotlin.ir.expressions.IrBlockBody
+import org.jetbrains.kotlin.ir.expressions.IrBreak
+import org.jetbrains.kotlin.ir.expressions.IrCall
+import org.jetbrains.kotlin.ir.expressions.IrContinue
+import org.jetbrains.kotlin.ir.expressions.IrExpression
+import org.jetbrains.kotlin.ir.expressions.IrReturn
+import org.jetbrains.kotlin.ir.expressions.IrStatementOrigin
+import org.jetbrains.kotlin.ir.expressions.impl.IrBlockImpl
+import org.jetbrains.kotlin.ir.expressions.impl.IrCallImpl
+import org.jetbrains.kotlin.ir.expressions.impl.IrConstImpl
+import org.jetbrains.kotlin.ir.expressions.impl.IrFunctionReferenceImpl
+import org.jetbrains.kotlin.ir.expressions.impl.IrGetValueImpl
+import org.jetbrains.kotlin.ir.expressions.impl.IrTryImpl
+import org.jetbrains.kotlin.ir.symbols.impl.IrSimpleFunctionSymbolImpl
+import org.jetbrains.kotlin.ir.types.IrType
+import org.jetbrains.kotlin.ir.types.classifierOrNull
+import org.jetbrains.kotlin.ir.types.getClass
+import org.jetbrains.kotlin.ir.types.isUnit
+import org.jetbrains.kotlin.ir.types.makeNullable
+import org.jetbrains.kotlin.ir.util.DeepCopySymbolRemapper
+import org.jetbrains.kotlin.ir.visitors.IrElementVisitor
+import org.jetbrains.kotlin.ir.visitors.transformChildrenVoid
+import org.jetbrains.kotlin.js.resolve.diagnostics.findPsi
+import org.jetbrains.kotlin.psi.KtFunctionLiteral
+import org.jetbrains.kotlin.psi2ir.findFirstFunction
+import org.jetbrains.kotlin.resolve.BindingTrace
+import org.jetbrains.kotlin.resolve.descriptorUtil.fqNameSafe
+import org.jetbrains.kotlin.resolve.inline.InlineUtil
+import org.jetbrains.kotlin.types.typeUtil.isUnit
+import org.jetbrains.kotlin.types.typeUtil.makeNullable
+import org.jetbrains.kotlin.util.OperatorNameConventions
+
+class ComposeObservePatcher(
+    context: JvmBackendContext,
+    symbolRemapper: DeepCopySymbolRemapper,
+    bindingTrace: BindingTrace
+) :
+    AbstractComposeLowering(context, symbolRemapper, bindingTrace),
+    FileLoweringPass,
+    ModuleLoweringPass {
+
+    override fun lower(module: IrModuleFragment) {
+        module.transformChildrenVoid(this)
+    }
+
+    override fun lower(irFile: IrFile) {
+        irFile.transformChildrenVoid(this)
+    }
+
+    override fun visitFunction(declaration: IrFunction): IrStatement {
+        super.visitFunction(declaration)
+        return visitFunctionForComposerParam(declaration)
+    }
+
+    private fun visitFunctionForComposerParam(declaration: IrFunction): IrStatement {
+        // Only insert observe scopes in non-empty composable function
+        if (declaration.body == null)
+            return declaration
+
+        val descriptor = declaration.descriptor
+
+        // Do not insert observe scope in an inline function
+        if (descriptor.isInline)
+            return declaration
+
+        // Do not insert an observe scope in an inline composable lambda
+        descriptor.findPsi()?.let { psi ->
+            (psi as? KtFunctionLiteral)?.let {
+                if (InlineUtil.isInlinedArgument(
+                        it,
+                        context.state.bindingContext,
+                        false
+                    )
+                )
+                    return declaration
+                if (it.isEmitInline(context.state.bindingContext)) {
+                    return declaration
+                }
+            }
+        }
+
+        // Do not insert an observe scope if the function has a return result
+        if (descriptor.returnType.let { it == null || !it.isUnit() })
+            return declaration
+
+        // Do not insert an observe scope if the function hasn't been transformed by the
+        // ComposerParamTransformer and has a synthetic "composer param" as its last parameter
+        val param = declaration.valueParameters.lastOrNull()
+        if (param == null || !param.isComposerParam()) return declaration
+
+        // Check if the descriptor has restart scope calls resolved
+        if (descriptor is SimpleFunctionDescriptor &&
+            // Lambdas that make are not lowered earlier should be ignored.
+            // All composable lambdas are already lowered to a class with an invoke() method.
+            declaration.origin != IrDeclarationOrigin.LOCAL_FUNCTION_FOR_LAMBDA &&
+            declaration.origin != IrDeclarationOrigin.LOCAL_FUNCTION_NO_CLOSURE) {
+
+            return functionWithRestartGroup(
+                declaration
+            ) { irGet(param) }
+        }
+
+        return declaration
+    }
+
+    private fun functionWithRestartGroup(
+        original: IrFunction,
+        getComposer: DeclarationIrBuilder.() -> IrExpression
+    ): IrStatement {
+        val oldBody = original.body
+
+        val startRestartGroupDescriptor = composerTypeDescriptor
+            .unsubstitutedMemberScope
+            .findFirstFunction(KtxNameConventions.STARTRESTARTGROUP.identifier) {
+                it.valueParameters.size == 1
+            }
+
+        val endRestartGroupDescriptor = composerTypeDescriptor
+            .unsubstitutedMemberScope
+            .findFirstFunction(KtxNameConventions.ENDRESTARTGROUP.identifier) {
+                it.valueParameters.size == 0
+            }
+
+        // Create call to get the composer
+        val unitType = context.irBuiltIns.unitType
+
+        val irBuilder = DeclarationIrBuilder(context, original.symbol)
+
+        // Create call to startRestartGroup
+        val startRestartGroup = irMethodCall(
+            irBuilder.getComposer(),
+            startRestartGroupDescriptor
+        ).apply {
+            putValueArgument(
+                0,
+                keyExpression(
+                    symbol.descriptor,
+                    original.startOffset,
+                    context.builtIns.intType.toIrType()
+                )
+            )
+        }
+
+        // Create call to endRestartGroup
+        val endRestartGroup = irMethodCall(irBuilder.getComposer(), endRestartGroupDescriptor)
+
+        val updateScopeDescriptor =
+            endRestartGroupDescriptor.returnType?.memberScope?.getContributedFunctions(
+                UPDATE_SCOPE,
+                NoLookupLocation.FROM_BACKEND
+            )?.singleOrNull { it.valueParameters.first().type.arguments.size == 2 }
+                ?: error("new updateScope not found in result type of endRestartGroup")
+        val updateScopeArgument:
+                    (outerBuilder: IrBlockBuilder) -> IrExpression =
+            if (original.isZeroParameterComposableUnitLambda()) { _ ->
+                // If we are in an invoke function for a callable class with no
+                // parameters then the `this` parameter can be used for the endRestartGroup.
+                // If isUnitInvoke() returns true then dispatchReceiverParameter is not
+                // null.
+                irBuilder.irGet(original.dispatchReceiverParameter!!)
+            } else { outerBuilder ->
+                // Create self-invoke lambda
+                val blockParameterDescriptor =
+                    updateScopeDescriptor.valueParameters.singleOrNull()
+                        ?: error("expected a single block parameter for updateScope")
+                val blockParameterType = blockParameterDescriptor.type
+                val selfSymbol = original.symbol
+
+                val lambdaDescriptor = AnonymousFunctionDescriptor(
+                    original.descriptor,
+                    Annotations.EMPTY,
+                    CallableMemberDescriptor.Kind.DECLARATION,
+                    SourceElement.NO_SOURCE,
+                    false
+                )
+
+                val passedInComposerParameter = ValueParameterDescriptorImpl(
+                    containingDeclaration = lambdaDescriptor,
+                    original = null,
+                    index = 0,
+                    annotations = Annotations.EMPTY,
+                    name = KtxNameConventions.COMPOSER_PARAMETER,
+                    outType = composerTypeDescriptor.defaultType.makeNullable(),
+                    declaresDefaultValue = false,
+                    isCrossinline = false,
+                    isNoinline = false,
+                    varargElementType = null,
+                    source = SourceElement.NO_SOURCE
+                )
+
+                lambdaDescriptor.apply {
+                    initialize(
+                        null,
+                        null,
+                        emptyList(),
+                        listOf(passedInComposerParameter),
+                        blockParameterType,
+                        Modality.FINAL,
+                        Visibilities.LOCAL
+                    )
+                }
+
+                val fn = IrFunctionImpl(
+                    UNDEFINED_OFFSET, UNDEFINED_OFFSET,
+                    IrDeclarationOrigin.LOCAL_FUNCTION_FOR_LAMBDA,
+                    IrSimpleFunctionSymbolImpl(lambdaDescriptor),
+                    context.irBuiltIns.unitType
+                ).also { fn ->
+                    fn.parent = original
+                    val localIrBuilder = DeclarationIrBuilder(context, fn.symbol)
+                    fn.addValueParameter(
+                        KtxNameConventions.COMPOSER_PARAMETER.identifier,
+                        composerTypeDescriptor.defaultType.toIrType().makeNullable()
+                    )
+                    fn.body = localIrBuilder.irBlockBody {
+                        // Call the function again with the same parameters
+                        +irReturn(irCall(selfSymbol).apply {
+                            symbol.descriptor
+                                .valueParameters
+                                .filter { !it.isComposerParam() }
+                                .forEachIndexed {
+                                    index, valueParameter ->
+                                val value = original.valueParameters[index].symbol
+                                putValueArgument(
+                                    index, IrGetValueImpl(
+                                        UNDEFINED_OFFSET,
+                                        UNDEFINED_OFFSET,
+                                        valueParameter.type.toIrType(),
+                                        value
+                                    )
+                                )
+                            }
+                            putValueArgument(
+                                descriptor.valueParameters.size - 1,
+                                IrGetValueImpl(
+                                    UNDEFINED_OFFSET,
+                                    UNDEFINED_OFFSET,
+                                    composerTypeDescriptor.defaultType.toIrType(),
+                                    fn.valueParameters[0].symbol
+                                )
+                            )
+                            symbol.descriptor.dispatchReceiverParameter?.let {
+                                // Ensure we get the correct type by trying to avoid
+                                // going through a KotlinType if possible.
+                                val parameter = original.dispatchReceiverParameter
+                                    ?: error("Expected dispatch receiver on declaration")
+                                val receiver = irGet(
+                                    parameter.type,
+                                    parameter.symbol
+                                )
+
+                                // Save the dispatch receiver into a temporary created in
+                                // the outer scope because direct references to the
+                                // receiver sometimes cause an invalid name, "$<this>", to
+                                // be generated.
+                                val tmp = outerBuilder.irTemporary(
+                                    value = receiver,
+                                    nameHint = "rcvr",
+                                    irType = parameter.type
+                                )
+                                dispatchReceiver = irGet(tmp)
+                            }
+                            symbol.descriptor.extensionReceiverParameter?.let {
+                                    receiverDescriptor ->
+                                extensionReceiver = irGet(
+                                    receiverDescriptor.type.toIrType(),
+                                    original.extensionReceiverParameter?.symbol
+                                        ?: error(
+                                            "Expected extension receiver on declaration"
+                                        )
+                                )
+                            }
+                            symbol.descriptor.typeParameters.forEachIndexed { index, descriptor ->
+                                putTypeArgument(index, descriptor.defaultType.toIrType())
+                            }
+                        })
+                    }
+                }
+                irBuilder.irBlock(origin = IrStatementOrigin.LAMBDA) {
+                    +fn
+                    +IrFunctionReferenceImpl(
+                        UNDEFINED_OFFSET,
+                        UNDEFINED_OFFSET,
+                        blockParameterType.toIrType(),
+                        fn.symbol,
+                        fn.symbol.descriptor,
+                        0,
+                        IrStatementOrigin.LAMBDA
+                    )
+                }
+            }
+
+        val endRestartGroupCallBlock = irBuilder.irBlock(
+            UNDEFINED_OFFSET,
+            UNDEFINED_OFFSET
+        ) {
+            val result = irTemporary(endRestartGroup)
+            val updateScopeSymbol = referenceSimpleFunction(updateScopeDescriptor)
+            +irIfThen(irNot(irEqeqeq(irGet(result.type, result.symbol), irNull())),
+                irCall(updateScopeSymbol).apply {
+                    dispatchReceiver = irGet(result.type, result.symbol)
+                    putValueArgument(
+                        0,
+                        updateScopeArgument(this@irBlock)
+                    )
+                }
+            )
+        }
+
+        when (oldBody) {
+            is IrBlockBody -> {
+                val earlyReturn = findPotentialEarly(oldBody)
+                if (earlyReturn != null) {
+                    if (earlyReturn is IrReturn &&
+                        oldBody.statements.lastOrNull() == earlyReturn) {
+                        // Transform block from:
+                        // {
+                        //   ...
+                        //   return value
+                        // }
+                        // to:
+                        // {
+                        //  composer.startRestartGroup()
+                        //  ...
+                        //  val tmp = value
+                        //  composer.endRestartGroup()
+                        //  return tmp
+                        // }
+                        original.body = irBuilder.irBlockBody {
+                            +startRestartGroup
+                            oldBody.statements
+                                .take(oldBody.statements.size - 1)
+                                .forEach { +it }
+                            val temp = irTemporary(earlyReturn.value)
+                            +endRestartGroupCallBlock
+                            +irReturn(irGet(temp))
+                        }
+                    } else {
+                        // Transform the block into
+                        // composer.startRestartGroup()
+                        // try {
+                        //   ... old statements ...
+                        // } finally {
+                        //    composer.endRestartGroup()
+                        // }
+                        original.body = irBuilder.irBlockBody {
+                            +IrTryImpl(
+                                oldBody.startOffset, oldBody.endOffset, unitType,
+                                IrBlockImpl(
+                                    UNDEFINED_OFFSET,
+                                    UNDEFINED_OFFSET,
+                                    unitType
+                                ).apply {
+                                    statements.add(startRestartGroup)
+                                    statements.addAll(oldBody.statements)
+                                },
+                                catches = emptyList(),
+                                finallyExpression = endRestartGroupCallBlock
+                            )
+                        }
+                    }
+                } else {
+                    // Insert the start and end calls into the block
+                    oldBody.statements.add(0, startRestartGroup)
+                    oldBody.statements.add(endRestartGroupCallBlock)
+                }
+                return original
+            }
+            else -> {
+                // Composable function do not use IrExpressionBody as they are converted
+                // by the call lowering to IrBlockBody to introduce the call temporaries.
+                error("Encountered IrExpressionBOdy when IrBlockBody was expected")
+            }
+        }
+    }
+
+    private fun irCall(descriptor: FunctionDescriptor): IrCall {
+        val type = descriptor.returnType?.toIrType() ?: error("Expected a return type")
+        val symbol = referenceFunction(descriptor)
+        return IrCallImpl(
+            UNDEFINED_OFFSET,
+            UNDEFINED_OFFSET,
+            type,
+            symbol,
+            descriptor
+        )
+    }
+
+    private fun irMethodCall(target: IrExpression, descriptor: FunctionDescriptor): IrCall {
+        return irCall(descriptor).apply {
+            dispatchReceiver = target
+        }
+    }
+
+    private fun keyExpression(
+        descriptor: CallableMemberDescriptor,
+        sourceOffset: Int,
+        intType: IrType
+    ): IrExpression {
+        val sourceKey = getKeyValue(descriptor, sourceOffset)
+        return IrConstImpl.int(
+            UNDEFINED_OFFSET,
+            UNDEFINED_OFFSET,
+            intType,
+            sourceKey
+        )
+    }
+
+    private fun IrFunction.isZeroParameterComposableUnitLambda(): Boolean {
+        if (name.isSpecial || name != OperatorNameConventions.INVOKE || !returnType.isUnit())
+            return false
+        val type = dispatchReceiverParameter?.type ?: return false
+        return valueParameters.size == 1 &&
+                    type.getClass()?.superTypes?.any {
+                        val fqName = it.classifierOrNull?.descriptor?.fqNameSafe
+                        fqName == ComposeFqNames.Function1
+                    } == true &&
+                    valueParameters[0].isComposerParam()
+    }
+}
+
+private fun findPotentialEarly(block: IrBlockBody): IrExpression? {
+    var result: IrExpression? = null
+    block.accept(object : IrElementVisitor<Unit, Unit> {
+        override fun visitElement(element: IrElement, data: Unit) {
+            if (result == null)
+                element.acceptChildren(this, Unit)
+        }
+
+        override fun visitBreak(jump: IrBreak, data: Unit) {
+            result = jump
+        }
+
+        override fun visitContinue(jump: IrContinue, data: Unit) {
+            result = jump
+        }
+
+        override fun visitReturn(expression: IrReturn, data: Unit) {
+            result = expression
+        }
+
+        override fun visitDeclaration(declaration: IrDeclaration, data: Unit) {
+            // Skip bodies of declarations
+        }
+    }, Unit)
+    return result
+}
+
+internal fun getKeyValue(descriptor: DeclarationDescriptor, startOffset: Int): Int =
+    descriptor.fqNameSafe.toString().hashCode() xor startOffset

plugins/compose-compiler-hosted/src/main/java/androidx/compose/plugins/kotlin/compiler/lower/ComposeResolutionMetadataTransformer.kt

@@ -0,0 +1,60 @@
+/*
+ * Copyright 2020 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.compose.plugins.kotlin.compiler.lower
+
+import androidx.compose.plugins.kotlin.ComposableCallableDescriptor
+import androidx.compose.plugins.kotlin.ComposableEmitDescriptor
+import androidx.compose.plugins.kotlin.analysis.ComposeWritableSlices
+import androidx.compose.plugins.kotlin.irTrace
+import org.jetbrains.kotlin.backend.common.BackendContext
+import org.jetbrains.kotlin.backend.common.FileLoweringPass
+import org.jetbrains.kotlin.backend.jvm.JvmBackendContext
+import org.jetbrains.kotlin.ir.declarations.IrFile
+import org.jetbrains.kotlin.ir.declarations.IrModuleFragment
+import org.jetbrains.kotlin.ir.expressions.IrCall
+import org.jetbrains.kotlin.ir.expressions.IrExpression
+import org.jetbrains.kotlin.ir.visitors.IrElementTransformerVoid
+import org.jetbrains.kotlin.ir.visitors.transformChildrenVoid
+
+class ComposeResolutionMetadataTransformer(val context: JvmBackendContext) :
+    IrElementTransformerVoid(),
+    FileLoweringPass,
+    ModuleLoweringPass {
+
+    override fun lower(module: IrModuleFragment) {
+        module.transformChildrenVoid(this)
+    }
+
+    override fun lower(irFile: IrFile) {
+        irFile.transformChildrenVoid(this)
+    }
+
+    override fun visitCall(expression: IrCall): IrExpression {
+
+        val descriptor = expression.symbol.descriptor
+
+        if (descriptor is ComposableEmitDescriptor) {
+            context.irTrace.record(
+                ComposeWritableSlices.COMPOSABLE_EMIT_METADATA,
+                expression,
+                descriptor
+            )
+        }
+
+        return super.visitCall(expression)
+    }
+}

plugins/compose-compiler-hosted/src/main/java/androidx/compose/plugins/kotlin/compiler/lower/ComposerIntrinsicTransformer.kt

@@ -0,0 +1,56 @@
+/*
+ * Copyright 2019 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.compose.plugins.kotlin.compiler.lower
+
+import androidx.compose.plugins.kotlin.ComposeFqNames
+import org.jetbrains.kotlin.backend.common.BackendContext
+import org.jetbrains.kotlin.backend.common.FileLoweringPass
+import org.jetbrains.kotlin.backend.jvm.JvmBackendContext
+import org.jetbrains.kotlin.ir.declarations.IrFile
+import org.jetbrains.kotlin.ir.declarations.IrModuleFragment
+import org.jetbrains.kotlin.ir.expressions.IrCall
+import org.jetbrains.kotlin.ir.expressions.IrExpression
+import org.jetbrains.kotlin.ir.visitors.IrElementTransformerVoid
+import org.jetbrains.kotlin.ir.visitors.transformChildrenVoid
+import org.jetbrains.kotlin.resolve.descriptorUtil.fqNameSafe
+
+class ComposerIntrinsicTransformer(val context: JvmBackendContext) :
+    IrElementTransformerVoid(),
+    FileLoweringPass,
+    ModuleLoweringPass {
+
+    override fun lower(module: IrModuleFragment) {
+        module.transformChildrenVoid(this)
+    }
+
+    override fun lower(irFile: IrFile) {
+        irFile.transformChildrenVoid(this)
+    }
+
+    override fun visitCall(expression: IrCall): IrExpression {
+        if (expression.symbol.descriptor.fqNameSafe == ComposeFqNames.CurrentComposerIntrinsic) {
+            // since this call was transformed by the ComposerParamTransformer, the first argument
+            // to this call is the composer itself. We just replace this expression with the
+            // argument expression and we are good.
+            assert(expression.valueArgumentsCount == 1)
+            val composerExpr = expression.getValueArgument(0)
+            if (composerExpr == null) error("Expected non-null composer argument")
+            return composerExpr
+        }
+        return super.visitCall(expression)
+    }
+}

plugins/compose-compiler-hosted/src/main/java/androidx/compose/plugins/kotlin/compiler/lower/ComposerLambdaMemoization.kt

@@ -0,0 +1,387 @@
+/*
+ * Copyright 2020 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.compose.plugins.kotlin.compiler.lower
+
+import androidx.compose.plugins.kotlin.ComposeUtils
+import androidx.compose.plugins.kotlin.analysis.ComposeWritableSlices
+import androidx.compose.plugins.kotlin.hasUntrackedAnnotation
+import androidx.compose.plugins.kotlin.irTrace
+import androidx.compose.plugins.kotlin.isEmitInline
+import org.jetbrains.kotlin.backend.common.lower.createIrBuilder
+import org.jetbrains.kotlin.backend.common.peek
+import org.jetbrains.kotlin.backend.common.pop
+import org.jetbrains.kotlin.backend.common.push
+import org.jetbrains.kotlin.backend.jvm.JvmBackendContext
+import org.jetbrains.kotlin.incremental.components.NoLookupLocation
+import org.jetbrains.kotlin.ir.IrStatement
+import org.jetbrains.kotlin.ir.UNDEFINED_OFFSET
+import org.jetbrains.kotlin.ir.builders.irBlock
+import org.jetbrains.kotlin.ir.builders.irCall
+import org.jetbrains.kotlin.ir.builders.irGet
+import org.jetbrains.kotlin.ir.builders.irReturn
+import org.jetbrains.kotlin.ir.builders.irTemporary
+import org.jetbrains.kotlin.ir.declarations.IrClass
+import org.jetbrains.kotlin.ir.declarations.IrFunction
+import org.jetbrains.kotlin.ir.declarations.IrModuleFragment
+import org.jetbrains.kotlin.ir.declarations.IrValueDeclaration
+import org.jetbrains.kotlin.ir.declarations.IrVariable
+import org.jetbrains.kotlin.ir.declarations.copyAttributes
+import org.jetbrains.kotlin.ir.expressions.IrExpression
+import org.jetbrains.kotlin.ir.expressions.IrFunctionAccessExpression
+import org.jetbrains.kotlin.ir.expressions.IrFunctionExpression
+import org.jetbrains.kotlin.ir.expressions.IrFunctionReference
+import org.jetbrains.kotlin.ir.expressions.IrValueAccessExpression
+import org.jetbrains.kotlin.ir.expressions.impl.IrFunctionReferenceImpl
+import org.jetbrains.kotlin.ir.expressions.impl.IrVarargImpl
+import org.jetbrains.kotlin.ir.types.toKotlinType
+import org.jetbrains.kotlin.ir.util.DeepCopySymbolRemapper
+import org.jetbrains.kotlin.ir.util.patchDeclarationParents
+import org.jetbrains.kotlin.ir.visitors.transformChildrenVoid
+import org.jetbrains.kotlin.js.resolve.diagnostics.findPsi
+import org.jetbrains.kotlin.name.FqName
+import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlin.psi.KtFunctionLiteral
+import org.jetbrains.kotlin.resolve.BindingTrace
+import org.jetbrains.kotlin.resolve.descriptorUtil.module
+import org.jetbrains.kotlin.resolve.inline.InlineUtil
+import org.jetbrains.kotlin.types.typeUtil.isUnit
+
+private class CaptureCollector {
+    val captures = mutableSetOf<IrValueDeclaration>()
+
+    fun recordCapture(local: IrValueDeclaration) {
+        captures.add(local)
+    }
+}
+
+private sealed class MemoizationContext {
+    open val composable get() = false
+    open fun declareLocal(local: IrValueDeclaration?) { }
+    open fun recordCapture(local: IrValueDeclaration?) { }
+    open fun pushCollector(collector: CaptureCollector) { }
+    open fun popCollector(collector: CaptureCollector) { }
+}
+
+private class ClassContext : MemoizationContext()
+
+private class FunctionContext(val declaration: IrFunction, override val composable: Boolean) :
+    MemoizationContext() {
+    val locals = mutableSetOf<IrValueDeclaration>()
+    var collectors = mutableListOf<CaptureCollector>()
+
+    init {
+        declaration.valueParameters.forEach {
+            declareLocal(it)
+        }
+    }
+
+    override fun declareLocal(local: IrValueDeclaration?) {
+        if (local != null) {
+            locals.add(local)
+        }
+    }
+
+    override fun recordCapture(local: IrValueDeclaration?) {
+        if (local != null && collectors.isNotEmpty() && locals.contains(local)) {
+            for (collector in collectors) {
+                collector.recordCapture(local)
+            }
+        }
+    }
+
+    override fun pushCollector(collector: CaptureCollector) {
+        collectors.add(collector)
+    }
+
+    override fun popCollector(collector: CaptureCollector) {
+        require(collectors.lastOrNull() == collector)
+        collectors.removeAt(collectors.size - 1)
+    }
+}
+
+class ComposerLambdaMemoization(
+    context: JvmBackendContext,
+    symbolRemapper: DeepCopySymbolRemapper,
+    bindingTrace: BindingTrace
+) :
+    AbstractComposeLowering(context, symbolRemapper, bindingTrace),
+    ModuleLoweringPass {
+
+    private val memoizationContextStack = mutableListOf<MemoizationContext>()
+
+    override fun lower(module: IrModuleFragment) = module.transformChildrenVoid(this)
+
+    override fun visitClass(declaration: IrClass): IrStatement {
+        memoizationContextStack.push(ClassContext())
+        val result = super.visitClass(declaration)
+        memoizationContextStack.pop()
+        return result
+    }
+
+    override fun visitFunction(declaration: IrFunction): IrStatement {
+        val descriptor = declaration.descriptor
+        val composable = descriptor.isComposable() &&
+                // Don't memoize in an inline function
+                !descriptor.isInline &&
+                // Don't memoize if in a composable that returns a value
+                // TODO(b/150390108): Consider allowing memoization in effects
+                descriptor.returnType.let { it != null && it.isUnit() }
+
+        memoizationContextStack.push(FunctionContext(declaration, composable))
+        val result = super.visitFunction(declaration)
+        memoizationContextStack.pop()
+        return result
+    }
+
+    override fun visitVariable(declaration: IrVariable): IrStatement {
+        memoizationContextStack.peek()?.declareLocal(declaration)
+        return super.visitVariable(declaration)
+    }
+
+    override fun visitValueAccess(expression: IrValueAccessExpression): IrExpression {
+        memoizationContextStack.forEach { it.recordCapture(expression.symbol.owner) }
+        return super.visitValueAccess(expression)
+    }
+
+    override fun visitFunctionReference(expression: IrFunctionReference): IrExpression {
+        // Memoize the instance created by using the :: prefix
+        val result = super.visitFunctionReference(expression)
+        val memoizationContext = memoizationContextStack.peek() as? FunctionContext
+            ?: return result
+        if (expression.valueArgumentsCount != 0) {
+            // If this syntax is as a curry syntax in the future, don't memoize.
+            // The syntax <expr>::<method>(<params>) and ::<function>(<params>) is reserved for
+            // future use. This ensures we don't try to memoize this syntax without knowing
+            // its meaning.
+
+            // The most likely correct implementation is to treat the parameters exactly as the
+            // receivers are treated below.
+            return result
+        }
+        if (memoizationContext.composable) {
+            // Memoize the reference for <expr>::<method>
+            val dispatchReceiver = expression.dispatchReceiver
+            val extensionReceiver = expression.extensionReceiver
+            if ((dispatchReceiver != null || extensionReceiver != null) &&
+                        dispatchReceiver.isNullOrStable() &&
+                        extensionReceiver.isNullOrStable()) {
+                // Save the receivers into a temporaries and memoize the function reference using
+                // the resulting temporaries
+                val builder = context.createIrBuilder(
+                    memoizationContext.declaration.symbol,
+                    startOffset = expression.startOffset,
+                    endOffset = expression.endOffset
+                )
+                return builder.irBlock(
+                    resultType = expression.type
+                ) {
+                    val captures = mutableListOf<IrValueDeclaration>()
+
+                    val tempDispatchReceiver = dispatchReceiver?.let {
+                        val tmp = irTemporary(it)
+                        captures.add(tmp)
+                        tmp
+                    }
+                    val tempExtensionReceiver = extensionReceiver?.let {
+                        val tmp = irTemporary(it)
+                        captures.add(tmp)
+                        tmp
+                    }
+
+                    +rememberExpression(
+                        memoizationContext,
+                        IrFunctionReferenceImpl(
+                            startOffset,
+                            endOffset,
+                            expression.type,
+                            expression.symbol,
+                            expression.descriptor,
+                            expression.typeArgumentsCount).copyAttributes(expression).apply {
+                            this.dispatchReceiver = tempDispatchReceiver?.let { irGet(it) }
+                            this.extensionReceiver = tempExtensionReceiver?.let { irGet(it) }
+                        },
+                        captures
+                    )
+                }
+            } else if (dispatchReceiver == null) {
+                return rememberExpression(memoizationContext, result, emptyList())
+            }
+        }
+        return result
+    }
+
+    override fun visitFunctionExpression(expression: IrFunctionExpression): IrExpression {
+        // Start recording variables captured in this scope
+        val composableFunction = memoizationContextStack.peek() as? FunctionContext
+            ?: return super.visitFunctionExpression(expression)
+        if (!composableFunction.composable) return super.visitFunctionExpression(expression)
+
+        val collector = CaptureCollector()
+        startCollector(collector)
+        // Wrap composable functions expressions or memoize non-composable function expressions
+        val result = super.visitFunctionExpression(expression)
+        stopCollector(collector)
+
+        // If the ancestor converted this then return
+        val functionExpression = result as? IrFunctionExpression ?: return result
+
+        // Ensure we don't transform targets of an inline function
+        if (functionExpression.isInlineArgument()) return functionExpression
+        if (functionExpression.isComposable()) {
+            return functionExpression
+        }
+
+        if (functionExpression.function.descriptor.hasUntrackedAnnotation())
+            return functionExpression
+
+        return rememberExpression(
+            composableFunction,
+            functionExpression,
+            collector.captures.toList()
+        )
+    }
+
+    private fun startCollector(collector: CaptureCollector) {
+        for (memoizationContext in memoizationContextStack) {
+            memoizationContext.pushCollector(collector)
+        }
+    }
+
+    private fun stopCollector(collector: CaptureCollector) {
+        for (memoizationContext in memoizationContextStack) {
+            memoizationContext.popCollector(collector)
+        }
+    }
+
+    private fun rememberExpression(
+        memoizationContext: FunctionContext,
+        expression: IrExpression,
+        captures: List<IrValueDeclaration>
+    ): IrExpression {
+
+        if (captures.any {
+            !((it as? IrVariable)?.isVar != true && it.type.toKotlinType().isStable())
+        }) return expression
+        val rememberParameterCount = captures.size + 1 // One additional parameter for the lambda
+        val declaration = memoizationContext.declaration
+        val descriptor = declaration.descriptor
+        val module = descriptor.module
+        val rememberFunctions = module
+            .getPackage(FqName(ComposeUtils.generateComposePackageName()))
+            .memberScope
+            .getContributedFunctions(
+                Name.identifier("remember"),
+                NoLookupLocation.FROM_BACKEND
+            )
+        val directRememberFunction = // Exclude the varargs version
+            rememberFunctions.singleOrNull {
+                it.valueParameters.size == rememberParameterCount &&
+                        // Exclude the varargs version
+                        it.valueParameters.firstOrNull()?.varargElementType == null
+            }
+        val rememberFunctionDescriptor = directRememberFunction
+            ?: rememberFunctions.single {
+                // Use the varargs version
+                it.valueParameters.firstOrNull()?.varargElementType != null
+            }
+
+        val rememberFunctionSymbol = referenceSimpleFunction(rememberFunctionDescriptor)
+
+        val irBuilder = context.createIrBuilder(
+            symbol = declaration.symbol,
+            startOffset = expression.startOffset,
+            endOffset = expression.endOffset
+        )
+
+        return irBuilder.irCall(
+            callee = rememberFunctionSymbol,
+            descriptor = rememberFunctionDescriptor,
+            type = expression.type
+        ).apply {
+            // The result type type parameter is first, followed by the argument types
+            putTypeArgument(0, expression.type)
+            val lambdaArgumentIndex = if (directRememberFunction != null) {
+                // Call to the non-vararg version
+                for (i in 1..captures.size) {
+                    putTypeArgument(i, captures[i - 1].type)
+                }
+
+                // condition arguments are the first `arg.size` arguments
+                for (i in captures.indices) {
+                    putValueArgument(i, irBuilder.irGet(captures[i]))
+                }
+                // The lambda is the last parameter
+                captures.size
+            } else {
+                val parameterType = rememberFunctionDescriptor.valueParameters[0].type.toIrType()
+                // Call to the vararg version
+                putValueArgument(0,
+                    IrVarargImpl(
+                        startOffset = UNDEFINED_OFFSET,
+                        endOffset = UNDEFINED_OFFSET,
+                        type = parameterType,
+                        varargElementType = context.irBuiltIns.anyType,
+                        elements = captures.map {
+                            irBuilder.irGet(it)
+                        }
+                    )
+                )
+                1
+            }
+
+            putValueArgument(lambdaArgumentIndex, irBuilder.irLambdaExpression(
+                descriptor = irBuilder.createFunctionDescriptor(
+                    rememberFunctionDescriptor.valueParameters.last().type
+                ),
+                type = rememberFunctionDescriptor.valueParameters.last().type.toIrType(),
+                body = {
+                    +irReturn(expression)
+                }
+            ))
+        }.patchDeclarationParents(declaration).markAsSynthetic()
+    }
+
+    private fun <T : IrFunctionAccessExpression> T.markAsSynthetic(): T {
+        // Mark it so the ComposableCallTransformer will insert the correct code around this call
+        context.state.irTrace.record(
+            ComposeWritableSlices.IS_SYNTHETIC_COMPOSABLE_CALL,
+            this,
+            true
+        )
+        return this
+    }
+
+    private fun IrFunctionExpression.isInlineArgument(): Boolean {
+        function.descriptor.findPsi()?.let { psi ->
+            (psi as? KtFunctionLiteral)?.let {
+                if (InlineUtil.isInlinedArgument(
+                        it,
+                        context.state.bindingContext,
+                        false
+                    )
+                )
+                    return true
+                if (it.isEmitInline(context.state.bindingContext)) {
+                    return true
+                }
+            }
+        }
+        return false
+    }
+
+    fun IrExpression?.isNullOrStable() = this == null || type.toKotlinType().isStable()
+}

plugins/compose-compiler-hosted/src/main/java/androidx/compose/plugins/kotlin/compiler/lower/ComposerParamTransformer.kt

@@ -0,0 +1,694 @@
+/*
+ * Copyright 2019 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.compose.plugins.kotlin.compiler.lower
+
+import androidx.compose.plugins.kotlin.KtxNameConventions
+import androidx.compose.plugins.kotlin.analysis.ComposeWritableSlices
+import androidx.compose.plugins.kotlin.hasComposableAnnotation
+import androidx.compose.plugins.kotlin.irTrace
+import androidx.compose.plugins.kotlin.isEmitInline
+import org.jetbrains.kotlin.backend.common.FileLoweringPass
+import org.jetbrains.kotlin.backend.common.descriptors.WrappedFunctionDescriptorWithContainerSource
+import org.jetbrains.kotlin.backend.common.descriptors.WrappedPropertyGetterDescriptor
+import org.jetbrains.kotlin.backend.common.descriptors.WrappedPropertySetterDescriptor
+import org.jetbrains.kotlin.backend.common.descriptors.WrappedSimpleFunctionDescriptor
+import org.jetbrains.kotlin.backend.common.ir.copyTo
+import org.jetbrains.kotlin.backend.common.ir.copyTypeParametersFrom
+import org.jetbrains.kotlin.backend.common.ir.isExpect
+import org.jetbrains.kotlin.backend.common.lower.DeclarationIrBuilder
+import org.jetbrains.kotlin.backend.common.lower.irThrow
+import org.jetbrains.kotlin.backend.jvm.JvmBackendContext
+import org.jetbrains.kotlin.backend.jvm.ir.isInlineParameter
+import org.jetbrains.kotlin.descriptors.FunctionDescriptor
+import org.jetbrains.kotlin.descriptors.Modality
+import org.jetbrains.kotlin.descriptors.PropertyGetterDescriptor
+import org.jetbrains.kotlin.descriptors.PropertySetterDescriptor
+import org.jetbrains.kotlin.ir.IrStatement
+import org.jetbrains.kotlin.ir.UNDEFINED_OFFSET
+import org.jetbrains.kotlin.ir.builders.declarations.addValueParameter
+import org.jetbrains.kotlin.ir.builders.irBlockBody
+import org.jetbrains.kotlin.ir.builders.irCall
+import org.jetbrains.kotlin.ir.builders.irString
+import org.jetbrains.kotlin.ir.declarations.IrClass
+import org.jetbrains.kotlin.ir.declarations.IrDeclaration
+import org.jetbrains.kotlin.ir.declarations.IrDeclarationContainer
+import org.jetbrains.kotlin.ir.declarations.IrDeclarationOrigin
+import org.jetbrains.kotlin.ir.declarations.IrDeclarationOriginImpl
+import org.jetbrains.kotlin.ir.declarations.IrFile
+import org.jetbrains.kotlin.ir.declarations.IrFunction
+import org.jetbrains.kotlin.ir.declarations.IrModuleFragment
+import org.jetbrains.kotlin.ir.declarations.IrOverridableDeclaration
+import org.jetbrains.kotlin.ir.declarations.IrProperty
+import org.jetbrains.kotlin.ir.declarations.IrSimpleFunction
+import org.jetbrains.kotlin.ir.declarations.IrValueParameter
+import org.jetbrains.kotlin.ir.declarations.copyAttributes
+import org.jetbrains.kotlin.ir.declarations.impl.IrFunctionImpl
+import org.jetbrains.kotlin.ir.expressions.IrCall
+import org.jetbrains.kotlin.ir.expressions.IrConstructorCall
+import org.jetbrains.kotlin.ir.expressions.IrExpression
+import org.jetbrains.kotlin.ir.expressions.IrGetValue
+import org.jetbrains.kotlin.ir.expressions.IrReturn
+import org.jetbrains.kotlin.ir.expressions.IrStatementOrigin
+import org.jetbrains.kotlin.ir.expressions.copyTypeArgumentsFrom
+import org.jetbrains.kotlin.ir.expressions.impl.IrCallImpl
+import org.jetbrains.kotlin.ir.expressions.impl.IrConstImpl
+import org.jetbrains.kotlin.ir.expressions.impl.IrConstructorCallImpl
+import org.jetbrains.kotlin.ir.expressions.impl.IrGetValueImpl
+import org.jetbrains.kotlin.ir.expressions.impl.IrReturnImpl
+import org.jetbrains.kotlin.ir.symbols.IrSimpleFunctionSymbol
+import org.jetbrains.kotlin.ir.symbols.impl.IrSimpleFunctionSymbolImpl
+import org.jetbrains.kotlin.ir.types.createType
+import org.jetbrains.kotlin.ir.types.isString
+import org.jetbrains.kotlin.ir.types.makeNullable
+import org.jetbrains.kotlin.ir.util.DeepCopySymbolRemapper
+import org.jetbrains.kotlin.ir.util.constructors
+import org.jetbrains.kotlin.ir.util.deepCopyWithSymbols
+import org.jetbrains.kotlin.ir.util.dump
+import org.jetbrains.kotlin.ir.util.explicitParameters
+import org.jetbrains.kotlin.ir.util.findAnnotation
+import org.jetbrains.kotlin.ir.util.functions
+import org.jetbrains.kotlin.ir.util.patchDeclarationParents
+import org.jetbrains.kotlin.ir.util.properties
+import org.jetbrains.kotlin.ir.visitors.IrElementTransformerVoid
+import org.jetbrains.kotlin.ir.visitors.acceptVoid
+import org.jetbrains.kotlin.ir.visitors.transformChildrenVoid
+import org.jetbrains.kotlin.js.resolve.diagnostics.findPsi
+import org.jetbrains.kotlin.load.java.JvmAbi
+import org.jetbrains.kotlin.name.FqName
+import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlin.psi.KtFunctionLiteral
+import org.jetbrains.kotlin.psi2ir.findFirstFunction
+import org.jetbrains.kotlin.resolve.BindingTrace
+import org.jetbrains.kotlin.resolve.DescriptorUtils
+import org.jetbrains.kotlin.resolve.inline.InlineUtil
+import org.jetbrains.kotlin.serialization.deserialization.descriptors.DescriptorWithContainerSource
+import org.jetbrains.kotlin.util.OperatorNameConventions
+
+private const val DEBUG_LOG = false
+
+class ComposerParamTransformer(
+    context: JvmBackendContext,
+    symbolRemapper: DeepCopySymbolRemapper,
+    bindingTrace: BindingTrace
+) :
+    AbstractComposeLowering(context, symbolRemapper, bindingTrace),
+    FileLoweringPass,
+    ModuleLoweringPass {
+
+    override fun lower(module: IrModuleFragment) {
+        module.transformChildrenVoid(this)
+
+        module.acceptVoid(symbolRemapper)
+
+        val typeRemapper = ComposerTypeRemapper(
+            context,
+            symbolRemapper,
+            typeTranslator,
+            composerTypeDescriptor
+        )
+        // for each declaration, we create a deepCopy transformer It is important here that we
+        // use the "preserving metadata" variant since we are using this copy to *replace* the
+        // originals, or else the module we would produce wouldn't have any metadata in it.
+        val transformer = DeepCopyIrTreeWithSymbolsPreservingMetadata(
+            context,
+            symbolRemapper,
+            typeRemapper,
+            typeTranslator
+        ).also { typeRemapper.deepCopy = it }
+        module.transformChildren(
+            transformer,
+            null
+        )
+        // just go through and patch all of the parents to make sure things are properly wired
+        // up.
+        module.patchDeclarationParents()
+    }
+
+    private val transformedFunctions: MutableMap<IrFunction, IrFunction> = mutableMapOf()
+
+    private val transformedFunctionSet = mutableSetOf<IrFunction>()
+
+    private val composerType = composerTypeDescriptor.defaultType.toIrType()
+
+    override fun lower(irFile: IrFile) {
+        if (DEBUG_LOG) {
+            println("""
+                =========
+                BEFORE
+                =========
+            """.trimIndent())
+            println(irFile.dump())
+        }
+        irFile.transform(this, null)
+        if (DEBUG_LOG) {
+            println("""
+                =========
+                AFTER TRANSFORM
+                =========
+            """.trimIndent())
+            println(irFile.dump())
+        }
+        irFile.remapComposableTypesWithComposerParam()
+        if (DEBUG_LOG) {
+            println("""
+                =========
+                AFTER TYPE REMAPPING
+                =========
+            """.trimIndent())
+            println(irFile.dump())
+        }
+    }
+
+    override fun visitFunction(declaration: IrFunction): IrStatement {
+        return super.visitFunction(declaration.withComposerParamIfNeeded())
+    }
+
+    override fun visitFile(declaration: IrFile): IrFile {
+        val originalFunctions = mutableListOf<IrFunction>()
+        val originalProperties = mutableListOf<Pair<IrProperty, IrSimpleFunction>>()
+        loop@for (child in declaration.declarations) {
+            when (child) {
+                is IrFunction -> originalFunctions.add(child)
+                is IrProperty -> {
+                    val getter = child.getter ?: continue@loop
+                    originalProperties.add(child to getter)
+                }
+            }
+        }
+        val result = super.visitFile(declaration)
+        result.patchWithSyntheticComposableDecoys(originalFunctions, originalProperties)
+        return result
+    }
+
+    override fun visitClass(declaration: IrClass): IrStatement {
+        val originalFunctions = declaration.functions.toList()
+        val originalProperties = declaration
+            .properties
+            .mapNotNull { p -> p.getter?.let { p to it } }
+            .toList()
+        val result = super.visitClass(declaration)
+        if (result !is IrClass) error("expected IrClass")
+        result.patchWithSyntheticComposableDecoys(originalFunctions, originalProperties)
+        return result
+    }
+
+    fun IrDeclarationContainer.patchWithSyntheticComposableDecoys(
+        originalFunctions: List<IrFunction>,
+        originalProperties: List<Pair<IrProperty, IrSimpleFunction>>
+    ) {
+        for (function in originalFunctions) {
+            if (transformedFunctions.containsKey(function) && function.isComposable()) {
+                declarations.add(function.copyAsComposableDecoy())
+            }
+        }
+        for ((property, getter) in originalProperties) {
+            if (transformedFunctions.containsKey(getter) && property.hasComposableAnnotation()) {
+                val newGetter = property.getter
+                assert(getter !== newGetter)
+                assert(newGetter != null)
+                // NOTE(lmr): the compiler seems to turn a getter with a single parameter into a
+                // setter, even though it's in the "getter" position. As a result, we will put
+                // the original parameter-less getter in the "getter" position, and add the
+                // single-parameter getter to the class itself.
+                property.getter = getter.copyAsComposableDecoy().also { it.parent = this }
+                declarations.add(newGetter!!)
+                newGetter.parent = this
+            }
+        }
+    }
+
+    fun IrCall.withComposerParamIfNeeded(composerParam: IrValueParameter): IrCall {
+        val isComposableLambda = isComposableLambdaInvoke()
+        if (!symbol.descriptor.isComposable() && !isComposableLambda)
+            return this
+        val ownerFn = when {
+            isComposableLambda ->
+                (symbol.owner as IrSimpleFunction).lambdaInvokeWithComposerParamIfNeeded()
+            else -> (symbol.owner as IrSimpleFunction).withComposerParamIfNeeded()
+        }
+        if (!transformedFunctionSet.contains(ownerFn))
+            return this
+        if (symbol.owner == ownerFn)
+            return this
+        return IrCallImpl(
+            startOffset,
+            endOffset,
+            type,
+            ownerFn.symbol,
+            ownerFn.symbol.descriptor,
+            typeArgumentsCount,
+            valueArgumentsCount + 1, // +1 for the composer param
+            origin,
+            superQualifierSymbol
+        ).also {
+            it.copyAttributes(this)
+            context.irTrace.record(
+                ComposeWritableSlices.IS_COMPOSABLE_CALL,
+                it,
+                true
+            )
+            it.copyTypeArgumentsFrom(this)
+            it.dispatchReceiver = dispatchReceiver
+            it.extensionReceiver = extensionReceiver
+            for (i in 0 until valueArgumentsCount) {
+                it.putValueArgument(i, getValueArgument(i))
+            }
+            it.putValueArgument(
+                valueArgumentsCount,
+                IrGetValueImpl(
+                    UNDEFINED_OFFSET,
+                    UNDEFINED_OFFSET,
+                    composerParam.symbol
+                )
+            )
+        }
+    }
+
+    // Transform `@Composable fun foo(params): RetType` into `fun foo(params, $composer: Composer): RetType`
+    fun IrFunction.withComposerParamIfNeeded(): IrFunction {
+        // don't transform functions that themselves were produced by this function. (ie, if we
+        // call this with a function that has the synthetic composer parameter, we don't want to
+        // transform it further).
+        if (transformedFunctionSet.contains(this)) return this
+
+        if (origin == COMPOSABLE_DECOY_IMPL) return this
+
+        // if not a composable fn, nothing we need to do
+        if (!descriptor.isComposable()) return this
+
+        // emit children lambdas are marked composable, but technically they are unit lambdas... so
+        // we don't want to transform them
+        if (isEmitInlineChildrenLambda()) return this
+
+        // if this function is an inlined lambda passed as an argument to an inline function (and
+        // is NOT a composable lambda), then we don't want to transform it. Ideally, this
+        // wouldn't have gotten this far because the `isComposable()` check above should return
+        // false, but right now the composable annotation checker seems to produce a
+        // false-positive here. It is important that we *DO NOT* transform this, but we should
+        // probably fix the annotation checker instead.
+        // TODO(b/147250515)
+        if (isNonComposableInlinedLambda()) return this
+
+        // we don't bother transforming expect functions. They exist only for type resolution and
+        // don't need to be transformed to have a composer parameter
+        if (isExpect) return this
+
+        // cache the transformed function with composer parameter
+        return transformedFunctions[this] ?: copyWithComposerParam()
+    }
+
+    fun IrFunction.lambdaInvokeWithComposerParamIfNeeded(): IrFunction {
+        if (transformedFunctionSet.contains(this)) return this
+        return transformedFunctions.getOrPut(this) {
+            lambdaInvokeWithComposerParam().also { transformedFunctionSet.add(it) }
+        }
+    }
+
+    fun IrFunction.lambdaInvokeWithComposerParam(): IrFunction {
+        val descriptor = descriptor
+        val argCount = descriptor.valueParameters.size
+        val newFnClass = context.irIntrinsics.symbols.externalSymbolTable
+            .referenceClass(context.builtIns.getFunction(argCount + 1))
+        val newDescriptor = newFnClass.descriptor.unsubstitutedMemberScope.findFirstFunction(
+            OperatorNameConventions.INVOKE.identifier
+        ) { true }
+
+        return IrFunctionImpl(
+            startOffset,
+            endOffset,
+            origin,
+            newDescriptor,
+            newDescriptor.returnType?.toIrType()!!
+        ).also { fn ->
+            fn.parent = newFnClass.owner
+
+            fn.copyTypeParametersFrom(this)
+            fn.dispatchReceiverParameter = dispatchReceiverParameter?.copyTo(fn)
+            fn.extensionReceiverParameter = extensionReceiverParameter?.copyTo(fn)
+            newDescriptor.valueParameters.forEach { p ->
+                fn.addValueParameter(p.name.identifier, p.type.toIrType())
+            }
+            assert(fn.body == null) { "expected body to be null" }
+        }
+    }
+
+    private fun IrFunction.copyAsComposableDecoy(): IrSimpleFunction {
+        if (origin == IrDeclarationOrigin.FAKE_OVERRIDE) return this as IrSimpleFunction
+        return copy().also { fn ->
+            fn.origin = COMPOSABLE_DECOY_IMPL
+            (fn as IrFunctionImpl).metadata = metadata
+            val errorClass = getTopLevelClass(FqName("kotlin.NotImplementedError"))
+            val errorCtor = errorClass.owner.constructors.single {
+                it.valueParameters.size == 1 &&
+                        it.valueParameters.single().type.isString()
+            }
+            if (this is IrOverridableDeclaration<*>) {
+                overriddenSymbols.mapTo(fn.overriddenSymbols) { it as IrSimpleFunctionSymbol }
+            }
+            // the decoy cannot have default expressions in its parameters, since they might be
+            // composable and if they are, it wouldn't have a composer param to use
+            fn.valueParameters.clear()
+            valueParameters.mapTo(fn.valueParameters) { p -> p.copyTo(fn, defaultValue = null) }
+            fn.body = DeclarationIrBuilder(context, fn.symbol).irBlockBody {
+                +irThrow(
+                    irCall(errorCtor).apply {
+                        putValueArgument(
+                            0,
+                            irString(
+                                "Composable functions cannot be called without a " +
+                                        "composer. If you are getting this error, it " +
+                                        "is likely because of a misconfigured compiler"
+                            )
+                        )
+                    }
+                )
+            }
+        }
+    }
+
+    private fun wrapDescriptor(descriptor: FunctionDescriptor): WrappedSimpleFunctionDescriptor {
+        return when (descriptor) {
+            is PropertyGetterDescriptor ->
+                WrappedPropertyGetterDescriptor(
+                    descriptor.annotations,
+                    descriptor.source
+                )
+            is PropertySetterDescriptor ->
+                WrappedPropertySetterDescriptor(
+                    descriptor.annotations,
+                    descriptor.source
+                )
+            is DescriptorWithContainerSource ->
+                WrappedFunctionDescriptorWithContainerSource(descriptor.containerSource)
+            else ->
+                WrappedSimpleFunctionDescriptor(sourceElement = descriptor.source)
+        }
+    }
+
+    private fun IrFunction.copy(
+        isInline: Boolean = this.isInline,
+        modality: Modality = descriptor.modality
+    ): IrSimpleFunction {
+        // TODO(lmr): use deepCopy instead?
+        val descriptor = descriptor
+        val newDescriptor = wrapDescriptor(descriptor)
+
+        return IrFunctionImpl(
+            startOffset,
+            endOffset,
+            origin,
+            IrSimpleFunctionSymbolImpl(newDescriptor),
+            name,
+            visibility,
+            modality,
+            returnType,
+            isInline,
+            isExternal,
+            descriptor.isTailrec,
+            descriptor.isSuspend
+        ).also { fn ->
+            newDescriptor.bind(fn)
+            if (this is IrSimpleFunction) {
+                fn.correspondingPropertySymbol = correspondingPropertySymbol
+            }
+            fn.parent = parent
+            fn.copyTypeParametersFrom(this)
+            fn.dispatchReceiverParameter = dispatchReceiverParameter?.copyTo(fn)
+            fn.extensionReceiverParameter = extensionReceiverParameter?.copyTo(fn)
+            valueParameters.mapTo(fn.valueParameters) { p ->
+                // Composable lambdas will always have `IrGet`s of all of their parameters
+                // generated, since they are passed into the restart lambda. This causes an
+                // interesting corner case with "anonymous parameters" of composable functions.
+                // If a parameter is anonymous (using the name `_`) in user code, you can usually
+                // make the assumption that it is never used, but this is technically not the
+                // case in composable lambdas. The synthetic name that kotlin generates for
+                // anonymous parameters has an issue where it is not safe to dex, so we sanitize
+                // the names here to ensure that dex is always safe.
+                p.copyTo(fn, name = dexSafeName(p.name))
+            }
+            annotations.mapTo(fn.annotations) { a -> a }
+            fn.body = body?.deepCopyWithSymbols(this)
+        }
+    }
+
+    private fun dexSafeName(name: Name): Name {
+        return if (name.isSpecial && name.asString().contains(' ')) {
+            val sanitized = name
+                .asString()
+                .replace(' ', '$')
+                .replace('<', '$')
+                .replace('>', '$')
+            Name.identifier(sanitized)
+        } else name
+    }
+
+    private fun jvmNameAnnotation(name: String): IrConstructorCall {
+        val jvmName = getTopLevelClass(DescriptorUtils.JVM_NAME)
+        val type = jvmName.createType(false, emptyList())
+        val ctor = jvmName.constructors.first()
+        return IrConstructorCallImpl(
+            UNDEFINED_OFFSET,
+            UNDEFINED_OFFSET,
+            type,
+            ctor,
+            ctor.descriptor,
+            0, 0, 1
+        ).also {
+            it.putValueArgument(0, IrConstImpl.string(
+                UNDEFINED_OFFSET,
+                UNDEFINED_OFFSET,
+                builtIns.stringType,
+                name
+            ))
+        }
+    }
+
+    private fun IrFunction.copyWithComposerParam(): IrFunction {
+        assert(explicitParameters.lastOrNull()?.name != KtxNameConventions.COMPOSER_PARAMETER) {
+            "Attempted to add composer param to $this, but it has already been added."
+        }
+        return copy().also { fn ->
+            val oldFn = this
+
+            // NOTE: it's important to add these here before we recurse into the body in
+            // order to avoid an infinite loop on circular/recursive calls
+            transformedFunctionSet.add(fn)
+            transformedFunctions[oldFn] = fn
+
+            // The overridden symbols might also be composable functions, so we want to make sure
+            // and transform them as well
+            if (this is IrOverridableDeclaration<*>) {
+                overriddenSymbols.mapTo(fn.overriddenSymbols) {
+                    it as IrSimpleFunctionSymbol
+                    val owner = it.owner
+                    val newOwner = owner.withComposerParamIfNeeded()
+                    newOwner.symbol as IrSimpleFunctionSymbol
+                }
+            }
+
+            // if we are transforming a composable property, the jvm signature of the
+            // corresponding getters and setters have a composer parameter. Since Kotlin uses the
+            // lack of a parameter to determine if it is a getter, this breaks inlining for
+            // composable property getters since it ends up looking for the wrong jvmSignature.
+            // In this case, we manually add the appropriate "@JvmName" annotation so that the
+            // inliner doesn't get confused.
+            val descriptor = descriptor
+            if (descriptor is PropertyGetterDescriptor &&
+                fn.annotations.findAnnotation(DescriptorUtils.JVM_NAME) == null
+            ) {
+                val name = JvmAbi.getterName(descriptor.correspondingProperty.name.identifier)
+                fn.annotations.add(jvmNameAnnotation(name))
+            }
+
+            // same thing for the setter
+            if (descriptor is PropertySetterDescriptor &&
+                fn.annotations.findAnnotation(DescriptorUtils.JVM_NAME) == null
+            ) {
+                val name = JvmAbi.setterName(descriptor.correspondingProperty.name.identifier)
+                fn.annotations.add(jvmNameAnnotation(name))
+            }
+
+            val valueParametersMapping = explicitParameters
+                .zip(fn.explicitParameters)
+                .toMap()
+
+            val composerParam = fn.addValueParameter(
+                KtxNameConventions.COMPOSER_PARAMETER.identifier,
+                composerType.makeNullable()
+            )
+            fn.transformChildrenVoid(object : IrElementTransformerVoid() {
+                var isNestedScope = false
+                override fun visitGetValue(expression: IrGetValue): IrGetValue {
+                    val newParam = valueParametersMapping[expression.symbol.owner]
+                    return if (newParam != null) {
+                        IrGetValueImpl(
+                            expression.startOffset,
+                            expression.endOffset,
+                            expression.type,
+                            newParam.symbol,
+                            expression.origin
+                        )
+                    } else expression
+                }
+
+                override fun visitReturn(expression: IrReturn): IrExpression {
+                    if (expression.returnTargetSymbol == oldFn.symbol) {
+                        // update the return statement to point to the new function, or else
+                        // it will be interpreted as a non-local return
+                        return super.visitReturn(IrReturnImpl(
+                            expression.startOffset,
+                            expression.endOffset,
+                            expression.type,
+                            fn.symbol,
+                            expression.value
+                        ))
+                    }
+                    return super.visitReturn(expression)
+                }
+
+                override fun visitFunction(declaration: IrFunction): IrStatement {
+                    val wasNested = isNestedScope
+                    try {
+                        // we don't want to pass the composer parameter in to composable calls
+                        // inside of nested scopes.... *unless* the scope was inlined.
+                        isNestedScope = if (declaration.isInlinedLambda()) wasNested else true
+                        return super.visitFunction(declaration)
+                    } finally {
+                        isNestedScope = wasNested
+                    }
+                }
+
+                override fun visitCall(expression: IrCall): IrExpression {
+                    val expr = if (!isNestedScope) {
+                        expression.withComposerParamIfNeeded(composerParam).also { call ->
+                            if (
+                                fn.isInline &&
+                                call !== expression &&
+                                call.isInlineParameterLambdaInvoke()
+                            ) {
+                                context.irTrace.record(
+                                    ComposeWritableSlices.IS_INLINE_COMPOSABLE_CALL,
+                                    call,
+                                    true
+                                )
+                            }
+                        }
+                    } else
+                        expression
+                    return super.visitCall(expr)
+                }
+            })
+        }
+    }
+
+    fun IrCall.isInlineParameterLambdaInvoke(): Boolean {
+        if (origin != IrStatementOrigin.INVOKE) return false
+        val lambda = dispatchReceiver as? IrGetValue
+        val valueParameter = lambda?.symbol?.owner as? IrValueParameter
+        return valueParameter?.isInlineParameter() == true
+    }
+
+    fun IrCall.isComposableLambdaInvoke(): Boolean {
+        return origin == IrStatementOrigin.INVOKE &&
+                dispatchReceiver?.type?.hasComposableAnnotation() == true
+    }
+
+    fun IrFunction.isNonComposableInlinedLambda(): Boolean {
+        descriptor.findPsi()?.let { psi ->
+            (psi as? KtFunctionLiteral)?.let {
+                val arg = InlineUtil.getInlineArgumentDescriptor(
+                    it,
+                    context.state.bindingContext
+                ) ?: return false
+
+                return !arg.type.hasComposableAnnotation()
+            }
+        }
+        return false
+    }
+
+    fun IrFunction.isInlinedLambda(): Boolean {
+        descriptor.findPsi()?.let { psi ->
+            (psi as? KtFunctionLiteral)?.let {
+                if (InlineUtil.isInlinedArgument(
+                        it,
+                        context.state.bindingContext,
+                        false
+                    )
+                )
+                    return true
+                if (it.isEmitInline(context.state.bindingContext)) {
+                    return true
+                }
+            }
+        }
+        return false
+    }
+
+    fun IrFunction.isEmitInlineChildrenLambda(): Boolean {
+        descriptor.findPsi()?.let { psi ->
+            (psi as? KtFunctionLiteral)?.let {
+                if (it.isEmitInline(context.state.bindingContext)) {
+                    return true
+                }
+            }
+        }
+        return false
+    }
+
+    fun IrFile.remapComposableTypesWithComposerParam() {
+        // NOTE(lmr): this operation is somewhat fragile, and the order things are done here is
+        // important.
+        val originalDeclarations = declarations.toList()
+
+        // The symbolRemapper needs to traverse everything to gather symbols, so we run this first.
+        acceptVoid(symbolRemapper)
+
+        // Now that we have all of the symbols, we can clear the existing declarations, since
+        // we are going to be putting new versions of them into the file.
+        declarations.clear()
+
+        originalDeclarations.mapTo(declarations) { d ->
+            val typeRemapper = ComposerTypeRemapper(
+                context,
+                symbolRemapper,
+                typeTranslator,
+                composerTypeDescriptor
+            )
+            // for each declaration, we create a deepCopy transformer It is important here that we
+            // use the "preserving metadata" variant since we are using this copy to *replace* the
+            // originals, or else the module we would produce wouldn't have any metadata in it.
+            val transformer = DeepCopyIrTreeWithSymbolsPreservingMetadata(
+                context,
+                symbolRemapper,
+                typeRemapper,
+                typeTranslator
+            ).also { typeRemapper.deepCopy = it }
+            val result = d.transform(
+                transformer,
+                null
+            ) as IrDeclaration
+            // just go through and patch all of the parents to make sure things are properly wired
+            // up.
+            result.patchDeclarationParents(this)
+            result
+        }
+    }
+}
+
+internal val COMPOSABLE_DECOY_IMPL =
+    object : IrDeclarationOriginImpl("COMPOSABLE_DECOY_IMPL", isSynthetic = true) {}

plugins/compose-compiler-hosted/src/main/java/androidx/compose/plugins/kotlin/compiler/lower/FakeJvmSymbols.kt

@@ -0,0 +1,85 @@
+/*
+ * Copyright 2020 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.compose.plugins.kotlin.compiler.lower
+
+import org.jetbrains.kotlin.backend.common.ir.createImplicitParameterDeclarationWithWrappedDescriptor
+import org.jetbrains.kotlin.descriptors.ClassKind
+import org.jetbrains.kotlin.descriptors.Modality
+import org.jetbrains.kotlin.descriptors.ModuleDescriptor
+import org.jetbrains.kotlin.descriptors.impl.EmptyPackageFragmentDescriptor
+import org.jetbrains.kotlin.ir.builders.declarations.addConstructor
+import org.jetbrains.kotlin.ir.builders.declarations.addFunction
+import org.jetbrains.kotlin.ir.builders.declarations.addTypeParameter
+import org.jetbrains.kotlin.ir.builders.declarations.addValueParameter
+import org.jetbrains.kotlin.ir.builders.declarations.buildClass
+import org.jetbrains.kotlin.ir.declarations.IrClass
+import org.jetbrains.kotlin.ir.declarations.IrPackageFragment
+import org.jetbrains.kotlin.ir.declarations.impl.IrExternalPackageFragmentImpl
+import org.jetbrains.kotlin.ir.descriptors.IrBuiltIns
+import org.jetbrains.kotlin.ir.symbols.IrClassSymbol
+import org.jetbrains.kotlin.ir.symbols.impl.IrExternalPackageFragmentSymbolImpl
+import org.jetbrains.kotlin.ir.types.defaultType
+import org.jetbrains.kotlin.name.FqName
+import org.jetbrains.kotlin.types.Variance
+
+class FakeJvmSymbols(val module: ModuleDescriptor, val irBuiltIns: IrBuiltIns) {
+
+    fun getJvmFunctionClass(parameterCount: Int): IrClassSymbol =
+        jvmFunctionClasses(parameterCount)
+
+    private val jvmFunctionClasses = { n: Int ->
+        createClass(FqName("kotlin.jvm.functions.Function$n"), ClassKind.INTERFACE) { klass ->
+            for (i in 1..n) {
+                klass.addTypeParameter("P$i", irBuiltIns.anyNType, Variance.IN_VARIANCE)
+            }
+            val returnType = klass.addTypeParameter("R", irBuiltIns.anyNType, Variance.OUT_VARIANCE)
+
+            klass.addFunction("invoke", returnType.defaultType, Modality.ABSTRACT).apply {
+                for (i in 1..n) {
+                    addValueParameter("p$i", klass.typeParameters[i - 1].defaultType)
+                }
+            }
+        }
+    }
+
+    private fun createClass(
+        fqName: FqName,
+        classKind: ClassKind = ClassKind.CLASS,
+        block: (IrClass) -> Unit = {}
+    ): IrClassSymbol =
+        buildClass {
+            name = fqName.shortName()
+            kind = classKind
+        }.apply {
+            parent = createPackage(FqName(fqName.parent().asString()))
+            createImplicitParameterDeclarationWithWrappedDescriptor()
+            block(this)
+        }.symbol
+
+    private fun createPackage(fqName: FqName): IrPackageFragment =
+        IrExternalPackageFragmentImpl(
+            IrExternalPackageFragmentSymbolImpl(
+                EmptyPackageFragmentDescriptor(module, fqName)
+            )
+        )
+
+    val lambdaClass: IrClassSymbol = createClass(FqName("kotlin.jvm.internal.Lambda")) { klass ->
+        klass.addConstructor().apply {
+            addValueParameter("arity", irBuiltIns.intType)
+        }
+    }
+}

plugins/compose-compiler-hosted/src/main/java/androidx/compose/plugins/kotlin/compiler/lower/ModuleLoweringPass.kt

@@ -0,0 +1,23 @@
+/*
+ * Copyright 2020 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.compose.plugins.kotlin.compiler.lower
+
+import org.jetbrains.kotlin.ir.declarations.IrModuleFragment
+
+interface ModuleLoweringPass {
+    fun lower(module: IrModuleFragment)
+}

plugins/compose-compiler-hosted/src/main/java/androidx/compose/plugins/kotlin/compiler/lower/package-info.java

@@ -0,0 +1,26 @@
+/*
+ * Copyright 2018 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.
+ */
+
+/**
+ * This package contains all the logic related to lowering a PSI containing Compose Components into IR.
+ * The entry-point for this package is ComponentClassLowering, which will generate all supporting
+ * synthetics.
+ * Each synthetic class of type [ClassName] lives in a file called [ClassName]Generator.
+ *
+ * Anything beginning with the token `lower` may modify IR
+ * Anything beginning with the token `generate` may only produce (return) IR
+ */
+package androidx.compose.plugins.kotlin.compiler.lower;

plugins/compose-compiler-hosted/src/main/java/androidx/compose/plugins/kotlin/frames/FrameUtils.kt

@@ -0,0 +1,23 @@
+package androidx.compose.plugins.kotlin.frames
+
+import org.jetbrains.kotlin.descriptors.ClassDescriptor
+import org.jetbrains.kotlin.descriptors.ModuleDescriptor
+import org.jetbrains.kotlin.descriptors.findClassAcrossModuleDependencies
+import org.jetbrains.kotlin.name.ClassId
+import org.jetbrains.kotlin.name.FqName
+import org.jetbrains.kotlin.name.Name
+import androidx.compose.plugins.kotlin.ComposeUtils
+import org.jetbrains.kotlin.resolve.descriptorUtil.fqNameSafe
+import org.jetbrains.kotlin.resolve.descriptorUtil.getSuperClassNotAny
+
+internal val composePackageName = FqName(ComposeUtils.generateComposePackageName())
+internal val framesPackageName = composePackageName.child(Name.identifier("frames"))
+internal val abstractRecordClassName = framesPackageName.child(Name.identifier("AbstractRecord"))
+internal val recordClassName = framesPackageName.child(Name.identifier("Record"))
+internal val componentClassName = composePackageName.child(Name.identifier("Component"))
+internal val modelClassName = composePackageName.child(Name.identifier("Model"))
+internal val framedTypeName = framesPackageName.child(Name.identifier("Framed"))
+internal fun ClassDescriptor.isFramed(): Boolean =
+    getSuperClassNotAny()?.fqNameSafe == componentClassName
+internal fun ModuleDescriptor.findTopLevel(name: FqName) =
+    findClassAcrossModuleDependencies(ClassId.topLevel(name)) ?: error("Could not find $name")

plugins/compose-compiler-hosted/src/main/java/androidx/compose/plugins/kotlin/frames/analysis/FrameMetadata.kt

@@ -0,0 +1,262 @@
+package androidx.compose.plugins.kotlin.frames.analysis
+
+import org.jetbrains.kotlin.builtins.DefaultBuiltIns
+import org.jetbrains.kotlin.descriptors.CallableMemberDescriptor
+import org.jetbrains.kotlin.descriptors.ClassDescriptor
+import org.jetbrains.kotlin.descriptors.Modality
+import org.jetbrains.kotlin.descriptors.PropertyDescriptor
+import org.jetbrains.kotlin.descriptors.ReceiverParameterDescriptor
+import org.jetbrains.kotlin.descriptors.SimpleFunctionDescriptor
+import org.jetbrains.kotlin.descriptors.SourceElement
+import org.jetbrains.kotlin.descriptors.TypeParameterDescriptor
+import org.jetbrains.kotlin.descriptors.Visibilities
+import org.jetbrains.kotlin.descriptors.Visibility
+import org.jetbrains.kotlin.descriptors.annotations.Annotations
+import org.jetbrains.kotlin.descriptors.impl.PropertyDescriptorImpl
+import org.jetbrains.kotlin.descriptors.impl.PropertyGetterDescriptorImpl
+import org.jetbrains.kotlin.descriptors.impl.PropertySetterDescriptorImpl
+import org.jetbrains.kotlin.descriptors.impl.SimpleFunctionDescriptorImpl
+import org.jetbrains.kotlin.descriptors.impl.ValueParameterDescriptorImpl
+import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlin.resolve.BindingContext
+import org.jetbrains.kotlin.resolve.descriptorUtil.builtIns
+import org.jetbrains.kotlin.resolve.hasBackingField
+import org.jetbrains.kotlin.types.KotlinType
+
+/**
+ * Helpers for creating the new properties, fields and properties of the framed class and the framed
+ * record
+ *
+ * - Framed properties are all public properties of the class.
+ * - A framed record has
+ *   - a corresponding property for each of the framed class' public properties
+ *   - a `create` method that creates a new instance of the record
+ *   - an `assign` method that copies values from the value into the record
+ * - All framed properties redirect to the current readable or writable record for the class
+ *   corresponding to the current open frame.
+ *
+ *   For example, given the declaration:
+ *
+ *   @Model
+ *   class MyModel {
+ *     var some: String = "Default some"
+ *     var data: String = "Default data"
+ *
+ *   }
+ *
+ *   The class is transformed into something like:
+ *
+ *   class MyModel: Framed {
+ *     var some: String
+ *       get() = (_readable(next) as MyModel_Record).some
+ *       set(value) { (_writable(next) as MyModel_Record).some = value }
+ *     var data: String
+ *       get() = ((_readable(next) as MyModel_Record).data
+ *       set(value) { (_writable(next, this) as MyModel_Record).data = value }
+ *
+ *     private var _firstFrameRecord: MyModelRecord? = null
+ *
+ *     override var firstFrameRecord: Record get() = _firstFrameRecord
+ *     override fun prependFrameRecord(value: Record) {
+ *       value.next = _firstFrameRecord
+ *       _firstFrameRecord = value
+ *     }
+ *
+ *     init {
+ *       next = MyModel_Record()
+ *       (next as MyModel_Record).some = "Default some"
+ *       (next as MyModel_Record).data = "Default data"
+ *     }
+ *   }
+ *
+ *   class MyModel_Record : AbstractRecord {
+ *     @JvmField var some: String
+ *     @JvmField var data: String
+ *
+ *     override fun create(): Record = MyModel_Record()
+ *     override fun assign(value: Record) {
+ *       some = (value as MyModel_Record).some
+ *       data = (value as MyModel_Record).data
+ *     }
+ *   }
+ */
+class FrameMetadata(private val framedClassDescriptor: ClassDescriptor) {
+    private val builtIns = DefaultBuiltIns.Instance
+
+    /**
+     * Get the list of properties on the framed class that should be framed
+     */
+    fun getFramedProperties(bindingContext: BindingContext) =
+        framedClassDescriptor.unsubstitutedMemberScope.getContributedDescriptors().mapNotNull {
+            if (it is PropertyDescriptor &&
+                it.kind == CallableMemberDescriptor.Kind.DECLARATION &&
+                it.isVar && it.hasBackingField(bindingContext)
+            ) it else null
+        }
+
+    /**
+     * Get the list of the record's properties (on for each public property of the framed object)
+     */
+    fun getRecordPropertyDescriptors(
+        recordClassDescriptor: ClassDescriptor,
+        bindingContext: BindingContext
+    ): List<PropertyDescriptor> {
+        return getFramedProperties(bindingContext).map {
+            syntheticProperty(
+                recordClassDescriptor,
+                it.name,
+                it.returnType!!
+            )
+        }
+    }
+
+    /**
+     * Get the list of the record's methods (create and assign)
+     */
+    fun getRecordMethodDescriptors(
+        container: ClassDescriptor,
+        recordDescriptor: ClassDescriptor
+    ): List<SimpleFunctionDescriptor> {
+        return listOf(
+            // fun create(): <record>
+            syntheticMethod(
+                container,
+                "create",
+                recordDescriptor.defaultType
+            ),
+            // fun assign(value: <record>)
+            syntheticMethod(
+                container, "assign", container.builtIns.unitType,
+                Parameter(
+                    "value",
+                    recordDescriptor.defaultType
+                )
+            )
+        )
+    }
+
+    fun firstFrameDescriptor(recordTypeDescriptor: ClassDescriptor): SimpleFunctionDescriptor =
+        syntheticMethod(
+            framedClassDescriptor,
+            "getFirstFrameRecord",
+            recordTypeDescriptor.defaultType
+        )
+
+    fun prependFrameRecordDescriptor(
+        recordTypeDescriptor: ClassDescriptor
+    ): SimpleFunctionDescriptor =
+        syntheticMethod(
+            framedClassDescriptor, "prependFrameRecord", builtIns.unitType,
+            Parameter(
+                "value",
+                recordTypeDescriptor.defaultType
+            )
+        )
+}
+
+private fun syntheticProperty(
+    container: ClassDescriptor,
+    name: Name,
+    type: KotlinType,
+    visibility: Visibility = Visibilities.PUBLIC,
+    readonly: Boolean = false
+): PropertyDescriptor =
+    PropertyDescriptorImpl.create(
+        container, Annotations.EMPTY, Modality.OPEN, Visibilities.PUBLIC, true,
+        name, CallableMemberDescriptor.Kind.SYNTHESIZED, SourceElement.NO_SOURCE,
+        false, false, true, true, false,
+        false
+    ).apply {
+        val getter = PropertyGetterDescriptorImpl(
+            this,
+            Annotations.EMPTY,
+            Modality.OPEN,
+            visibility,
+            false,
+            false,
+            false,
+            CallableMemberDescriptor.Kind.SYNTHESIZED,
+            null,
+            SourceElement.NO_SOURCE
+        ).apply { initialize(type) }
+        val setter = if (readonly) null else PropertySetterDescriptorImpl(
+            this,
+            Annotations.EMPTY,
+            Modality.OPEN,
+            visibility,
+            false,
+            false,
+            false,
+            CallableMemberDescriptor.Kind.SYNTHESIZED,
+            null,
+            SourceElement.NO_SOURCE
+        ).apply {
+            initialize(
+                PropertySetterDescriptorImpl.createSetterParameter(this, type, Annotations.EMPTY)
+            )
+        }
+        initialize(getter, setter)
+        setType(
+            type,
+            emptyList<TypeParameterDescriptor>(),
+            container.thisAsReceiverParameter,
+            null as ReceiverParameterDescriptor?
+        )
+    }
+
+private data class Parameter(val name: Name, val type: KotlinType) {
+    constructor(name: String, type: KotlinType) : this(Name.identifier(name), type)
+}
+
+private fun syntheticMethod(
+    container: ClassDescriptor,
+    name: Name,
+    returnType: KotlinType,
+    vararg parameters: Parameter
+): SimpleFunctionDescriptor =
+    SimpleFunctionDescriptorImpl.create(
+        container,
+        Annotations.EMPTY,
+        name,
+        CallableMemberDescriptor.Kind.SYNTHESIZED,
+        SourceElement.NO_SOURCE
+    ).apply {
+        val parameterDescriptors = parameters.map {
+            ValueParameterDescriptorImpl(
+                this,
+                null,
+                0,
+                Annotations.EMPTY,
+                it.name,
+                it.type,
+                false,
+                false,
+                false,
+                null,
+                SourceElement.NO_SOURCE
+            )
+        }
+
+        initialize(
+            null,
+            container.thisAsReceiverParameter,
+            emptyList<TypeParameterDescriptor>(),
+            parameterDescriptors,
+            returnType,
+            Modality.FINAL,
+            Visibilities.PUBLIC
+        )
+    }
+
+private fun syntheticMethod(
+    container: ClassDescriptor,
+    name: String,
+    returnType: KotlinType,
+    vararg parameters: Parameter
+): SimpleFunctionDescriptor =
+    syntheticMethod(
+        container,
+        Name.identifier(name),
+        returnType,
+        *parameters
+    )

plugins/compose-compiler-hosted/src/main/java/androidx/compose/plugins/kotlin/frames/analysis/FrameModelChecker.kt

@@ -0,0 +1,53 @@
+package androidx.compose.plugins.kotlin.frames.analysis
+
+import org.jetbrains.kotlin.container.StorageComponentContainer
+import org.jetbrains.kotlin.container.useInstance
+import org.jetbrains.kotlin.descriptors.ClassDescriptor
+import org.jetbrains.kotlin.descriptors.DeclarationDescriptor
+import org.jetbrains.kotlin.descriptors.ModuleDescriptor
+import org.jetbrains.kotlin.diagnostics.reportFromPlugin
+import org.jetbrains.kotlin.extensions.StorageComponentContainerContributor
+import org.jetbrains.kotlin.lexer.KtTokens
+import org.jetbrains.kotlin.psi.KtClass
+import org.jetbrains.kotlin.psi.KtDeclaration
+import androidx.compose.plugins.kotlin.ComposeUtils
+import androidx.compose.plugins.kotlin.analysis.ComposeDefaultErrorMessages
+import androidx.compose.plugins.kotlin.analysis.ComposeErrors
+import org.jetbrains.kotlin.platform.TargetPlatform
+import org.jetbrains.kotlin.platform.jvm.isJvm
+import org.jetbrains.kotlin.resolve.checkers.DeclarationChecker
+import org.jetbrains.kotlin.resolve.checkers.DeclarationCheckerContext
+
+open class FrameModelChecker : DeclarationChecker, StorageComponentContainerContributor {
+
+    override fun registerModuleComponents(
+        container: StorageComponentContainer,
+        platform: TargetPlatform,
+        moduleDescriptor: ModuleDescriptor
+    ) {
+        if (!platform.isJvm()) return
+        container.useInstance(FrameModelChecker())
+    }
+
+    override fun check(
+        declaration: KtDeclaration,
+        descriptor: DeclarationDescriptor,
+        context: DeclarationCheckerContext
+    ) {
+        if (descriptor is ClassDescriptor) {
+            if (!descriptor.isModelClass) return
+
+            if (declaration.hasModifier(KtTokens.OPEN_KEYWORD) ||
+                declaration.hasModifier(KtTokens.ABSTRACT_KEYWORD)) {
+                val element = (declaration as? KtClass)?.nameIdentifier ?: declaration
+                context.trace.reportFromPlugin(
+                    ComposeErrors.OPEN_MODEL.on(element),
+                    ComposeDefaultErrorMessages
+                )
+            }
+        }
+    }
+}
+
+private val MODEL_FQNAME = ComposeUtils.composeFqName("Model")
+val DeclarationDescriptor.isModelClass: Boolean get() = annotations.hasAnnotation(MODEL_FQNAME)

plugins/compose-compiler-hosted/src/main/java/androidx/compose/plugins/kotlin/frames/analysis/FramePackageAnalysisHandlerExtension.kt

@@ -0,0 +1,106 @@
+package androidx.compose.plugins.kotlin.frames.analysis
+
+import com.intellij.openapi.project.Project
+import org.jetbrains.kotlin.analyzer.AnalysisResult
+import org.jetbrains.kotlin.container.ComponentProvider
+import org.jetbrains.kotlin.container.get
+import org.jetbrains.kotlin.context.ProjectContext
+import org.jetbrains.kotlin.descriptors.ClassDescriptor
+import org.jetbrains.kotlin.descriptors.ModuleDescriptor
+import org.jetbrains.kotlin.descriptors.annotations.Annotated
+import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlin.psi.KtClassOrObject
+import org.jetbrains.kotlin.psi.KtFile
+import androidx.compose.plugins.kotlin.ComposeFlags
+import androidx.compose.plugins.kotlin.frames.FrameRecordClassDescriptor
+import androidx.compose.plugins.kotlin.frames.SyntheticFramePackageDescriptor
+import androidx.compose.plugins.kotlin.frames.abstractRecordClassName
+import androidx.compose.plugins.kotlin.frames.findTopLevel
+import androidx.compose.plugins.kotlin.frames.modelClassName
+import androidx.compose.plugins.kotlin.frames.recordClassName
+import org.jetbrains.kotlin.psi.KtDeclaration
+import org.jetbrains.kotlin.resolve.BindingTrace
+import org.jetbrains.kotlin.resolve.jvm.extensions.AnalysisHandlerExtension
+import org.jetbrains.kotlin.resolve.lazy.ResolveSession
+
+private fun doAnalysis(
+    module: ModuleDescriptor,
+    bindingTrace: BindingTrace,
+    files: Collection<KtFile>,
+    resolveSession: ResolveSession
+) {
+    if (!ComposeFlags.FRAMED_MODEL_CLASSES) return
+    for (file in files) {
+        analyseDeclarations(module, bindingTrace, file.declarations, resolveSession)
+    }
+}
+
+private fun analyseDeclarations(
+    module: ModuleDescriptor,
+    bindingTrace: BindingTrace,
+    declarations: List<KtDeclaration>,
+    resolveSession: ResolveSession
+) {
+    for (declaration in declarations) {
+        val ktClass = declaration as? KtClassOrObject ?: continue
+
+        analyseDeclarations(module, bindingTrace, ktClass.declarations, resolveSession)
+
+        val framedDescriptor = resolveSession.resolveToDescriptor(declaration) as?
+                ClassDescriptor ?: continue
+        if (!framedDescriptor.hasModelAnnotation()) continue
+
+        val classFqName = ktClass.fqName!!
+        val recordFqName = classFqName.parent().child(Name.identifier(
+            "${classFqName.shortName()}\$Record")
+        )
+        val recordSimpleName = recordFqName.shortName()
+        val recordPackage =
+            SyntheticFramePackageDescriptor(
+                module,
+                recordFqName.parent()
+            )
+        val baseTypeDescriptor = module.findTopLevel(abstractRecordClassName)
+        val recordDescriptor = module.findTopLevel(recordClassName)
+        val baseType = baseTypeDescriptor.defaultType
+        val frameClass =
+            FrameRecordClassDescriptor(
+                recordSimpleName,
+                recordPackage,
+                recordDescriptor,
+                framedDescriptor,
+                listOf(baseType),
+                bindingTrace.bindingContext
+            )
+
+        recordPackage.setClassDescriptor(frameClass)
+        bindingTrace.record(FrameWritableSlices.RECORD_CLASS, classFqName, frameClass)
+        bindingTrace.record(
+            FrameWritableSlices.FRAMED_DESCRIPTOR,
+            classFqName,
+            framedDescriptor
+        )
+    }
+}
+
+class FramePackageAnalysisHandlerExtension : AnalysisHandlerExtension {
+    override fun doAnalysis(
+        project: Project,
+        module: ModuleDescriptor,
+        projectContext: ProjectContext,
+        files: Collection<KtFile>,
+        bindingTrace: BindingTrace,
+        componentProvider: ComponentProvider
+    ): AnalysisResult? {
+        val resolveSession = componentProvider.get<ResolveSession>()
+        doAnalysis(
+            module,
+            bindingTrace,
+            files,
+            resolveSession
+        )
+        return null
+    }
+}
+
+private fun Annotated.hasModelAnnotation() = annotations.findAnnotation(modelClassName) != null

plugins/compose-compiler-hosted/src/main/java/androidx/compose/plugins/kotlin/frames/analysis/FrameWritableSlices.kt

@@ -0,0 +1,15 @@
+package androidx.compose.plugins.kotlin.frames.analysis
+
+import org.jetbrains.kotlin.descriptors.ClassDescriptor
+import org.jetbrains.kotlin.name.FqName
+import androidx.compose.plugins.kotlin.frames.FrameRecordClassDescriptor
+import org.jetbrains.kotlin.util.slicedMap.BasicWritableSlice
+import org.jetbrains.kotlin.util.slicedMap.RewritePolicy
+import org.jetbrains.kotlin.util.slicedMap.WritableSlice
+
+object FrameWritableSlices {
+    val RECORD_CLASS: WritableSlice<FqName, FrameRecordClassDescriptor> =
+        BasicWritableSlice(RewritePolicy.DO_NOTHING)
+    val FRAMED_DESCRIPTOR: WritableSlice<FqName, ClassDescriptor> =
+        BasicWritableSlice(RewritePolicy.DO_NOTHING)
+}

plugins/compose-compiler-hosted/src/main/resources/META-INF/services/org.jetbrains.kotlin.compiler.plugin.CommandLineProcessor

@@ -0,0 +1 @@
+androidx.compose.plugins.kotlin.ComposeCommandLineProcessor

plugins/compose-compiler-hosted/src/main/resources/META-INF/services/org.jetbrains.kotlin.compiler.plugin.ComponentRegistrar

@@ -0,0 +1 @@
+androidx.compose.plugins.kotlin.ComposeComponentRegistrar

settings.gradle

@@ -283,6 +283,8 @@ include ":kotlin-build-common",
         ':kotlin-noarg:plugin-marker',
         ":test-instrumenter",
 
+        ":compose-compiler-plugin",
+
         ":kotlinx-serialization-compiler-plugin",
         ":kotlinx-serialization-ide-plugin",
         ":kotlin-serialization",
@@ -498,5 +500,7 @@ project(':kotlin-serialization').projectDir = file("$rootDir/libraries/tools/kot
 project(':kotlin-serialization-unshaded').projectDir = file("$rootDir/libraries/tools/kotlin-serialization-unshaded")
 project(':kotlin-serialization:plugin-marker').projectDir = file("$rootDir/libraries/tools/kotlin-serialization/plugin-marker")
 
+project(':compose-compiler-plugin').projectDir = file("$rootDir/plugins/compose-compiler-hosted")
+
 // Uncomment to use locally built protobuf-relocated
 // includeBuild("dependencies/protobuf")