From 0159d195399dbe34b32a92e6f4edb161b8a8112c Mon Sep 17 00:00:00 2001
From: Dmitry Jemerov <yole@jetbrains.com>
Date: Tue, 19 Dec 2017 15:22:55 +0100
Subject: [PATCH] Apply rainbow highlighting in anonymous initializers

 #KT-18839 Fixed
---
 .../idea/highlighter/KotlinRainbowVisitor.kt   |  4 +++-
 .../KotlinRainbowHighlighterTest.kt            | 18 ++++++++++++++++++
 2 files changed, 21 insertions(+), 1 deletion(-)

diff --git a/idea/src/org/jetbrains/kotlin/idea/highlighter/KotlinRainbowVisitor.kt b/idea/src/org/jetbrains/kotlin/idea/highlighter/KotlinRainbowVisitor.kt
index 353f689e8b0..2cb34d7df62 100644
--- a/idea/src/org/jetbrains/kotlin/idea/highlighter/KotlinRainbowVisitor.kt
+++ b/idea/src/org/jetbrains/kotlin/idea/highlighter/KotlinRainbowVisitor.kt
@@ -91,7 +91,9 @@ class KotlinRainbowVisitor : RainbowVisitor() {
             return
         }
 
-        val context = target.getStrictParentOfType<KtDeclarationWithBody>() ?: return
+        val context = target.getStrictParentOfType<KtDeclarationWithBody>()
+                ?: target.getStrictParentOfType<KtAnonymousInitializer>()
+                ?: return
         addInfo(getInfo(context, rainbowElement, rainbowElement.text, attributesKeyToUse))
     }
 
diff --git a/idea/tests/org/jetbrains/kotlin/idea/highlighter/KotlinRainbowHighlighterTest.kt b/idea/tests/org/jetbrains/kotlin/idea/highlighter/KotlinRainbowHighlighterTest.kt
index 8875983355b..57cf613f83c 100644
--- a/idea/tests/org/jetbrains/kotlin/idea/highlighter/KotlinRainbowHighlighterTest.kt
+++ b/idea/tests/org/jetbrains/kotlin/idea/highlighter/KotlinRainbowHighlighterTest.kt
@@ -83,6 +83,24 @@ class KotlinRainbowHighlighterTest : KotlinLightCodeInsightFixtureTestCase() {
         """)
     }
 
+    fun testInitBlock() {
+        checkRainbow("""
+            class Some {
+                init {
+                    val <rainbow color='ff000004'>x</rainbow> = 128
+                    println(<rainbow color='ff000004'>x</rainbow>)
+
+                    run {
+                        println(<rainbow color='ff000004'>x</rainbow>)
+                    }
+                    fun some() {
+                        val <rainbow color='ff000003'>b</rainbow> = 299
+                        println(<rainbow color='ff000003'>b</rainbow> + <rainbow color='ff000004'>x</rainbow>)
+                    }
+                }
+            }""")
+    }
+
     private fun checkRainbow(code: String) {
         myFixture.testRainbow("rainbow.kt", code, true, true)
     }