Add explicit type intention

src/main/kotlin/com/mdrobnak/lalrpop/intentions/AddExplicitTypeIntention.kt

@@ -0,0 +1,64 @@
+package com.mdrobnak.lalrpop.intentions
+
+import com.intellij.codeInsight.intention.IntentionAction
+import com.intellij.openapi.editor.Editor
+import com.intellij.openapi.project.Project
+import com.intellij.psi.PsiFile
+import com.mdrobnak.lalrpop.LpLanguage
+import com.mdrobnak.lalrpop.psi.LpAlternative
+import com.mdrobnak.lalrpop.psi.LpElementFactory
+import com.mdrobnak.lalrpop.psi.LpNonterminal
+import com.mdrobnak.lalrpop.psi.getContextAndResolveType
+
+class AddExplicitTypeIntention : IntentionAction {
+    override fun startInWriteAction(): Boolean = true
+
+    override fun getText(): String = "Add explicit type"
+
+    override fun getFamilyName(): String = "Add explicit type"
+
+    override fun isAvailable(project: Project, editor: Editor?, file: PsiFile?): Boolean {
+        if (editor == null || file == null || file.language != LpLanguage) return false
+
+        var element = file.findElementAt(editor.caretModel.primaryCaret.offset)
+        while (element != null) {
+            if (element is LpNonterminal) return true
+            element = element.parent
+        }
+        return false
+    }
+
+    override fun invoke(project: Project, editor: Editor?, file: PsiFile?) {
+        if (editor == null || file == null || file.language != LpLanguage) return
+
+        var element = file.findElementAt(editor.caretModel.primaryCaret.offset)
+
+        var nonterminal: LpNonterminal? = null
+        var alternativeToUse: LpAlternative? = null
+
+        while (element != null) {
+            if (element is LpAlternative) {
+                if (element.action == null) // TODO: maybe get from the rust plugin?
+                    alternativeToUse = element
+            } else if (element is LpNonterminal) {
+                nonterminal = element
+                break
+            }
+
+            element = element.parent
+        }
+
+        if (nonterminal == null) return
+
+        val inferredType =
+            alternativeToUse?.getContextAndResolveType(listOf()) ?: nonterminal.getContextAndResolveType(listOf())
+
+        val type = LpElementFactory(project).createNonterminalType(inferredType)
+        val typeRef = nonterminal.typeRef
+
+        if (typeRef != null)
+            typeRef.replace(type.second)
+        else
+            nonterminal.addRangeAfter(type.first, type.second, nonterminal.nonterminalName)
+    }
+}

src/main/kotlin/com/mdrobnak/lalrpop/psi/LpElementFactory.kt

@@ -43,4 +43,19 @@ class LpElementFactory(val project: Project) {
 
     fun createWhitespace(s: String): PsiWhiteSpace = createFromText("${s}grammar;")
         ?: error("Failed to create whitespace '$s'")
+
+    /**
+     * Returns a pair where the first element corresponds to the ':' and the second to the type_ref.
+     * Used with `nonterminal.addRangeAfter(pair.first, pair.second, nonterminal.name)` to add the type to a nonterminal
+     * where a type doesn't exist already.
+     */
+    fun createNonterminalType(type: String): Pair<PsiElement, LpTypeRef> {
+        val nonterminal = createFromText<LpNonterminal>("grammar;\n Nonterminal: $type = {};") ?: error("Failed to create nonterminal")
+
+        val typeRef = nonterminal.typeRef!!
+        val whitespace = typeRef.prevSibling
+        val colon = whitespace.prevSibling
+
+        return colon to typeRef
+    }
 }

src/main/kotlin/com/mdrobnak/lalrpop/psi/LpResolveType.kt

@@ -1,6 +1,7 @@
 package com.mdrobnak.lalrpop.psi
 
 import com.intellij.psi.PsiElement
+import com.mdrobnak.lalrpop.psi.util.lalrpopTypeResolutionContext
 
 data class NonterminalGenericArgument(val rustType: String, var name: String)
 
@@ -30,3 +31,6 @@ interface LpResolveType : PsiElement {
 //    fun resolveType(arguments: List<NonterminalGenericArgument>): String =
 //        this.resolveType(this.containingFile.lalrpopTypeResolutionContext(), arguments)
 }
+
+fun LpResolveType.getContextAndResolveType(arguments: List<NonterminalGenericArgument>): String =
+    this.resolveType(this.containingFile.lalrpopTypeResolutionContext(), arguments)

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

@@ -84,5 +84,10 @@
                          enabledByDefault="true"
                          level="ERROR"/>
 
+        <intentionAction id="com.mdrobnak.lalrpop.intentions.AddExplicitTypeIntention">
+            <className>com.mdrobnak.lalrpop.intentions.AddExplicitTypeIntention</className>
+            <category>LALRPOP</category>
+        </intentionAction>
+
     </extensions>
 </idea-plugin>

src/main/resources/intentionDescriptions/AddExplicitTypeIntention/after.lalrpop.template

@@ -0,0 +1,5 @@
+grammar;
+
+A = ();
+
+B: () = A;

src/main/resources/intentionDescriptions/AddExplicitTypeIntention/before.lalrpop.template

@@ -0,0 +1,5 @@
+grammar;
+
+A = ();
+
+B = A;

src/main/resources/intentionDescriptions/AddExplicitTypeIntention/description.html

@@ -0,0 +1,10 @@
+<html>
+<body>
+<p>
+    Adds the inferred type of the alternative the caret is on to the parent nonterminal, or tries to infer the
+    type from the other alternatives if this one has action code.
+</p>
+<!-- tooltip end -->
+<p></p>
+</body>
+</html>