Add explicit type intention

build.gradle

@@ -14,6 +14,11 @@ repositories {
     mavenCentral()
 }
 
+dependencies {
+    testImplementation "org.junit.jupiter:junit-jupiter-api:5.7.0"
+    testImplementation "org.junit.jupiter:junit-jupiter-engine:5.7.0"
+}
+
 // Add the generated source files into the build
 sourceSets.main.java.srcDirs 'src/main/gen'
 idea.module.generatedSourceDirs.add(file('src/main/gen'))
@@ -67,4 +72,10 @@ patchPluginXml {
       <b>0.1.0</b><br>
       Initial release of the LALRPOP plugin
       """
+}
+
+test {
+    useJUnitPlatform()
+
+    maxHeapSize = "1G"
 }

src/main/grammars/LalrpopParser.bnf

@@ -11,6 +11,8 @@
   elementTypeHolderClass="com.mdrobnak.lalrpop.psi.LpElementTypes"
   elementTypeClass="com.mdrobnak.lalrpop.psi.LpElementType"
   tokenTypeClass="com.mdrobnak.lalrpop.psi.LpTokenType"
+
+  generateTokenAccessors=true
 }
 
 grammar ::= SHEBANG_ATTRIBUTE* use_stmt*
@@ -99,7 +101,7 @@ alternative ::= annotation* symbol+ (IF cond)? action?
 }
 
 action ::= action_type CODE {
-    implements = "com.intellij.psi.PsiLanguageInjectionHost"
+    implements = ["com.intellij.psi.PsiLanguageInjectionHost" "com.mdrobnak.lalrpop.psi.LpResolveType"]
     mixin = "com.mdrobnak.lalrpop.psi.ext.LpActionMixin"
 }
 
@@ -168,7 +170,7 @@ type_of_symbol ::= POUND symbol POUND {
     implements = "com.mdrobnak.lalrpop.psi.LpResolveType"
     mixin = "com.mdrobnak.lalrpop.psi.ext.LpTypeOfSymbolMixin"
 }
-rust_reference ::= AND lifetime_rule? MUT? type_ref {
+rust_reference ::= AND LIFETIME? MUT? type_ref {
     implements = "com.mdrobnak.lalrpop.psi.LpResolveType"
     mixin = "com.mdrobnak.lalrpop.psi.ext.LpRustReferenceMixin"
 }
@@ -188,9 +190,7 @@ return_type ::= RSINGLEARROW type_ref
 
 type_generic_arguments ::= LESSTHAN <<comma type_ref_or_lifetime>> GREATERTHAN
 
-private type_ref_or_lifetime ::= type_ref | lifetime_rule
-
-lifetime_rule ::= LIFETIME
+type_ref_or_lifetime ::= type_ref | LIFETIME
 
 extern_token ::= EXTERN LBRACE associated_type* enum_token associated_type* RBRACE
     | EXTERN LBRACE associated_type* RBRACE {
@@ -215,9 +215,7 @@ enum_token ::= ENUM type_ref LBRACE <<comma conversion>> RBRACE {
     mixin = "com.mdrobnak.lalrpop.psi.ext.LpEnumTokenMixin"
 }
 
-associated_type ::= TYPE associated_type_name EQUALS type_ref SEMICOLON
-
-associated_type_name ::= ID
+associated_type ::= TYPE ID EQUALS type_ref SEMICOLON
 
 conversion ::= terminal USER_ACTION CODE
 

src/main/kotlin/com/mdrobnak/lalrpop/injectors/LpRustActionCodeInjector.kt

@@ -17,55 +17,18 @@ import org.rust.lang.RsLanguage
  */
 class LpRustActionCodeInjector : MultiHostInjector {
     override fun getLanguagesToInject(registrar: MultiHostRegistrar, context: PsiElement) {
-        if (!context.isValid || context !is LpActionImpl) {
+        if (!context.isValid || context !is LpAction) {
             return
         }
 
-        val codeNode = context.code ?: return
+        val codeNode = context.code
         val imports = context.containingFile.importCode
-        val nonterminal = context.parentOfType<LpNonterminal>()!!
-        val alternative = context.parentOfType<LpAlternative>()!!
 
-        val typeResolutionContext = context.containingFile.lalrpopTypeResolutionContext()
-
-        val inputs = alternative.selected.map { (it as LpSymbolImpl).getSelectedType(typeResolutionContext) }
-        val returnType = context.actionType.returnType(nonterminal.resolveType(typeResolutionContext, listOf()), typeResolutionContext)
-
-        val grammarDecl = PsiTreeUtil.findChildOfType(context.containingFile, LpGrammarDecl::class.java)
-
-        val grammarParams = grammarDecl?.grammarParams
-        val grammarParametersString =
-            grammarParams?.grammarParamList?.joinToString(separator = "") { "${it.name}: ${it.typeRef.text}," }
-                ?: ""
-
-        val grammarTypeParams = grammarDecl?.grammarTypeParams
-        val genericParameters = nonterminal.nonterminalName.nonterminalParams?.nonterminalParamList
-
-        fun addNullableIterators(a: List<String>?, b: List<String>?): List<String> = (a ?: listOf()) + (b ?: listOf())
-        fun List<String>.join(): String =
-            if (this.isEmpty()) "" else this.joinToString(prefix = "<", postfix = ">", separator = ", ") { it }
-
-        val grammarTypeParamsString = addNullableIterators(
-            grammarTypeParams?.typeParamList?.map { it.text },
-            genericParameters?.map { it.text }
-        ).join()
-
-        val arguments = inputs.mapIndexed { index, it ->
-            when (it) {
-                is LpSelectedType.WithName -> (if (it.isMutable) "mut " else "") + it.name + ": " + it.type
-                is LpSelectedType.WithoutName -> "__intellij_lalrpop_noname_$index: " + it.type
-            }
-        }.joinToString(", ")
-
-        val grammarWhereClauses = grammarDecl?.grammarWhereClauses
-        val grammarWhereClausesString =
-            grammarWhereClauses?.grammarWhereClauseList?.joinToString(prefix = "where ", separator = ", ") { it.text }
-                ?: ""
+        val rustFunctionHeader = context.actionCodeFunctionHeader(true)
 
         val prefix = "mod __intellij_lalrpop {\n" +
                 "$imports\n" +
-                "fn __intellij_lalrpop $grammarTypeParamsString ($grammarParametersString $arguments) -> $returnType" +
-                " $grammarWhereClausesString {\n"
+                "$rustFunctionHeader {\n"
 
         val suffix = "\n}\n}"
 
@@ -78,5 +41,5 @@ class LpRustActionCodeInjector : MultiHostInjector {
     }
 
     override fun elementsToInjectIn(): List<Class<out PsiElement>> =
-        listOf(LpActionImpl::class.java)
+        listOf(LpAction::class.java)
 }

src/main/kotlin/com/mdrobnak/lalrpop/inspections/CannotInferNonterminalTypeInspection.kt

@@ -1,10 +1,16 @@
 package com.mdrobnak.lalrpop.inspections
 
 import com.intellij.codeInspection.LocalInspectionTool
+import com.intellij.codeInspection.LocalQuickFix
+import com.intellij.codeInspection.ProblemDescriptor
 import com.intellij.codeInspection.ProblemsHolder
+import com.intellij.openapi.project.Project
 import com.intellij.psi.PsiElementVisitor
+import com.mdrobnak.lalrpop.psi.LpMacroArguments
 import com.mdrobnak.lalrpop.psi.LpNonterminal
 import com.mdrobnak.lalrpop.psi.LpVisitor
+import com.mdrobnak.lalrpop.psi.ext.setType
+import com.mdrobnak.lalrpop.psi.getContextAndResolveType
 
 object CannotInferNonterminalTypeInspection : LocalInspectionTool() {
     override fun buildVisitor(holder: ProblemsHolder, isOnTheFly: Boolean): PsiElementVisitor {
@@ -17,9 +23,20 @@ object CannotInferNonterminalTypeInspection : LocalInspectionTool() {
                     holder.registerProblem(
                         nonterminal,
                         "Cannot infer type of nonterminal",
-                    ) // TODO Quickfix: get from the rust plugin if available.
+                        InferFromRustPluginQuickFix,
+                    )
                 }
             }
         }
     }
-}
+}
+
+object InferFromRustPluginQuickFix : LocalQuickFix {
+    override fun getFamilyName(): String = "Get from action code"
+
+    override fun applyFix(project: Project, descriptor: ProblemDescriptor) {
+        (descriptor.psiElement as LpNonterminal).apply {
+            setType(getContextAndResolveType(LpMacroArguments.identity(nonterminalName.nonterminalParams)))
+        }
+    }
+}

src/main/kotlin/com/mdrobnak/lalrpop/inspections/WrongInferredTypeInspection.kt

@@ -7,18 +7,17 @@ import com.intellij.psi.PsiElementVisitor
 import com.intellij.psi.PsiFile
 import com.intellij.psi.PsiFileFactory
 import com.intellij.psi.util.PsiTreeUtil
-import com.jetbrains.rd.util.string.printToString
 import com.mdrobnak.lalrpop.injectors.findModuleDefinition
+import com.mdrobnak.lalrpop.psi.LpMacroArguments
 import com.mdrobnak.lalrpop.psi.LpNonterminal
 import com.mdrobnak.lalrpop.psi.LpVisitor
 import com.mdrobnak.lalrpop.psi.ext.importCode
+import com.mdrobnak.lalrpop.psi.ext.rustGenericUnitStructs
 import com.mdrobnak.lalrpop.psi.util.lalrpopTypeResolutionContext
 import org.rust.lang.RsLanguage
 import org.rust.lang.core.macros.RsExpandedElement
 import org.rust.lang.core.macros.setContext
-import org.rust.lang.core.psi.RsModDeclItem
 import org.rust.lang.core.psi.RsTypeAlias
-import org.rust.lang.core.psi.ext.childOfType
 import org.rust.lang.core.psi.ext.childrenOfType
 import org.rust.lang.core.resolve.ImplLookup
 
@@ -34,22 +33,20 @@ object WrongInferredTypeInspection : LocalInspectionTool() {
     ): PsiElementVisitor {
         return object : LpVisitor() {
             override fun visitNonterminal(nonterminal: LpNonterminal) {
-                // TODO: idea to implement this inspection for lalrpop macros:
-                // https://github.com/Mcat12/intellij-lalrpop/pull/22#discussion_r542538737
-                if (nonterminal.nonterminalName.nonterminalParams != null) return
-
                 val context = nonterminal.containingFile.lalrpopTypeResolutionContext()
 
-                val explicitType = nonterminal.typeRef?.resolveType(context, listOf())
+                val explicitType = nonterminal.typeRef?.resolveType(context, LpMacroArguments())
 
                 // LALRPOP will automatically supply action code of (), so we don't need to worry about inferring the wrong type.
                 // https://github.com/lalrpop/lalrpop/blob/8a96e9646b3d00c2226349efed832c4c25631c53/lalrpop/src/normalize/lower/mod.rs#L351-L354
                 if (explicitType == "()") return
 
                 var seenType = explicitType
 
+                val unitStructsGenerics by lazy { nonterminal.rustGenericUnitStructs() }
+
                 nonterminal.alternatives.alternativeList.filter { it.action == null }.forEach { alternative ->
-                    val alternativeType = alternative.resolveType(context, listOf())
+                    val alternativeType = alternative.resolveType(context, LpMacroArguments())
 
                     if (seenType == null) {
                         seenType = alternativeType
@@ -60,6 +57,7 @@ object WrongInferredTypeInspection : LocalInspectionTool() {
                             nonterminal.project,
                             nonterminal.containingFile,
                             nonterminal.containingFile.importCode,
+                            unitStructsGenerics,
                             seenType!!,
                             alternativeType
                         )
@@ -74,9 +72,24 @@ object WrongInferredTypeInspection : LocalInspectionTool() {
         }
     }
 
-    fun sameTypes(project: Project, lalrpopFile: PsiFile, importCode: String, type1: String, type2: String): Boolean {
+    fun sameTypes(
+        project: Project,
+        lalrpopFile: PsiFile,
+        importCode: String,
+        unitStructsGenerics: String,
+        type1: String,
+        type2: String
+    ): Boolean {
         val file = PsiFileFactory.getInstance(project)
-            .createFileFromText(RsLanguage, "mod __intellij_lalrpop {\n$importCode\ntype T1 = $type1;\ntype T2 = $type2;\n}")
+            .createFileFromText(
+                RsLanguage,
+                "mod __intellij_lalrpop {\n" +
+                        "    $importCode\n" +
+                        "    $unitStructsGenerics\n" +
+                        "    type T1 = $type1;\n" +
+                        "    type T2 = $type2;\n" +
+                        "}"
+            )
         val aliases = PsiTreeUtil.findChildrenOfType(file, RsTypeAlias::class.java)
 
         if (aliases.size != 2) return false

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

@@ -0,0 +1,34 @@
+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.intellij.psi.util.parentOfType
+import com.mdrobnak.lalrpop.LpLanguage
+import com.mdrobnak.lalrpop.psi.LpMacroArguments
+import com.mdrobnak.lalrpop.psi.LpNonterminal
+import com.mdrobnak.lalrpop.psi.ext.setType
+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
+
+        return file.findElementAt(editor.caretModel.primaryCaret.offset)?.parentOfType<LpNonterminal>(withSelf = true) != null
+    }
+
+    override fun invoke(project: Project, editor: Editor?, file: PsiFile?) {
+        if (editor == null || file == null || file.language != LpLanguage) return
+
+        file.findElementAt(editor.caretModel.primaryCaret.offset)?.parentOfType<LpNonterminal>(withSelf = true)?.apply {
+            setType(getContextAndResolveType(LpMacroArguments.identity(nonterminalName.nonterminalParams)))
+        }
+    }
+}

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, type = $type")
+
+        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,8 +1,45 @@
 package com.mdrobnak.lalrpop.psi
 
+import com.intellij.openapi.project.Project
 import com.intellij.psi.PsiElement
+import com.mdrobnak.lalrpop.psi.util.lalrpopTypeResolutionContext
+import org.rust.lang.core.macros.setContext
+import org.rust.lang.core.psi.RsPsiFactory
+import org.rust.lang.core.psi.RsTypeParameterList
+import org.rust.lang.core.psi.ext.RsElement
+import org.rust.lang.core.types.Substitution
+import org.rust.lang.core.types.infer.RsInferenceContext
+import org.rust.lang.core.types.toTypeSubst
+import org.rust.lang.core.types.ty.TyTypeParameter
+import org.rust.lang.core.types.ty.TyUnit
+import org.rust.lang.core.types.type
 
-data class NonterminalGenericArgument(val rustType: String, var name: String)
+data class LpMacroArgument(val rustType: String, val name: String)
+data class LpMacroArguments(val arguments: List<LpMacroArgument> = listOf()) : List<LpMacroArgument> by arguments {
+    fun getSubstitution(
+        params: RsTypeParameterList?,
+        project: Project,
+        inferenceContext: RsInferenceContext,
+        expandedElementContext: RsElement
+    ): Substitution =
+        params?.typeParameterList?.map { param ->
+            TyTypeParameter.named(param) to (arguments.find { arg -> arg.name == param.identifier.text }?.rustType?.let {
+                RsPsiFactory(project).createType(it).let { typeRef ->
+                    typeRef.setContext(expandedElementContext)
+
+                    inferenceContext.fullyResolve(typeRef.type)
+                }
+            } ?: TyUnit)
+        }.orEmpty().toMap().toTypeSubst()
+
+    companion object {
+        fun identity(params: LpNonterminalParams?): LpMacroArguments =
+            LpMacroArguments(params?.nonterminalParamList?.map {
+                val name = it.name!!
+                LpMacroArgument(name, name)
+            }.orEmpty())
+    }
+}
 
 data class LpTypeResolutionContext(
     val locationType: String = "usize",
@@ -25,8 +62,9 @@ interface LpResolveType : PsiElement {
      * And referenced with Nonterminal<A, B> in another symbol, the list of arguments should be the resolved types of
      * "A" and "B", in this order.
      */
-    fun resolveType(context: LpTypeResolutionContext, arguments: List<NonterminalGenericArgument>): String
+    fun resolveType(context: LpTypeResolutionContext, arguments: LpMacroArguments): String
 
-//    fun resolveType(arguments: List<NonterminalGenericArgument>): String =
-//        this.resolveType(this.containingFile.lalrpopTypeResolutionContext(), arguments)
 }
+
+fun LpResolveType.getContextAndResolveType(arguments: LpMacroArguments): String =
+    this.resolveType(this.containingFile.lalrpopTypeResolutionContext(), arguments)

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

@@ -1,6 +1,6 @@
 package com.mdrobnak.lalrpop.psi
 
 sealed class LpSelectedType {
-    data class WithName(val isMutable: Boolean, val name: String, val type: String): LpSelectedType()
+    data class WithName(val name: String, val type: String, val isMutable: Boolean = false): LpSelectedType()
     data class WithoutName(val type: String): LpSelectedType()
 }