Resolve top-level Enums, Structs, Errors, and user-defined types from imports w/ various fixes (intellij-solidity#324)

Author: apottere

src/main/kotlin/me/serce/solidity/ide/annotation/annotator.kt

@@ -49,6 +49,7 @@ class SolidityAnnotator : Annotator {
       is SolStructDefinition -> element.identifier?.let { applyColor(holder, it, SolColor.STRUCT_NAME) }
       is SolEnumDefinition -> element.identifier?.let { applyColor(holder, it, SolColor.ENUM_NAME) }
       is SolEventDefinition -> element.identifier?.let { applyColor(holder, it, SolColor.EVENT_NAME) }
+      is SolUserDefinedValueTypeDefinition -> element.identifier?.let { applyColor(holder, it, SolColor.USER_DEFINED_VALUE_TYPE) }
       is SolConstantVariableDeclaration -> applyColor(holder, element.identifier, SolColor.CONSTANT)
       is SolStateVariableDeclaration -> {
         if (element.mutationModifier?.textMatches("constant") == true) {
@@ -74,7 +75,13 @@ class SolidityAnnotator : Annotator {
         else -> when(SolResolver.resolveTypeNameUsingImports(element).firstOrNull()) {
           is SolErrorDefinition -> applyColor(holder, element.referenceNameElement, SolColor.ERROR_NAME)
           is SolEventDefinition -> applyColor(holder, element.referenceNameElement, SolColor.EVENT_NAME)
-          else -> applyColor(holder, element.referenceNameElement, SolColor.FUNCTION_CALL)
+          else -> element.firstChild.let {
+            if (it is SolPrimaryExpression && SolResolver.resolveTypeNameUsingImports(element.firstChild).filterIsInstance<SolStructDefinition>().isNotEmpty()) {
+              applyColor(holder, element.referenceNameElement, SolColor.STRUCT_NAME)
+            } else {
+              applyColor(holder, element.referenceNameElement, SolColor.FUNCTION_CALL)
+            }
+          }
         }
       }
     }

src/main/kotlin/me/serce/solidity/lang/resolve/engine.kt

@@ -24,11 +24,11 @@ object SolResolver {
       val result = if (element is SolFunctionCallElement) {
         resolveError(element) +
           resolveEvent(element) +
-          resolveContractUsingImports(element, element.containingFile, true) +
+          resolveContract(element) +
           resolveEnum(element) +
           resolveUserDefinedValueType(element)
       } else {
-        resolveContractUsingImports(element, element.containingFile, true) +
+        resolveContract(element) +
           resolveEnum(element) +
           resolveStruct(element) +
           resolveUserDefinedValueType(element)
@@ -39,17 +39,18 @@ object SolResolver {
   /**
    * @param withAliases aliases are not recursive, so count them only at the first level of recursion
    */
-  private fun resolveContractUsingImports(
+  private fun <T: SolNamedElement>resolveUsingImports(
+    target: Class<T>,
     element: PsiElement,
     file: PsiFile,
-    withAliases: Boolean
-  ): Set<SolContractDefinition> =
-    RecursionManager.doPreventingRecursion(ResolveContractKey(element.nameOrText, file), true) {
+    withAliases: Boolean,
+  ): Set<T> =
+    RecursionManager.doPreventingRecursion(ResolveUsingImportsKey(element.nameOrText, file), true) {
       if (element is SolUserDefinedTypeName && element.findIdentifiers().size > 1) {
         emptySet()
       } else {
         val inFile = file.children
-          .filterIsInstance<SolContractDefinition>()
+          .filterIsInstance(target)
           .filter { it.name == element.nameOrText }
 
         val resolvedViaAlias = when (withAliases) {
@@ -64,36 +65,38 @@ object SolResolver {
                   }
                 }
             }.flatMap { (alias, resolvedFile) ->
-              resolveContractUsingImports(alias, resolvedFile.containingFile, false)
+              resolveUsingImports(target, alias, resolvedFile.containingFile, false)
             }
           else -> emptyList()
         }
 
         val imported = file.children
           .filterIsInstance<SolImportDirective>()
           .mapNotNull { nullIfError { it.importPath?.reference?.resolve()?.containingFile } }
-          .flatMap { resolveContractUsingImports(element, it, false) }
+          .flatMap { resolveUsingImports(target, element, it, false) }
 
         (inFile + resolvedViaAlias + imported).toSet()
       }
     } ?: emptySet()
 
+  private fun resolveContract(element: PsiElement): Set<SolContractDefinition> =
+    resolveUsingImports(SolContractDefinition::class.java, element, element.containingFile, true)
   private fun resolveEnum(element: PsiElement): Set<SolNamedElement> =
-    resolveInFile<SolEnumDefinition>(element) + resolveInnerType<SolEnumDefinition>(element) { it.enumDefinitionList }
+    resolveInnerType<SolEnumDefinition>(element) { it.enumDefinitionList } + resolveUsingImports(SolEnumDefinition::class.java, element, element.containingFile, true)
 
   private fun resolveStruct(element: PsiElement): Set<SolNamedElement> =
-    resolveInFile<SolStructDefinition>(element) + resolveInnerType<SolStructDefinition>(element) { it.structDefinitionList }
+    resolveInnerType<SolStructDefinition>(element) { it.structDefinitionList } + resolveUsingImports(SolStructDefinition::class.java, element, element.containingFile, true)
 
   private fun resolveUserDefinedValueType(element: PsiElement): Set<SolNamedElement> =
-    resolveInFile<SolUserDefinedValueTypeDefinition>(element) + resolveInnerType<SolUserDefinedValueTypeDefinition>(
+    resolveInnerType<SolUserDefinedValueTypeDefinition>(
       element,
-      { it.userDefinedValueTypeDefinitionList })
+      { it.userDefinedValueTypeDefinitionList }) + resolveUsingImports(SolUserDefinedValueTypeDefinition::class.java, element, element.containingFile, true)
 
   private fun resolveEvent(element: PsiElement): Set<SolNamedElement> =
     resolveInnerType<SolEventDefinition>(element) { it.eventDefinitionList }
 
   private fun resolveError(element: PsiElement): Set<SolNamedElement> =
-    resolveInnerType<SolErrorDefinition>(element) { it.errorDefinitionList }
+    resolveInnerType<SolErrorDefinition>(element) { it.errorDefinitionList } + resolveUsingImports(SolErrorDefinition::class.java, element, element.containingFile, true)
 
   private inline fun <reified T : SolNamedElement> resolveInFile(element: PsiElement) : Set<T> {
     return element.parentOfType<SolidityFile>()
@@ -343,7 +346,7 @@ object SolResolver {
   }
 }
 
-data class ResolveContractKey(val name: String?, val file: PsiFile)
+data class ResolveUsingImportsKey(val name: String?, val file: PsiFile)
 
 private fun <T> Sequence<T>.takeWhileInclusive(pred: (T) -> Boolean): Sequence<T> {
   var shouldContinue = true

src/main/kotlin/me/serce/solidity/lang/resolve/ref/refs.kt

@@ -93,6 +93,15 @@ class SolFunctionCallReference(element: SolFunctionCallExpression) : SolReferenc
   }
 
   fun resolveFunctionCall(): Collection<SolCallable> {
+    if (element.parent is SolRevertStatement) {
+      return SolResolver.resolveTypeNameUsingImports(element).filterIsInstance<SolErrorDefinition>()
+    }
+    if (element.firstChild is SolPrimaryExpression) {
+      val structs = SolResolver.resolveTypeNameUsingImports(element.firstChild).filterIsInstance<SolStructDefinition>()
+      if (structs.isNotEmpty()) {
+        return structs
+      }
+    }
     val resolved: Collection<SolCallable> = when (val expr = element.expression) {
       is SolPrimaryExpression -> {
         val regular = expr.varLiteral?.let { SolResolver.resolveVarLiteral(it) }

src/main/kotlin/me/serce/solidity/lang/types/inference.kt

@@ -73,6 +73,7 @@ private fun getSolTypeFromUserDefinedTypeName(type: SolUserDefinedTypeName): Sol
         is SolContractDefinition -> SolContract(it)
         is SolStructDefinition -> SolStruct(it)
         is SolEnumDefinition -> SolEnum(it)
+        is SolUserDefinedValueTypeDefinition -> getSolType(it.elementaryTypeName)
         else -> null
       }
     }