diff --git a/idea/idea-analysis/src/org/jetbrains/kotlin/idea/caches/project/getModuleInfo.kt b/idea/idea-analysis/src/org/jetbrains/kotlin/idea/caches/project/getModuleInfo.kt
index cb7ef495533..cc1c111a491 100644
--- a/idea/idea-analysis/src/org/jetbrains/kotlin/idea/caches/project/getModuleInfo.kt
+++ b/idea/idea-analysis/src/org/jetbrains/kotlin/idea/caches/project/getModuleInfo.kt
@@ -5,6 +5,7 @@
 
 package org.jetbrains.kotlin.idea.caches.project
 
+import com.intellij.ide.scratch.ScratchFileService
 import com.intellij.openapi.module.Module
 import com.intellij.openapi.module.ModuleManager
 import com.intellij.openapi.project.Project
@@ -21,6 +22,7 @@ import org.jetbrains.kotlin.asJava.classes.KtLightClassForFacade
 import org.jetbrains.kotlin.asJava.elements.KtLightElement
 import org.jetbrains.kotlin.idea.caches.lightClasses.KtLightClassForDecompiledDeclaration
 import org.jetbrains.kotlin.idea.core.script.ScriptDependenciesManager
+import org.jetbrains.kotlin.idea.core.script.scriptRelatedModuleName
 import org.jetbrains.kotlin.idea.util.ProjectRootsUtil
 import org.jetbrains.kotlin.idea.util.isInSourceContentWithoutInjected
 import org.jetbrains.kotlin.idea.util.isKotlinBinary
@@ -230,6 +232,11 @@ private inline fun <T> collectInfosByVirtualFile(
         it.toIdeaModuleInfo(project, virtualFile, treatAsLibrarySource)?.let(onOccurrence)
     }
 
