JetCodeFragment: increment modification count after inserting import

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

@@ -24,8 +24,9 @@ public class JetBlockCodeFragment(
         project: Project,
         name: String,
         text: CharSequence,
+        imports: String?,
         context: PsiElement?
-) : JetCodeFragment(project, name, text, JetNodeTypes.BLOCK_CODE_FRAGMENT, context) {
+) : JetCodeFragment(project, name, text, imports, JetNodeTypes.BLOCK_CODE_FRAGMENT, context) {
 
     override fun getContentElement() = findChildByClass(javaClass<JetBlockExpression>())
                                                 ?: throw IllegalStateException("Block expression should be parsed for BlockCodeFragment")

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

@@ -26,24 +26,28 @@ import com.intellij.testFramework.LightVirtualFile
 import org.jetbrains.kotlin.idea.JetFileType
 import java.util.HashSet
 import com.intellij.openapi.util.Key
+import com.intellij.psi.impl.PsiModificationTrackerImpl
+import com.intellij.psi.util.PsiModificationTracker
 import org.jetbrains.kotlin.types.JetType
+import java.util.LinkedHashSet
 
 public abstract class JetCodeFragment(
         private val _project: Project,
         name: String,
         text: CharSequence,
+        imports: String?, // Should be separated by JetCodeFragment.IMPORT_SEPARATOR
         elementType: IElementType,
         private val context: PsiElement?
 ): JetFile((PsiManager.getInstance(_project) as PsiManagerEx).getFileManager().createFileViewProvider(LightVirtualFile(name, JetFileType.INSTANCE, text), true), false), JavaCodeFragment {
 
     private var viewProvider = super<JetFile>.getViewProvider() as SingleRootFileViewProvider
-    private var myImports = HashSet<String>();
+    private var myImports = LinkedHashSet<String>();
 
     {
         getViewProvider().forceCachedPsi(this)
         init(TokenType.CODE_FRAGMENT, elementType)
         if (context != null) {
-            addImportsFromString(getImportsForElement(context))
+            initImports(context, imports)
         }
     }
 
@@ -98,6 +102,11 @@ public abstract class JetCodeFragment(
     override fun addImportsFromString(imports: String?) {
         if (imports == null || imports.isEmpty()) return
 
+        // We should increment modification tracker after inserting import in code fragment to invalidate resolve caches.
+        // Without this modification references with new import won't be resolved without any modification in code fragment.
+        // Also shorten references won't work.
+        (PsiModificationTracker.SERVICE.getInstance(getProject()) as PsiModificationTrackerImpl).incOutOfCodeBlockModificationCounter()
+
         myImports.addAll(imports.split(IMPORT_SEPARATOR))
     }
 
@@ -119,17 +128,27 @@ public abstract class JetCodeFragment(
         return true
     }
 
+    private fun initImports(context: PsiElement, imports: String?) {
+        val containingFile = context.getContainingFile()
+        if (containingFile !is JetFile) return
+
+        val importListForContextElement = containingFile.getImportList()
+        if (importListForContextElement != null) {
+            myImports.addAll(importListForContextElement.getImports().map { it.getText() })
+        }
+
+        val packageName = containingFile.getPackageDirective()?.getFqName()?.asString()
+        if (packageName != null && packageName.isNotEmpty()) {
+            myImports.add("import $packageName.*")
+        }
+
+        if (imports != null && !imports.isEmpty()) {
+            myImports.addAll(imports.split(IMPORT_SEPARATOR))
+        }
+    }
+
     class object {
         public val IMPORT_SEPARATOR: String = ","
         public val RUNTIME_TYPE_EVALUATOR: Key<Function1<JetExpression, JetType?>> = Key.create("RUNTIME_TYPE_EVALUATOR")
-
-        public fun getImportsForElement(elementAtCaret: PsiElement): String {
-            val containingFile = elementAtCaret.getContainingFile()
-            if (containingFile !is JetFile) return ""
-
-            return containingFile.getImportList()?.getImports()
-                        ?.map { it.getText() }
-                        ?.join(JetCodeFragment.IMPORT_SEPARATOR) ?: ""
-        }
     }
 }

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

@@ -24,8 +24,9 @@ public class JetExpressionCodeFragment(
         project: Project,
         name: String,
         text: CharSequence,
+        imports: String?,
         context: PsiElement?
-) : JetCodeFragment(project, name, text, JetNodeTypes.EXPRESSION_CODE_FRAGMENT, context) {
+) : JetCodeFragment(project, name, text, imports, JetNodeTypes.EXPRESSION_CODE_FRAGMENT, context) {
 
     override fun getContentElement() = findChildByClass(javaClass<JetExpression>())
 }

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

@@ -323,11 +323,11 @@ public class JetPsiFactory(private val project: Project) {
     }
 
     public fun createExpressionCodeFragment(text: String, context: PsiElement?): JetExpressionCodeFragment {
-        return JetExpressionCodeFragment(project, "fragment.kt", text, context)
+        return JetExpressionCodeFragment(project, "fragment.kt", text, null, context)
     }
 
     public fun createBlockCodeFragment(text: String, context: PsiElement?): JetBlockCodeFragment {
-        return JetBlockCodeFragment(project, "fragment.kt", text, context)
+        return JetBlockCodeFragment(project, "fragment.kt", text, null, context)
     }
 
     public fun createReturn(text: String): JetReturnExpression {

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

@@ -25,7 +25,7 @@ import org.jetbrains.kotlin.types.JetType;
 
 public class JetTypeCodeFragment extends JetCodeFragment {
     public JetTypeCodeFragment(Project project, String name, CharSequence text, PsiElement context) {
-        super(project, name, text, JetNodeTypes.TYPE_CODE_FRAGMENT, context);
+        super(project, name, text, null, JetNodeTypes.TYPE_CODE_FRAGMENT, context);
     }
 
     @Nullable