diff --git a/idea-plugin/src/main/kotlin/com/android/tools/compose/ComposableAnnotator.kt b/idea-plugin/src/main/kotlin/com/android/tools/compose/ComposableAnnotator.kt
index 7a8935778f..d39d71890a 100644
--- a/idea-plugin/src/main/kotlin/com/android/tools/compose/ComposableAnnotator.kt
+++ b/idea-plugin/src/main/kotlin/com/android/tools/compose/ComposableAnnotator.kt
@@ -16,78 +16,32 @@
 
 package com.android.tools.compose
 
-import com.android.tools.modules.*
-import com.intellij.lang.annotation.AnnotationHolder
-import com.intellij.lang.annotation.Annotator
-import com.intellij.lang.annotation.HighlightSeverity
 import com.intellij.openapi.editor.DefaultLanguageHighlighterColors
 import com.intellij.openapi.editor.colors.TextAttributesKey
-import com.intellij.openapi.progress.ProcessCanceledException
-import com.intellij.openapi.roots.ProjectFileIndex
-import com.intellij.openapi.util.Key
 import com.intellij.psi.PsiElement
-import org.jetbrains.kotlin.analyzer.AnalysisResult
-import org.jetbrains.kotlin.idea.caches.resolve.analyzeWithAllCompilerChecks
-import org.jetbrains.kotlin.psi.KtCallExpression
-import org.jetbrains.kotlin.psi.KtFile
-import org.jetbrains.kotlin.resolve.BindingContext
-import org.jetbrains.kotlin.resolve.calls.util.getResolvedCall
+import org.jetbrains.kotlin.descriptors.DeclarationDescriptor
+import org.jetbrains.kotlin.idea.highlighter.HighlighterExtension
+import org.jetbrains.kotlin.resolve.calls.model.ResolvedCall
 
 // Used to apply styles for calls to @Composable functions.
-class ComposableAnnotator : Annotator {
+class ComposableHighlighter : HighlighterExtension() {
+    override fun highlightDeclaration(elementToHighlight: PsiElement, descriptor: DeclarationDescriptor): TextAttributesKey? {
+        return null
+    }
+
+    override fun highlightCall(elementToHighlight: PsiElement, resolvedCall: ResolvedCall<*>): TextAttributesKey? {
+        return if (resolvedCall.isComposableInvocation()) COMPOSABLE_CALL_TEXT_ATTRIBUTES_KEY else null
+    }
+
     companion object TextAttributeRegistry {
         val COMPOSABLE_CALL_TEXT_ATTRIBUTES_KEY: TextAttributesKey
-        val COMPOSABLE_CALL_TEXT_ATTRIBUTES_NAME = "ComposableCallTextAttributes"
-        private val ANALYSIS_RESULT_KEY = Key<AnalysisResult>(
-            "ComposableAnnotator.DidAnnotateKey"
-        )
-        private val CAN_CONTAIN_COMPOSABLE_KEY = Key<Boolean>(
-            "ComposableAnnotator.CanContainComposable"
-        )
+        const val COMPOSABLE_CALL_TEXT_ATTRIBUTES_NAME = "ComposableCallTextAttributes"
 
         init {
             COMPOSABLE_CALL_TEXT_ATTRIBUTES_KEY = TextAttributesKey.createTextAttributesKey(
-              COMPOSABLE_CALL_TEXT_ATTRIBUTES_NAME,
-              DefaultLanguageHighlighterColors.FUNCTION_CALL)
+                COMPOSABLE_CALL_TEXT_ATTRIBUTES_NAME,
+                DefaultLanguageHighlighterColors.FUNCTION_CALL)
         }
     }
+}
 
-    override fun annotate(element: PsiElement, holder: AnnotationHolder) {
-        if (element !is KtCallExpression) return
-
-        // AnnotationHolder.currentAnnotationSession applies to a single file.
-        var canContainComposable = holder.currentAnnotationSession.getUserData(CAN_CONTAIN_COMPOSABLE_KEY)
-        if (canContainComposable == null) {
-            // isComposeEnabled doesn't work for library sources, we check all kt library sources files. File check only once on opening.
-            canContainComposable = element.inComposeModule() ||
-                    (element.containingFile.virtualFile != null &&
-                            ProjectFileIndex.getInstance(element.project)
-                                .isInLibrarySource(element.containingFile.virtualFile))
-            holder.currentAnnotationSession.putUserData(CAN_CONTAIN_COMPOSABLE_KEY, canContainComposable)
-        }
-
-        if (!canContainComposable) return
-
-        // AnnotationHolder.currentAnnotationSession applies to a single file.
-        var analysisResult = holder.currentAnnotationSession.getUserData(
-            ANALYSIS_RESULT_KEY
-        )
-        if (analysisResult == null) {
-            val ktFile = element.containingFile as? KtFile ?: return
-            analysisResult = ktFile.analyzeWithAllCompilerChecks()
-            holder.currentAnnotationSession.putUserData(
-              ANALYSIS_RESULT_KEY, analysisResult
-            )
-        }
-        if (analysisResult.isError()) {
-            throw ProcessCanceledException(analysisResult.error)
-        }
-        if (!shouldStyleCall(analysisResult.bindingContext, element)) return
-        val elementToStyle = element.calleeExpression ?: return
-        holder.newSilentAnnotation(HighlightSeverity.INFORMATION).range(elementToStyle).textAttributes(COMPOSABLE_CALL_TEXT_ATTRIBUTES_KEY).create()
-    }
-
-    private fun shouldStyleCall(bindingContext: BindingContext, element: KtCallExpression): Boolean {
-        return element.getResolvedCall(bindingContext)?.isComposableInvocation() == true
-    }
-}
\ No newline at end of file
diff --git a/idea-plugin/src/main/kotlin/com/android/tools/compose/ComposableDeclarationChecker.kt b/idea-plugin/src/main/kotlin/com/android/tools/compose/ComposableDeclarationChecker.kt
index 8e62519b25..95563d2480 100644
--- a/idea-plugin/src/main/kotlin/com/android/tools/compose/ComposableDeclarationChecker.kt
+++ b/idea-plugin/src/main/kotlin/com/android/tools/compose/ComposableDeclarationChecker.kt
@@ -39,6 +39,7 @@ import org.jetbrains.kotlin.psi.KtProperty
 import org.jetbrains.kotlin.psi.KtPropertyAccessor
 import org.jetbrains.kotlin.resolve.checkers.DeclarationChecker
 import org.jetbrains.kotlin.resolve.checkers.DeclarationCheckerContext
