re-using evaluators logic for resolving iterator type in HaxeResolveUtil

HaxeFoundation · Dec 10, 2023 · 53c81ea · 53c81ea
1 parent 1a4efc7
commit 53c81ea
Show file tree

Hide file tree

Showing 2 changed files with 34 additions and 46 deletions.
diff --git a/src/main/java/com/intellij/plugins/haxe/model/type/HaxeExpressionEvaluator.java b/src/main/java/com/intellij/plugins/haxe/model/type/HaxeExpressionEvaluator.java
@@ -782,30 +782,7 @@ else if (subelement instanceof HaxeForStatement forStatement) {
             }
           }
           else if (subelement instanceof HaxeIteratorkey || subelement instanceof HaxeIteratorValue) {
-            HaxeForStatement forStatement = PsiTreeUtil.getParentOfType(subelement, HaxeForStatement.class);
-            HaxeGenericResolver forResolver = HaxeGenericResolverUtil.generateResolverFromScopeParents(forStatement);
-
-            HaxeIterable iterable = forStatement.getIterable();
-            var keyValueIteratorType = HaxeTypeResolver.getPsiElementType(iterable, element, forResolver);
-
-            var iteratorType = keyValueIteratorType.getClassType();
-            if (iteratorType.isTypeDef()) {
-              iteratorType = iteratorType.fullyResolveTypeDefClass();
-            }
-            var iteratorTypeResolver = iteratorType.getGenericResolver();
-
-            HaxeMethodModel iteratorReturnType = (HaxeMethodModel)iteratorType.getHaxeClassModel().getMember("next", iteratorTypeResolver);
-            HaxeGenericResolver nextResolver = iteratorReturnType.getGenericResolver(null);
-            nextResolver.addAll(iteratorTypeResolver);
-
-            ResultHolder returnType = iteratorReturnType.getReturnType(nextResolver);
-            SpecificHaxeClassReference type = returnType.getClassType();
-            HaxeGenericResolver genericResolver = type.getGenericResolver();
-            if (subelement instanceof HaxeIteratorkey ) {
-              typeHolder = type.getHaxeClassModel().getMember("key", null).getResultType(genericResolver);
-            }else {
-              typeHolder = type.getHaxeClassModel().getMember("value", null).getResultType(genericResolver);
-            }
+            typeHolder = findIteratorType(element, subelement);
           }
 
 
@@ -1412,6 +1389,34 @@ else if (getter instanceof HaxeExternInterfaceDeclaration interfaceDeclaration)
     return SpecificHaxeClassReference.getUnknown(element).createHolder();
   }
 
+  public static ResultHolder findIteratorType(PsiElement reference, PsiElement iteratorElement) {
+    HaxeForStatement forStatement = PsiTreeUtil.getParentOfType(iteratorElement, HaxeForStatement.class);
+    HaxeGenericResolver forResolver = HaxeGenericResolverUtil.generateResolverFromScopeParents(forStatement);
+
+    HaxeIterable iterable = forStatement.getIterable();
+    var keyValueIteratorType = HaxeTypeResolver.getPsiElementType(iterable, reference, forResolver);
+
+    var iteratorType = keyValueIteratorType.getClassType();
+    if (iteratorType.isTypeDef()) {
+      iteratorType = iteratorType.fullyResolveTypeDefClass();
+    }
+    var iteratorTypeResolver = iteratorType.getGenericResolver();
+
+    HaxeMethodModel iteratorReturnType = (HaxeMethodModel)iteratorType.getHaxeClassModel().getMember("next", iteratorTypeResolver);
+    HaxeGenericResolver nextResolver = iteratorReturnType.getGenericResolver(null);
+    nextResolver.addAll(iteratorTypeResolver);
+
+    ResultHolder returnType = iteratorReturnType.getReturnType(nextResolver);
+    SpecificHaxeClassReference type = returnType.getClassType();
+    HaxeGenericResolver genericResolver = type.getGenericResolver();
+    if (iteratorElement instanceof HaxeIteratorkey ) {
+      return type.getHaxeClassModel().getMember("key", null).getResultType(genericResolver);
+    }else  if (iteratorElement instanceof  HaxeIteratorValue){
+      return type.getHaxeClassModel().getMember("value", null).getResultType(genericResolver);
+    }
+    return null;
+  }
+
   private static boolean isUntypedReturn(HaxeReturnStatement statement) {
     PsiElement child = statement.getFirstChild();
     while(child != null) {

diff --git a/src/main/java/com/intellij/plugins/haxe/util/HaxeResolveUtil.java b/src/main/java/com/intellij/plugins/haxe/util/HaxeResolveUtil.java
@@ -59,6 +59,7 @@
 
 import static com.intellij.plugins.haxe.lang.lexer.HaxeTokenTypeSets.DOC_COMMENT;
 import static com.intellij.plugins.haxe.model.type.HaxeExpressionEvaluator.evaluate;
+import static com.intellij.plugins.haxe.model.type.HaxeExpressionEvaluator.findIteratorType;
 import static com.intellij.plugins.haxe.util.HaxeDebugLogUtil.traceAs;
 
 /**
@@ -710,31 +711,13 @@ private static HaxeResolveResult getHaxeClassResolveResultInternal(@Nullable Psi
           return HaxeResolveResult.EMPTY;
         }
         final HaxeExpression expression = iterable.getExpression();
-        if (expression instanceof HaxeReference) {
-          final HaxeResolveResult resolveResult = ((HaxeReference)expression).resolveHaxeClass();
-          final HaxeClass resolveResultHaxeClass = resolveResult.getHaxeClass();
-          final HaxeGenericResolver resolver = resolveResult.getGenericResolver();
-          final HaxeGenericSpecialization resultSpecialization = resolveResult.getSpecialization();
+        if (expression instanceof HaxeReference reference) {
 
-          // find keyValue iterator type
-          HaxeResolveResult keyValueIteratorResult =
-            getResolveMethodReturnType(resolver, resolveResultHaxeClass, "keyValueIterator",
-                                       resultSpecialization.getInnerSpecialization(resolveResultHaxeClass));
-
-
-          HaxeClass iteratorClass = keyValueIteratorResult.getHaxeClass();
-          HaxeResolveResult iteratorResult =
-            getResolveMethodReturnType(resolver, iteratorClass, "next", keyValueIteratorResult.getSpecialization());
-
-          HaxeClass keyValueType = iteratorResult.getHaxeClass();
-
-          if (element instanceof HaxeIteratorkey) {
-           return  resolveFieldType(resolver, keyValueType, "key",  iteratorResult.getSpecialization());
+          ResultHolder type = findIteratorType(reference, element);
+          if (type!= null && !type.isUnknown() && type.isClassType()) {
+            return type.getClassType().asResolveResult();
           }
 
-          if (element instanceof  HaxeIteratorValue) {
-            return  resolveFieldType(resolver, keyValueType, "value", iteratorResult.getSpecialization());
-          }
         }
         return HaxeResolveResult.EMPTY;
     }