+    if (ScratchFileService.isInScratchRoot(virtualFile)) {
+        val scratchModule = virtualFile.scriptRelatedModuleName?.let { ModuleManager.getInstance(project).findModuleByName(it) }
+        onOccurrence(scratchModule?.testSourceInfo() ?: scratchModule?.productionSourceInfo())
+    }
+
     val scriptDefinition = getScriptDefinition(virtualFile, project)
     if (scriptDefinition != null) {
         onOccurrence(ScriptModuleInfo(project, virtualFile, scriptDefinition))
diff --git a/idea/idea-core/src/org/jetbrains/kotlin/idea/core/script/ModuleNameFileAttribute.kt b/idea/idea-core/src/org/jetbrains/kotlin/idea/core/script/ModuleNameFileAttribute.kt
new file mode 100644
index 00000000000..c8dcdce437a
--- /dev/null
+++ b/idea/idea-core/src/org/jetbrains/kotlin/idea/core/script/ModuleNameFileAttribute.kt
@@ -0,0 +1,31 @@
+/*
+ * Copyright 2010-2018 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license
+ * that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.idea.core.script
+
+import com.intellij.openapi.vfs.VirtualFile
+import com.intellij.openapi.vfs.VirtualFileWithId
+import com.intellij.openapi.vfs.newvfs.FileAttribute
+import kotlin.reflect.KProperty
+
+var VirtualFile.scriptRelatedModuleName: String? by ScratchModuleNameProperty()
+
+private val moduleNameAttribute = FileAttribute("kotlin-script-moduleName", 1, false)
+private class ScratchModuleNameProperty {
+
+    operator fun getValue(file: VirtualFile, property: KProperty<*>): String? {
+        if (file !is VirtualFileWithId) return null
+
+        return moduleNameAttribute.readAttributeBytes(file)?.let { String(it) }
+    }
+
+    operator fun setValue(file: VirtualFile, property: KProperty<*>, newValue: String?) {
+        if (file !is VirtualFileWithId) return
+
+        if (newValue != null) {
+            moduleNameAttribute.writeAttributeBytes(file, newValue.toByteArray())
+        }
+    }
+}
\ No newline at end of file
diff --git a/idea/idea-jvm/src/org/jetbrains/kotlin/idea/scratch/ScratchFileModuleInfoProvider.kt b/idea/idea-jvm/src/org/jetbrains/kotlin/idea/scratch/ScratchFileModuleInfoProvider.kt
index a47e387c04c..d4c30b3b558 100644
--- a/idea/idea-jvm/src/org/jetbrains/kotlin/idea/scratch/ScratchFileModuleInfoProvider.kt
+++ b/idea/idea-jvm/src/org/jetbrains/kotlin/idea/scratch/ScratchFileModuleInfoProvider.kt
@@ -22,23 +22,17 @@ import com.intellij.openapi.components.AbstractProjectComponent
 import com.intellij.openapi.diagnostic.Logger
 import com.intellij.openapi.fileEditor.FileEditorManager
 import com.intellij.openapi.fileEditor.FileEditorManagerListener
-import com.intellij.openapi.module.Module
 import com.intellij.openapi.module.ModuleManager
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.roots.ProjectRootManager
 import com.intellij.openapi.vfs.VirtualFile
-import com.intellij.openapi.vfs.newvfs.FileAttribute
 import com.intellij.psi.PsiManager
 import org.jetbrains.kotlin.idea.KotlinFileType
-import org.jetbrains.kotlin.idea.caches.project.NotUnderContentRootModuleInfo
-import org.jetbrains.kotlin.idea.caches.project.moduleInfo
-import org.jetbrains.kotlin.idea.caches.project.productionSourceInfo
-import org.jetbrains.kotlin.idea.caches.project.testSourceInfo
+import org.jetbrains.kotlin.idea.core.script.scriptRelatedModuleName
 import org.jetbrains.kotlin.idea.util.application.runWriteAction
 import org.jetbrains.kotlin.parsing.KotlinParserDefinition.Companion.STD_SCRIPT_EXT
 import org.jetbrains.kotlin.parsing.KotlinParserDefinition.Companion.STD_SCRIPT_SUFFIX
 import org.jetbrains.kotlin.psi.KtFile
-import kotlin.reflect.KProperty
 
 class ScratchFileModuleInfoProvider(project: Project) : AbstractProjectComponent(project) {
     private val LOG = Logger.getInstance(this.javaClass)
@@ -74,17 +68,15 @@ class ScratchFileModuleInfoProvider(project: Project) : AbstractProjectComponent
 
             val scratchPanel = getEditorWithScratchPanel(source, file)?.second
             scratchPanel?.addModuleListener { psiFile, module ->
-                psiFile.moduleInfo = getModuleInfo(module)
+                psiFile.virtualFile.scriptRelatedModuleName = module.name
+
                 // Drop caches for old module
                 ProjectRootManager.getInstance(myProject).incModificationCount()
                 // Force re-highlighting
                 DaemonCodeAnalyzer.getInstance(myProject).restart(psiFile)
-
-                psiFile.virtualFile.moduleName = module.name
             }
 
-            val module =
-                ktFile.virtualFile.moduleName?.let { ModuleManager.getInstance(myProject).findModuleByName(it) }
+            val module = ktFile.virtualFile.scriptRelatedModuleName?.let { ModuleManager.getInstance(myProject).findModuleByName(it) }
             if (module != null) {
                 scratchPanel?.setModule(module)
             } else {
@@ -95,23 +87,4 @@ class ScratchFileModuleInfoProvider(project: Project) : AbstractProjectComponent
             }
         }
     }
-
-    private fun getModuleInfo(it: Module?) =
-        it?.testSourceInfo() ?: it?.productionSourceInfo() ?: NotUnderContentRootModuleInfo
-}
-
-private var VirtualFile.moduleName: String? by ScratchModuleNameProperty()
-
-private val moduleNameAttribute = FileAttribute("ScratchModuleName", 1, false)
-private class ScratchModuleNameProperty {
-
-    operator fun getValue(file: VirtualFile, property: KProperty<*>): String? {
-        return moduleNameAttribute.readAttributeBytes(file)?.let { String(it) }
-    }
-
-    operator fun setValue(file: VirtualFile, property: KProperty<*>, newValue: String?) {
-        if (newValue != null) {
-            moduleNameAttribute.writeAttributeBytes(file, newValue.toByteArray())
-        }
-    }
 }
\ No newline at end of file