+import org.jetbrains.kotlin.resolve.descriptorUtil.module
 import org.jetbrains.kotlin.types.KotlinType
 
 class ComposableDeclarationChecker : DeclarationChecker, StorageComponentContainerContributor {
@@ -47,7 +48,6 @@ class ComposableDeclarationChecker : DeclarationChecker, StorageComponentContain
     platform: TargetPlatform,
     moduleDescriptor: ModuleDescriptor
   ) {
-    if (!platform.isJvm()) return
     container.useInstance(this)
   }
 
diff --git a/idea-plugin/src/main/kotlin/com/android/tools/compose/ComposeColorSettingsPage.kt b/idea-plugin/src/main/kotlin/com/android/tools/compose/ComposeColorSettingsPage.kt
index 6614e9347a..e121e46040 100644
--- a/idea-plugin/src/main/kotlin/com/android/tools/compose/ComposeColorSettingsPage.kt
+++ b/idea-plugin/src/main/kotlin/com/android/tools/compose/ComposeColorSettingsPage.kt
@@ -34,8 +34,8 @@ class ComposeColorSettingsPage : ColorSettingsPage {
     override fun getAdditionalHighlightingTagToDescriptorMap(): MutableMap<String,
             TextAttributesKey> {
         val attributes = HashMap<String, TextAttributesKey>()
-        attributes[ComposableAnnotator.COMPOSABLE_CALL_TEXT_ATTRIBUTES_NAME] =
-          ComposableAnnotator.COMPOSABLE_CALL_TEXT_ATTRIBUTES_KEY
+        attributes[ComposableHighlighter.COMPOSABLE_CALL_TEXT_ATTRIBUTES_NAME] =
+            ComposableHighlighter.COMPOSABLE_CALL_TEXT_ATTRIBUTES_KEY
         attributes["ANNOTATION"] = KotlinHighlightingColors.ANNOTATION
         attributes["KEYWORD"] = KotlinHighlightingColors.KEYWORD
         attributes["FUNCTION_DECLARATION"] = KotlinHighlightingColors.FUNCTION_DECLARATION
@@ -50,7 +50,7 @@ class ComposeColorSettingsPage : ColorSettingsPage {
     override fun getAttributeDescriptors(): Array<AttributesDescriptor> {
         // TODO: this needs to be localized.
         return arrayOf(AttributesDescriptor("Calls to @Compose functions",
-                                            ComposableAnnotator.COMPOSABLE_CALL_TEXT_ATTRIBUTES_KEY))
+            ComposableHighlighter.COMPOSABLE_CALL_TEXT_ATTRIBUTES_KEY))
     }
 
     override fun getColorDescriptors(): Array<ColorDescriptor> {
diff --git a/idea-plugin/src/main/resources/META-INF/plugin.xml b/idea-plugin/src/main/resources/META-INF/plugin.xml
index 2906ed56dd..ba5f25125e 100644
--- a/idea-plugin/src/main/resources/META-INF/plugin.xml
+++ b/idea-plugin/src/main/resources/META-INF/plugin.xml
@@ -26,6 +26,7 @@
         <quickFixContributor implementation="com.android.tools.compose.intentions.ComposeUnresolvedFunctionFixContributor"/>
         <additionalExtractableAnalyser implementation="com.android.tools.compose.ComposableFunctionExtractableAnalyser"/>
         <irGenerationExtension implementation="com.android.tools.compose.ComposePluginIrGenerationExtension"/>
+        <highlighterExtension implementation="com.android.tools.compose.ComposableHighlighter"/>
     </extensions>
 
     <extensions defaultExtensionNs="org.jetbrains.kotlin.extensions.internal">
@@ -34,10 +35,7 @@
 
     <extensions defaultExtensionNs="com.intellij">
         <dependencySupport coordinate="androidx.compose.runtime:runtime" kind="java" displayName="Jetpack Compose"/>
-
-        <annotator
-                language="kotlin"
-                implementationClass="com.android.tools.compose.ComposableAnnotator"/>
+        
         <additionalTextAttributes scheme="Default" file="colorschemes/IdeComposableAnnotatorColorSchemeDefault.xml"/>
         <colorSettingsPage implementation="com.android.tools.compose.ComposeColorSettingsPage"/>