Fixing issues after code review

compiler/frontend/src/org/jetbrains/kotlin/psi/JetClass.java

@@ -253,7 +253,7 @@ public class JetClass extends JetTypeParameterListOwnerStub<KotlinClassStub> imp
     }
 
     @Nullable
-    public PsiElement getClassOrTraitKeyword() {
-        return findChildByType(TokenSet.create(JetTokens.CLASS_KEYWORD, JetTokens.TRAIT_KEYWORD));
+    public PsiElement getClassOrInterfaceKeyword() {
+        return findChildByType(TokenSet.create(JetTokens.CLASS_KEYWORD, JetTokens.INTERFACE_KEYWORD));
     }
 }

compiler/frontend/src/org/jetbrains/kotlin/psi/JetPsiFactory.kt

@@ -28,6 +28,7 @@ import com.intellij.util.LocalTimeCounter
 import org.jetbrains.kotlin.analyzer.ModuleInfo
 import org.jetbrains.kotlin.idea.JetFileType
 import org.jetbrains.kotlin.lexer.JetKeywordToken
+import org.jetbrains.kotlin.lexer.JetModifierKeywordToken
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.name.renderName
@@ -225,7 +226,7 @@ public class JetPsiFactory(private val project: Project) {
         return createClass("class Foo {\n $decl \n}").getSecondaryConstructors().first()
     }
 
-    public fun createModifierList(modifier: JetKeywordToken): JetModifierList {
+    public fun createModifierList(modifier: JetModifierKeywordToken): JetModifierList {
         return createModifierList(modifier.getValue())
     }
 
@@ -233,6 +234,10 @@ public class JetPsiFactory(private val project: Project) {
         return createProperty(text + " val x").getModifierList()!!
     }
 
+    public fun createModifier(modifier: JetModifierKeywordToken): PsiElement {
+        return createModifierList(modifier.getValue()).getModifier(modifier)!!
+    }
+
     public fun createAnnotationEntry(text: String): JetAnnotationEntry {
         val modifierList = createProperty(text + " val x").getModifierList()
         return modifierList!!.getAnnotationEntries().first()

compiler/frontend/src/org/jetbrains/kotlin/psi/addRemoveModifier.kt

@@ -44,7 +44,7 @@ internal fun addModifier(owner: JetModifierListOwner, modifier: JetModifierKeywo
 }
 
 internal fun addModifier(modifierList: JetModifierList, modifier: JetModifierKeywordToken, defaultVisibilityModifier: JetModifierKeywordToken) {
-    val newModifier = JetPsiFactory(modifierList).createModifierList(modifier).getModifier(modifier)!!
+    val newModifier = JetPsiFactory(modifierList).createModifier(modifier)
     val modifierToReplace = MODIFIERS_TO_REPLACE[modifier]
             ?.map { modifierList.getModifier(it) }
             ?.filterNotNull()