diff --git a/idea/src/org/jetbrains/jet/plugin/formatter/kotlinSpacingRules.kt b/idea/src/org/jetbrains/jet/plugin/formatter/kotlinSpacingRules.kt
index d62faf931bc..fe7bedcddb9 100644
--- a/idea/src/org/jetbrains/jet/plugin/formatter/kotlinSpacingRules.kt
+++ b/idea/src/org/jetbrains/jet/plugin/formatter/kotlinSpacingRules.kt
@@ -27,9 +27,19 @@ import com.intellij.formatting.Spacing
 import com.intellij.lang.ASTNode
 import com.intellij.formatting.ASTBlock
 import org.jetbrains.jet.plugin.formatter.KotlinSpacingBuilder.CustomSpacingBuilder
+import com.intellij.formatting.SpacingBuilder
+import com.intellij.formatting.SpacingBuilder.RuleBuilder
 
 val MODIFIERS_LIST_ENTRIES = TokenSet.orSet(TokenSet.create(ANNOTATION_ENTRY, ANNOTATION), MODIFIER_KEYWORDS)
 
+fun SpacingBuilder.beforeInside(element: IElementType, tokenSet: TokenSet, spacingFun: RuleBuilder.() -> Unit) {
+    tokenSet.getTypes().forEach { inType -> beforeInside(element, inType).spacingFun() }
+}
+
+fun SpacingBuilder.afterInside(element: IElementType, tokenSet: TokenSet, spacingFun: RuleBuilder.() -> Unit) {
+    tokenSet.getTypes().forEach { inType -> afterInside(element, inType).spacingFun() }
+}
+
 fun createSpacingBuilder(settings: CodeStyleSettings): KotlinSpacingBuilder {
     val jetSettings = settings.getCustomSettings(javaClass<JetCodeStyleSettings>())!!
     val jetCommonSettings = settings.getCommonSettings(JetLanguage.INSTANCE)!!
@@ -97,28 +107,13 @@ fun createSpacingBuilder(settings: CodeStyleSettings): KotlinSpacingBuilder {
             betweenInside(WHILE_KEYWORD, LPAR, WHILE).spacing(1, 1, 0, false, 0)
             betweenInside(WHILE_KEYWORD, LPAR, DO_WHILE).spacing(1, 1, 0, false, 0)
 
-            // TODO: Ask for better API
-            // Type of the declaration colon
-            beforeInside(COLON, PROPERTY).spaceIf(jetSettings.SPACE_BEFORE_TYPE_COLON)
-            afterInside(COLON, PROPERTY).spaceIf(jetSettings.SPACE_AFTER_TYPE_COLON)
-            beforeInside(COLON, FUN).spaceIf(jetSettings.SPACE_BEFORE_TYPE_COLON)
-            afterInside(COLON, FUN).spaceIf(jetSettings.SPACE_AFTER_TYPE_COLON)
-            beforeInside(COLON, VALUE_PARAMETER).spaceIf(jetSettings.SPACE_BEFORE_TYPE_COLON)
-            afterInside(COLON, VALUE_PARAMETER).spaceIf(jetSettings.SPACE_AFTER_TYPE_COLON)
-            beforeInside(COLON, MULTI_VARIABLE_DECLARATION_ENTRY).spaceIf(jetSettings.SPACE_BEFORE_TYPE_COLON)
-            afterInside(COLON, MULTI_VARIABLE_DECLARATION_ENTRY).spaceIf(jetSettings.SPACE_AFTER_TYPE_COLON)
-            beforeInside(COLON, FUNCTION_LITERAL).spaceIf(jetSettings.SPACE_BEFORE_TYPE_COLON)
-            afterInside(COLON, FUNCTION_LITERAL).spaceIf(jetSettings.SPACE_AFTER_TYPE_COLON)
+            val TYPE_COLON_ELEMENTS = TokenSet.create(PROPERTY, FUN, VALUE_PARAMETER, MULTI_VARIABLE_DECLARATION_ENTRY, FUNCTION_LITERAL)
+            beforeInside(COLON, TYPE_COLON_ELEMENTS) { spaceIf(jetSettings.SPACE_BEFORE_TYPE_COLON) }
+            afterInside(COLON, TYPE_COLON_ELEMENTS) { spaceIf(jetSettings.SPACE_AFTER_TYPE_COLON) }
 
-            // Extends or constraint colon
-            beforeInside(COLON, TYPE_CONSTRAINT).spaceIf(jetSettings.SPACE_BEFORE_EXTEND_COLON)
-            afterInside(COLON, TYPE_CONSTRAINT).spaceIf(jetSettings.SPACE_AFTER_EXTEND_COLON)
-            beforeInside(COLON, CLASS).spaceIf(jetSettings.SPACE_BEFORE_EXTEND_COLON)
-            afterInside(COLON, CLASS).spaceIf(jetSettings.SPACE_AFTER_EXTEND_COLON)
-            beforeInside(COLON, OBJECT_DECLARATION).spaceIf(jetSettings.SPACE_BEFORE_EXTEND_COLON)
-            afterInside(COLON, OBJECT_DECLARATION).spaceIf(jetSettings.SPACE_AFTER_EXTEND_COLON)
-            beforeInside(COLON, TYPE_PARAMETER).spaceIf(jetSettings.SPACE_BEFORE_EXTEND_COLON)
-            afterInside(COLON, TYPE_PARAMETER).spaceIf(jetSettings.SPACE_AFTER_EXTEND_COLON)
+            val EXTEND_COLON_ELEMENTS = TokenSet.create(TYPE_CONSTRAINT, CLASS, OBJECT_DECLARATION, TYPE_PARAMETER)
+            beforeInside(COLON, EXTEND_COLON_ELEMENTS) { spaceIf(jetSettings.SPACE_BEFORE_EXTEND_COLON) }
+            afterInside(COLON, EXTEND_COLON_ELEMENTS) { spaceIf(jetSettings.SPACE_AFTER_EXTEND_COLON) }
 
             between(VALUE_ARGUMENT_LIST, FUNCTION_LITERAL_EXPRESSION).spaces(1)
             beforeInside(ARROW, FUNCTION_LITERAL).spaceIf(jetSettings.SPACE_BEFORE_LAMBDA_ARROW)