Tweaking capture variable type resolve to hopefully better handle arrays (#1233)

Author: m0rkeulv

src/main/java/com/intellij/plugins/haxe/lang/psi/HaxeResolver.java

@@ -1542,7 +1542,7 @@ private static List<PsiElement> tryResolveExtractedValue(HaxeReference reference
       boolean checkParent = true;
 
       HaxeSwitchStatement switchStatement = PsiTreeUtil.getParentOfType(reference, HaxeSwitchStatement.class);
-      PsiElement pathElement = buildExtractVarLiteralPath(parent, checkParent, objectPath, switchStatement);
+      PsiElement pathElement = buildExtractVarPath(parent, checkParent, objectPath, switchStatement);
 
       HaxeNamedComponent lastElement = null;
       Collections.reverse(objectPath);
@@ -1599,7 +1599,7 @@ private static List<PsiElement> tryResolveExtractedValue(HaxeReference reference
     return null;
   }
 
-  private static @Nullable PsiElement buildExtractVarLiteralPath(PsiElement extractedValue, boolean checkParent, Stack<Object> objectPath, HaxeSwitchStatement switchStatement) {
+  public static @Nullable PsiElement buildExtractVarPath(PsiElement extractedValue, boolean checkParent, Stack<Object> objectPath, HaxeSwitchStatement switchStatement) {
     PsiElement pastParent = extractedValue;
     PsiElement valueParent = extractedValue.getParent();
     PsiElement pathElement = null;
@@ -1609,6 +1609,10 @@ private static List<PsiElement> tryResolveExtractedValue(HaxeReference reference
       if (valueParent instanceof HaxeEnumObjectLiteralElement objectLiteralElement) {
         objectPath.add(extractObjectLiteralName(objectLiteralElement));
 
+      } else if (valueParent instanceof HaxeEnumExtractArrayLiteral arrayLiteral) {
+        List<PsiElement> caseExpressionList = Arrays.asList(arrayLiteral.getChildren());
+        objectPath.add(caseExpressionList.indexOf(pastParent));
+
       } else if (valueParent instanceof HaxeSwitchCaseExprArray caseExprArray) {
         List<PsiElement> caseExpressionList = Arrays.asList(caseExprArray.getChildren());
         objectPath.add(caseExpressionList.indexOf(pastParent));

src/main/java/com/intellij/plugins/haxe/lang/psi/impl/HaxeStringLiteralImpl.java

@@ -22,13 +22,14 @@
 
 import com.intellij.lang.ASTNode;
 import com.intellij.openapi.util.TextRange;
+import com.intellij.plugins.haxe.lang.psi.HaxeLiteralExpression;
 import com.intellij.psi.*;
 import com.intellij.psi.impl.source.resolve.reference.ReferenceProvidersRegistry;
 import com.intellij.psi.impl.source.tree.LeafElement;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
-abstract public class HaxeStringLiteralImpl extends HaxeReferenceImpl implements PsiLanguageInjectionHost, PsiLiteralValue {
+abstract public class HaxeStringLiteralImpl extends HaxeReferenceImpl implements PsiLanguageInjectionHost, PsiLiteralValue, HaxeLiteralExpression {
 
   public HaxeStringLiteralImpl(ASTNode node) {
     super(node);

src/main/java/com/intellij/plugins/haxe/model/evaluator/HaxeExpressionEvaluator.java

@@ -241,7 +241,7 @@ static ResultHolder _handle(@Nullable final PsiElement element,
     }
 
     if (element instanceof HaxeEnumExtractedValueReference extractedValue) {
-      return handleEnumExtractedValue(extractedValue, resolver);
+      return handleExtractedValue(extractedValue, resolver);
     }
 
 

src/main/java/com/intellij/plugins/haxe/model/evaluator/HaxeExpressionEvaluatorHandlers.java

@@ -39,6 +39,7 @@
 
 import static com.intellij.plugins.haxe.lang.lexer.HaxeTokenTypeSets.ONLY_COMMENTS;
 import static com.intellij.plugins.haxe.lang.lexer.HaxeTokenTypes.KUNTYPED;
+import static com.intellij.plugins.haxe.lang.psi.HaxeResolver.buildExtractVarPath;
 import static com.intellij.plugins.haxe.lang.psi.impl.HaxeReferenceImpl.getLiteralClassName;
 import static com.intellij.plugins.haxe.lang.psi.impl.HaxeReferenceImpl.tryToFindTypeFromCallExpression;
 import static com.intellij.plugins.haxe.lang.psi.impl.HaxeReferenceUtil.isStaticExtension;
@@ -856,28 +857,88 @@ public void run() {
   }
 
 
-  static ResultHolder handleEnumExtractedValue(@NotNull HaxeEnumExtractedValueReference extractedValue, @NotNull HaxeGenericResolver resolver) {
-    HaxeExtractorMatchExpression matchExpression = PsiTreeUtil.getParentOfType(extractedValue, HaxeExtractorMatchExpression.class, true, HaxeEnumArgumentExtractor.class);
+  static ResultHolder handleExtractedValue(@NotNull HaxeEnumExtractedValueReference extractedValueRef, @NotNull HaxeGenericResolver resolver) {
+    HaxeExtractorMatchExpression matchExpression = PsiTreeUtil.getParentOfType(extractedValueRef, HaxeExtractorMatchExpression.class, true, HaxeEnumArgumentExtractor.class);
     if (matchExpression != null) {
       HaxeSwitchCaseExpr match = matchExpression.getMatch();
-      if (PsiTreeUtil.isAncestor(match, extractedValue, false)) {
+      if (PsiTreeUtil.isAncestor(match, extractedValueRef, false)) {
         return evaluate(matchExpression.getExtractorExpression(), resolver).result;
       }
     }
 
-    HaxeEnumArgumentExtractor extractor = PsiTreeUtil.getParentOfType(extractedValue, HaxeEnumArgumentExtractor.class, true, HaxeEnumArgumentExtractor.class);
+    HaxeEnumArgumentExtractor extractor = PsiTreeUtil.getParentOfType(extractedValueRef, HaxeEnumArgumentExtractor.class, true, HaxeEnumArgumentExtractor.class);
     if (extractor != null) {
       HaxeEnumExtractorModel extractorModel = (HaxeEnumExtractorModel)extractor.getModel();
-      return extractorModel.resolveExtractedValueType(extractedValue);
+      return extractorModel.resolveExtractedValueType(extractedValueRef);
     }
 
-    PsiElement resolve = extractedValue.resolve();
+    PsiElement resolve = extractedValueRef.resolve();
     if(resolve != null) {
-      return evaluate(resolve).result;
+
+      ResultHolder result = evaluate(resolve).result;
+
+        if (extractedValueRef.getParent() instanceof HaxeEnumExtractedValue extractedValue) {
+          //TODO: mlo - We should probably try to  redo how switch extractors are parsed in the BNF
+          // to get more consistency between ExprArray and ArrayLiteral (we might not need both)
+          if (extractedValue.getParent() instanceof HaxeEnumExtractArrayLiteral
+                  || extractedValue.getParent() instanceof HaxeSwitchCaseExprArray) {
+
+            Stack<Object> objectPath = new Stack<>();
+
+            HaxeSwitchStatement switchStatement = PsiTreeUtil.getParentOfType(extractedValue, HaxeSwitchStatement.class);
+            buildExtractVarPath(extractedValue, true, objectPath, switchStatement);
+
+            Collections.reverse(objectPath);
+            // workaround for switch on literal array  (ex. switch (["mixed Types", 1, false]{...}))
+            // as we are not going to look for generics in these  situations to determine type.
+            boolean isLiteralInSwitchExpression = isLiteralInSwitchExpression(resolve, extractedValueRef);
+            int offset = isLiteralInSwitchExpression ? 1 : 0;
+
+            ResultHolder typePointer = result;
+              for (int i = offset; i < objectPath.size(); i++) {
+                  Object o = objectPath.get(i);
+                  if (typePointer == null || typePointer.isUnknown()) break;
+
+                  if (o instanceof Integer index) {
+                      typePointer = typeFromArray(typePointer, index);
+                    //TODO do we need support for objects here ?
+//                  } else if (o instanceof String name) {
+                    // typePointer = typeFromObjectName(typePointer,name);
+                  }
+              }
+            return typePointer != null ? typePointer : createUnknown(extractedValueRef);
+          }
+        }
+      return result;
     }
-    return createUnknown(extractedValue);
+    return createUnknown(extractedValueRef);
   }
 
+  private static boolean isLiteralInSwitchExpression(PsiElement resolve, @NotNull HaxeEnumExtractedValueReference extractedValueRef) {
+    HaxeLiteralExpression isLiteral = PsiTreeUtil.getParentOfType(resolve, HaxeLiteralExpression.class, false, HaxeSwitchStatement.class);
+    HaxeArrayLiteral isInsideArrayLiteral = PsiTreeUtil.getParentOfType(resolve, HaxeArrayLiteral.class, false, HaxeSwitchStatement.class);
+    HaxeSwitchStatement resolvedSwitchParent = PsiTreeUtil.getParentOfType(resolve, HaxeSwitchStatement.class);
+    if(resolvedSwitchParent == null) return false;
+    boolean isInSameSwitch = PsiTreeUtil.isAncestor(resolvedSwitchParent.getSwitchBlock(), extractedValueRef, false);
+    return isInSameSwitch && (isLiteral != null || isInsideArrayLiteral != null);
+  }
+
+  private static ResultHolder typeFromArray(ResultHolder typePointer, Integer index) {
+
+      SpecificHaxeClassReference classType = typePointer.getClassType();
+      if(classType != null){
+        if(classType.fullyResolveTypeDefAndUnwrapNullTypeReference() instanceof SpecificHaxeClassReference classReference) {
+          if(classReference.isArray()) {
+            @NotNull ResultHolder[] specifics = classReference.getSpecifics();
+            if(specifics.length == 1) return specifics[0];
+//          }else {
+            // TODO considder checking array access in other types for expected return value
+            // - verify if its supported first
+          }
+        }
+      }
+      return null;
+    }
 
 
   static ResultHolder handleForStatement(

src/test/java/com/intellij/plugins/haxe/ide/HaxeSemanticAnnotatorTest.java

@@ -1041,6 +1041,11 @@ public void testSwitchPatternMatching() throws Throwable {
     myFixture.enableInspections(HaxeUnresolvedSymbolInspection.class);
     doTestNoFixWithWarnings();
   }
+  @Test
+  public void testSwitchPatternMatchingArrays() throws Throwable {
+    myFixture.enableInspections(HaxeUnresolvedSymbolInspection.class);
+    doTestNoFixWithWarnings();
+  }
 
   @Test
   public void testSwitchStatements() throws Throwable {

src/test/resources/testData/annotation.semantic/SwitchPatternMatchingArrays.hx

@@ -0,0 +1,28 @@
+class PatternMachingTest {
+
+    public function test1DArray() {
+        var myArray = [1, 2, 3];
+        var s = switch (myArray) {
+            case [a, b, c]: a + b + c;
+            case a: 0;
+        }
+        trace(s);
+    }
+
+    public function test2DArray() {
+        var myArray = [[1, 2, 3], [4, 5, 6]];
+        var s = switch (myArray) {
+            case [[a, b, _],[_,c,_]]: a + b + c;
+            case a: 0;
+        }
+        trace(s);
+    }
+    public function testArrayInObject() {
+        var myObjectWithArray = {arr:[[1,2], [3]]};
+        var s = switch (myObjectWithArray) {
+            case {arr:[[a, b], [c]] }: a + b + c;
+            case a: 0;
+        }
+    }
+
+}