Features/#462 gutter icons module case

CHANGELOG.md

@@ -10,6 +10,7 @@ Released on ...
 - TODO  (see [#NNN](https://github.com/JetBrains-Research/snakecharm/issues/NNN)) 
 
 ### Fixed
+- Complex cases in creation of gutter icons for rule inheritance (see [#462](https://github.com/JetBrains-Research/snakecharm/issues/462), [#463](https://github.com/JetBrains-Research/snakecharm/issues/463))
 - Resolve for rule names in `use rule` section (see [#455](https://github.com/JetBrains-Research/snakecharm/issues/455))
 - Multiple args inspection in `workdir` case (see [#140](https://github.com/JetBrains-Research/snakecharm/issues/140))
 - `localrules` and `ruleorder` now take into account `use rule` (see [#448](https://github.com/JetBrains-Research/snakecharm/issues/448))

src/main/kotlin/com/jetbrains/snakecharm/SnakemakeRuleInheritanceMarkerProvider.kt

@@ -9,7 +9,8 @@ import com.intellij.psi.PsiElement
 import com.intellij.psi.search.GlobalSearchScope
 import com.intellij.psi.stubs.StubIndex
 import com.jetbrains.snakecharm.lang.psi.*
-import com.jetbrains.snakecharm.lang.psi.stubs.SmkModuleNameIndex.Companion.KEY
+import com.jetbrains.snakecharm.lang.psi.stubs.SmkUseInheritedRulesIndex
+import com.jetbrains.snakecharm.lang.psi.stubs.SmkUseInheritedRulesIndex.Companion.INHERITED_RULES_DECLARATION_VIA_WILDCARD
 
 class SnakemakeRuleInheritanceMarkerProvider : RelatedItemLineMarkerProvider() {
 
@@ -51,32 +52,32 @@ class SnakemakeRuleInheritanceMarkerProvider : RelatedItemLineMarkerProvider() {
         element: SmkRuleOrCheckpoint,
         result: MutableCollection<in RelatedItemLineMarkerInfo<*>>
     ) {
-        val currentFile = element.containingFile
         val identifier = when (val identifier = element.nameIdentifier) {
             is SmkUseNameIdentifier -> identifier.getNameBeforeWildcard()
             else -> identifier
         } ?: return
-        val modulesFromStub = StubIndex.getInstance().getAllKeys(KEY, element.project)
         val module = element.let { ModuleUtilCore.findModuleForPsiElement(it.originalElement) } ?: return
-        val files = mutableSetOf<SmkFile>()
-        for (smkModule in modulesFromStub) {
-            files.addAll(StubIndex.getElements(
-                KEY,
-                smkModule,
-                module.project,
-                GlobalSearchScope.moduleWithDependentsScope(module),
-                SmkModule::class.java
-            ).mapNotNull { it.containingFile as? SmkFile })
-        }
-        if (currentFile is SmkFile) {
-            files.add(currentFile)
-        }
+        val descendantsDefinedExplicitly = StubIndex.getElements(
+            SmkUseInheritedRulesIndex.KEY,
+            identifier.text,
+            module.project,
+            GlobalSearchScope.moduleWithDependentsScope(module),
+            SmkUse::class.java
+        )
+        val potentialDescendants = StubIndex.getElements(
+            SmkUseInheritedRulesIndex.KEY,
+            INHERITED_RULES_DECLARATION_VIA_WILDCARD,
+            module.project,
+            GlobalSearchScope.moduleWithDependentsScope(module),
+            SmkUse::class.java
+        )
         val overrides = mutableListOf<SmkRuleOrCheckpoint>()
-        files.forEach { file ->
-            file.collectUses().forEach { (_, psi) ->
-                if (psi.getImportedRules()?.firstOrNull { it == element } != null) {
-                    overrides.add(psi)
-                }
+        (descendantsDefinedExplicitly + potentialDescendants).forEach { descendant ->
+            // We don't save it in stub because it requires 'resolve'
+            // We compare resolve results even for descendantsDefinedExplicitly
+            // Because there may be rules with the same names
+            if (descendant.getImportedRules()?.contains(element) == true) {
+                overrides.add(descendant)
             }
         }
         if (overrides.isEmpty()) {

src/main/kotlin/com/jetbrains/snakecharm/lang/psi/impl/refs/SmkRuleOrCheckpointNameReference.kt

@@ -61,12 +61,18 @@ class SmkRuleOrCheckpointNameReference(
             return results
         }
         val allImportedFiles = smkFile.collectIncludedFiles()
+        val potentialModule = moduleRef?.reference?.resolve() as? SmkModule
         return results.filter { resolveResult ->
             // If we resolve module references, there must be only SmkModules
             (resolveResult.element is SmkModule && itIsModuleMameReference) ||
                     // We don't want to suggest local resolve result for the reference of rule, which was imported
                     (moduleRef != null // Module name reference is defined and resolve result is from another file
-                            && element.containingFile.originalFile != resolveResult.element?.containingFile?.originalFile)
+                            && element.containingFile.originalFile != resolveResult.element?.containingFile?.originalFile
+                            // Also, because of using indexes, we need to check if the resolve result file
+                            // Connected to the file, which was declared in moduleRef, via 'include' or 'module'
+                            // Later, allImportedFiles will be stored in cache
+                            && resolveResult.element?.containingFile?.originalFile in ((potentialModule?.getPsiFile() as? SmkFile)?.collectIncludedFiles()
+                        ?: listOf()))
                     // OR There are no 'from *name*' combination, so it hasn't been imported
                     || (moduleRef == null && resolveResult.element?.containingFile?.originalFile in allImportedFiles)
         }.toMutableList()

src/main/kotlin/com/jetbrains/snakecharm/lang/psi/impl/stubs/SmkRuleDescendantElementType.kt

@@ -0,0 +1,48 @@
+package com.jetbrains.snakecharm.lang.psi.impl.stubs
+
+import com.intellij.psi.stubs.*
+import com.jetbrains.python.psi.PyElement
+import com.jetbrains.python.psi.PyStubElementType
+import com.jetbrains.snakecharm.lang.psi.SmkRuleLike
+import com.jetbrains.snakecharm.lang.psi.stubs.RuleDescendantStub
+import com.jetbrains.snakecharm.lang.psi.stubs.SmkUseInheritedRulesIndex
+
+abstract class SmkRuleDescendantElementType<StubT : RuleDescendantStub<*>, PsiT : PyElement>(
+    debugName: String,
+    private val nameIndexKey: StubIndexKey<String, out SmkRuleLike<*>>?
+) :
+    PyStubElementType<StubT, PsiT>(debugName) {
+    abstract fun createStub(name: String?, inheritedRules: List<String?>, parentStub: StubElement<*>?): StubT
+
+    override fun serialize(stub: StubT, dataStream: StubOutputStream) {
+        dataStream.writeName(stub.name)
+        dataStream.writeInt(stub.getInheritedRulesNames().size)
+        stub.getInheritedRulesNames().forEach { name ->
+            dataStream.writeName(name)
+        }
+    }
+
+    override fun deserialize(dataStream: StubInputStream, parentStub: StubElement<*>?): StubT {
+        val name = dataStream.readNameString()
+        val size = dataStream.readInt()
+        val inheritedRules = mutableListOf<String?>()
+        for (x in 0 until size) {
+            inheritedRules.add(dataStream.readNameString())
+        }
+        return createStub(name, inheritedRules, parentStub)
+    }
+
+    override fun indexStub(stub: StubT, sink: IndexSink) {
+        if (nameIndexKey != null) {
+            stub.name?.let { name ->
+                sink.occurrence(nameIndexKey, name)
+            }
+        }
+        stub.getInheritedRulesNames().forEach { inheritedName ->
+            if (inheritedName != null) {
+                sink.occurrence(SmkUseInheritedRulesIndex.KEY, inheritedName)
+            }
+        }
+        super.indexStub(stub, sink)
+    }
+}

src/main/kotlin/com/jetbrains/snakecharm/lang/psi/impl/stubs/SmkRuleLikeStubImpl.kt

@@ -18,6 +18,16 @@ abstract class SmkRuleLikeStubImpl<StubT : NamedStub<PsiT>, PsiT>(
     override fun getName() = myName
 }
 
+abstract class SmkRuleDescendantStubImpl<StubT : RuleDescendantStub<PsiT>, PsiT>(
+    private val myName: String?,
+    private val inheritedRulesNames: List<String?>,
+    parent: StubElement<*>?,
+    type: IStubElementType<StubT, PsiT>
+) : StubBase<PsiT>(parent, type), RuleDescendantStub<PsiT> where PsiT : PsiNamedElement, PsiT : PyElement {
+    override fun getName() = myName
+    override fun getInheritedRulesNames() = inheritedRulesNames
+}
+
 class SmkCheckpointStubImpl(
     name: String?,
     parent: StubElement<*>?
@@ -42,5 +52,7 @@ class SmkModuleStubImpl(
 
 class SmkUseStubImpl(
     name: String?,
+    inheritedRulesNames: List<String?>,
     parent: StubElement<*>?
-) : SmkRuleLikeStubImpl<SmkUseStub, SmkUse>(name, parent, USE_DECLARATION_STATEMENT), SmkUseStub
+) : SmkRuleDescendantStubImpl<SmkUseStub, SmkUse>(name, inheritedRulesNames, parent, USE_DECLARATION_STATEMENT),
+    SmkUseStub

src/main/kotlin/com/jetbrains/snakecharm/lang/psi/impl/stubs/SmkUseElementType.kt

@@ -5,17 +5,25 @@ import com.intellij.psi.PsiElement
 import com.intellij.psi.stubs.StubElement
 import com.jetbrains.snakecharm.lang.psi.SmkUse
 import com.jetbrains.snakecharm.lang.psi.impl.SmkUseImpl
+import com.jetbrains.snakecharm.lang.psi.stubs.SmkUseInheritedRulesIndex.Companion.INHERITED_RULES_DECLARATION_VIA_WILDCARD
 import com.jetbrains.snakecharm.lang.psi.stubs.SmkUseNameIndex.Companion.KEY
 import com.jetbrains.snakecharm.lang.psi.stubs.SmkUseStub
 
 class SmkUseElementType
-    : SmkRuleLikeElementType<SmkUseStub, SmkUse>("SMK_USE_DECLARATION_STATEMENT", KEY) {
+    : SmkRuleDescendantElementType<SmkUseStub, SmkUse>("SMK_USE_DECLARATION_STATEMENT", KEY) {
 
-    override fun createStub(name: String?, parentStub: StubElement<*>?): SmkUseStub =
-        SmkUseStubImpl(name, parentStub)
+    lateinit var psi: SmkUse
+
+    override fun createStub(name: String?, inheritedRules: List<String?>, parentStub: StubElement<*>?) =
+        SmkUseStubImpl(name, inheritedRules, parentStub)
 
     override fun createStub(psi: SmkUse, parentStub: StubElement<out PsiElement>?): SmkUseStub =
-        createStub(psi.name, parentStub)
+        createStub(
+            psi.name,
+            psi.getDefinedReferencesOfImportedRuleNames()?.map { it.text } ?: listOf(
+                INHERITED_RULES_DECLARATION_VIA_WILDCARD
+            ),
+            parentStub).also { this.psi = psi }
 
     override fun createPsi(stub: SmkUseStub): SmkUse = SmkUseImpl(stub)
 

src/main/kotlin/com/jetbrains/snakecharm/lang/psi/stubs/RuleDescendantStub.kt

@@ -0,0 +1,11 @@
+package com.jetbrains.snakecharm.lang.psi.stubs
+
+import com.intellij.psi.PsiNamedElement
+import com.intellij.psi.stubs.NamedStub
+
+/**
+ * Stub of rule like element which inherits at least one rule
+ */
+interface RuleDescendantStub<T : PsiNamedElement> : NamedStub<T> {
+    fun getInheritedRulesNames(): List<String?>
+}

src/main/kotlin/com/jetbrains/snakecharm/lang/psi/stubs/SmkStubs.kt

@@ -7,4 +7,4 @@ interface SmkRuleStub : NamedStub<SmkRule>
 interface SmkCheckpointStub : NamedStub<SmkCheckPoint>
 interface SmkSubworkflowStub : NamedStub<SmkSubworkflow>
 interface SmkModuleStub : NamedStub<SmkModule>
-interface SmkUseStub : NamedStub<SmkUse>
+interface SmkUseStub : RuleDescendantStub<SmkUse>

src/main/kotlin/com/jetbrains/snakecharm/lang/psi/stubs/SmkUseInheritedRulesIndex.kt

@@ -0,0 +1,16 @@
+package com.jetbrains.snakecharm.lang.psi.stubs
+
+import com.intellij.psi.stubs.StringStubIndexExtension
+import com.intellij.psi.stubs.StubIndexKey
+import com.jetbrains.snakecharm.lang.psi.SmkUse
+
+class SmkUseInheritedRulesIndex :
+    StringStubIndexExtension<SmkUse>() {
+    override fun getKey() = KEY
+
+    companion object {
+        const val INHERITED_RULES_DECLARATION_VIA_WILDCARD = "*"
+
+        val KEY = StubIndexKey.createIndexKey<String, SmkUse>("Smk.use.inheritedRules")
+    }
+}

src/main/resources/META-INF/plugin.xml

@@ -94,6 +94,7 @@
     <lang.documentationProvider language="Snakemake" implementationClass="com.jetbrains.snakecharm.lang.documentation.SmkWrapperDocumentation"/>
     <lang.psiStructureViewFactory language="Snakemake" implementationClass="com.jetbrains.snakecharm.lang.structureView.SmkStructureViewFactory"/>
     <stubIndex implementation="com.jetbrains.snakecharm.lang.psi.stubs.SmkRuleNameIndex"/>
+    <stubIndex implementation="com.jetbrains.snakecharm.lang.psi.stubs.SmkUseInheritedRulesIndex"/>
     <stubIndex implementation="com.jetbrains.snakecharm.lang.psi.stubs.SmkCheckpointNameIndex"/>
     <stubIndex implementation="com.jetbrains.snakecharm.lang.psi.stubs.SmkUseNameIndex"/>
     <stubIndex implementation="com.jetbrains.snakecharm.lang.psi.stubs.SmkModuleNameIndex"/>

src/test/resources/features/highlighting/rule_inheritance_marker_provider.feature

@@ -35,9 +35,14 @@ Feature: tests line marker provider in case of rule inheritance
 
   Scenario Outline: Rule inheritance in case of module usage
     Given a snakemake project
+    And a file "doo.smk" with text
+    """
+    <rule_like> NAME: # doo.smk
+      input: "doo_name_inp"
+    """
     And a file "boo.smk" with text
     """
-    <rule_like> NAME:
+    <rule_like> NAME: # boo.smk
       input: "name_inp"
 
     <rule_like> NAME_2:
@@ -61,8 +66,11 @@ Feature: tests line marker provider in case of rule inheritance
 
     use rule * from MODULE as IMPLICIT_DEFINITION_OF_*
     """
+    When I change current file to <doo.smk>
+    And I put the caret at NAME: # doo.smk
+    Then I expect no markers
     When I change current file to <boo.smk>
-    And I put the caret at NAME:
+    And I put the caret at NAME: # boo.smk
     Then I expect marker of overridden section with references:
       | NAME_3                   | foo.smk |
       | NAME_4                   | foo.smk |
@@ -84,7 +92,7 @@ Feature: tests line marker provider in case of rule inheritance
     Then I expect no markers
     When I put the caret at EXPLICIT_DEFINITION_OF_*
     Then I expect marker of overriding section with references:
-      | NAME | boo.smk |
+      | NAME   | boo.smk |
       | NAME_2 | boo.smk |
     When I put the caret at IMPLICIT_DEFINITION_OF_*
     Then I expect marker of overriding section with references:
@@ -156,6 +164,71 @@ Feature: tests line marker provider in case of rule inheritance
     Then I expect marker of overriding section with references:
       | NAME   | doo.smk |
       | NAME_2 | doo.smk |
+    Examples:
+      | rule_like  |
+      | rule       |
+      | checkpoint |
+
+  Scenario Outline: Rule inheritance in case of include in the main file, issue #463
+    Given a snakemake project
+    And a file "boo.smk" with text
+    """
+    <rule_like> NAME:
+      input: "name_inp"
+
+    <rule_like> NAME_2:
+      input: "name_2_inp"
+    """
+    Given I open a file "foo.smk" with text
+    """
+    include: "boo.smk"
+
+    use rule NAME as NAME_3 with:
+      threads: 1
+    """
+    When I change current file to <boo.smk>
+    And I put the caret at NAME:
+    Then I expect marker of overridden section with references:
+      | NAME_3 | foo.smk |
+    When I change current file to <boo.smk>
+    And I put the caret at NAME_2:
+    Then I expect no markers
+    When I change current file to <foo.smk>
+    And I put the caret at NAME_3
+    Then I expect marker of overriding section with references:
+      | NAME | boo.smk |
+    Examples:
+      | rule_like  |
+      | rule       |
+      | checkpoint |
+
+  Scenario Outline: Rule inheritance in case of include in the main file, issue #462
+    Given a snakemake project
+    And a file "doo.smk" with text
+    """
+    <rule_like> NAME:
+      input: "name_inp"
+    """
+    And a file "boo.smk" with text
+    """
+    module DOO:
+      snakefile: "doo.smk"
+    """
+    Given I open a file "foo.smk" with text
+    """
+    include: "boo.smk"
+
+    use rule NAME from DOO as NAME_3 with:
+      threads: 1
+    """
+    When I change current file to <doo.smk>
+    And I put the caret at NAME:
+    Then I expect marker of overridden section with references:
+      | NAME_3 | foo.smk |
+    When I change current file to <foo.smk>
+    And I put the caret at NAME_3
+    Then I expect marker of overriding section with references:
+      | NAME | doo.smk |
     Examples:
       | rule_like  |
       | rule       |