allow overloading the index operator

de.peeeq.wurstscript/src/main/antlr/de/peeeq/wurstscript/antlr/Wurst.g4

@@ -25,7 +25,7 @@ jassGlobalsBlock:
 				;
 
 jassGlobalDecl:
-				  constant='constant'? typeExpr name=ID ('=' initial=expr)? NL
+				  constant='constant'? typeExpr name=(ID|CONTINUE|SKIP_|BREAK) ('=' initial=expr)? NL
 			  ;
 
 jassFuncDef:
@@ -35,7 +35,7 @@ jassFuncDef:
 			   'endfunction' NL
 		   ;
 
-jassLocal: 'local' typeExpr name=ID ('=' initial=expr)? NL;
+jassLocal: 'local' typeExpr name=(ID|CONTINUE|SKIP_|BREAK) ('=' initial=expr)? NL;
 
 jassStatements: stmts+=jassStatement*;
 
@@ -224,7 +224,7 @@ funcSignature:
 formalParameters: '(' (params+=formalParameter (',' params+=formalParameter)*)? ')';
 
 formalParameter:
-				 vararg=VARARG? typeExpr name=ID
+				 vararg=VARARG? typeExpr name=(ID|CONTINUE|SKIP_|BREAK)
 			   ;
 
 typeExpr:
@@ -238,7 +238,7 @@ typeExpr:
 varDef:
 		  modifiersWithDoc 
 		  ('var'|constant='constant' varType=typeExpr?|constant='let'|varType=typeExpr)
-		  name=ID ('=' variableInit)? NL
+		  name=(ID|CONTINUE|SKIP_|BREAK) ('=' variableInit)? NL
 	  ;		  
 
 variableInit: (arrayInit | initial=expr);
@@ -312,12 +312,12 @@ stmtWhile:
 
 localVarDef:
 		  (var='var'|let='let'|type=typeExpr)
-		  name=ID ('=' variableInit)?
+		  name=(ID|CONTINUE|SKIP_|BREAK) ('=' variableInit)?
 	  ;
 
 
 localVarDefInline:
-					 typeExpr? name=ID
+					 typeExpr? name=(ID|CONTINUE|SKIP_|BREAK)
 				 ;
 
 stmtSet:
@@ -335,12 +335,12 @@ exprAssignable:
 			  ;
 
 exprMemberVar: 
-				 expr dots=('.'|'..') varname=ID? indexes?
+				 expr dots=('.'|'..') varname=(ID|CONTINUE|SKIP_|BREAK)? indexes?
 			 ;
 
 
 exprVarAccess:
-				 varname=(ID|IT) indexes?
+				 varname=(ID|CONTINUE|SKIP_|BREAK|IT) indexes?
 			 ;
 
 
@@ -354,7 +354,7 @@ expr:
 	  | left=expr 'castTo' castToType=typeExpr
 	  | left=expr 'instanceof' instaneofType=typeExpr
 	  | receiver=expr dotsCall=('.'|'..') funcName=ID? typeArgs argumentList
-	  | receiver=expr dotsVar=('.'|'..') varName=ID? indexes?
+	  | receiver=expr dotsVar=('.'|'..') varName=(ID|CONTINUE|SKIP_|BREAK)? indexes?
       | left=expr op=('*'|'/'|'%'|'div'|'mod') right=expr
 	  | op='-' right=expr // TODO move unary minus one up to be compatible with Java etc.
 		                  // currently it is here to be backwards compatible with the old wurst parser
@@ -375,7 +375,7 @@ exprPrimary:
 	  | exprClosure
 	  | exprStatementsBlock
 	  | exprDestroy
-      | varname=(ID|IT) indexes?
+      | varname=(ID|CONTINUE|SKIP_|BREAK|IT) indexes?
       | atom=(INT
       | REAL
 	  | STRING
@@ -412,7 +412,7 @@ shortFormalParameters:
     | /* empty */
     ;
 
-shortFormalParameter: typeExpr? name=ID;
+shortFormalParameter: typeExpr? name=(ID|CONTINUE|SKIP_|BREAK);
 
 typeParams: ('<' (params+=typeParam (',' params+=typeParam)*)? '>')?;
 
@@ -447,6 +447,7 @@ exprList : exprs+=expr (',' exprs+=expr)*;
 
 
 
+
 nativeType: 'nativetype' name=ID ('extends' extended=ID)? NL;
 initBlock: 'init' NL statementsBlock; 
 nativeDef: modifiersWithDoc 'native' funcSignature NL; 

de.peeeq.wurstscript/src/main/java/de/peeeq/wurstio/languageserver/requests/GetDefinition.java

@@ -6,9 +6,12 @@
 import de.peeeq.wurstio.languageserver.WFile;
 import de.peeeq.wurstscript.WLogger;
 import de.peeeq.wurstscript.ast.*;
+import de.peeeq.wurstscript.attributes.AttrFuncDef;
 import de.peeeq.wurstscript.attributes.CofigOverridePackages;
+import de.peeeq.wurstscript.attributes.names.FuncLink;
 import de.peeeq.wurstscript.attributes.names.NameLink;
 import de.peeeq.wurstscript.parser.WPos;
+import de.peeeq.wurstscript.types.WurstTypeArray;
 import de.peeeq.wurstscript.types.WurstType;
 import de.peeeq.wurstscript.types.WurstTypeNamedScope;
 import de.peeeq.wurstscript.utils.Utils;
@@ -77,6 +80,10 @@ private List<? extends Location> execute2(ModelManager modelManager) {
                 return linkTo(originalDecl);
             }
         }
+        FunctionDefinition indexOpDecl = getIndexOperatorDeclarationAtPos(e);
+        if (indexOpDecl != null) {
+            return linkTo(indexOpDecl);
+        }
         if (e instanceof FuncRef) {
             FuncRef funcRef = (FuncRef) e;
             FunctionDefinition decl = funcRef.attrFuncDef();
@@ -117,6 +124,37 @@ private List<? extends Location> execute2(ModelManager modelManager) {
         return Collections.emptyList();
     }
 
+    private FunctionDefinition getIndexOperatorDeclarationAtPos(Element e) {
+        Element target = e;
+        if (target instanceof Indexes && target.getParent() instanceof NameRef) {
+            target = target.getParent();
+        }
+        if (!(target instanceof NameRef) || !(target instanceof AstElementWithIndexes)) {
+            return null;
+        }
+
+        NameRef nr = (NameRef) target;
+        AstElementWithIndexes withIndexes = (AstElementWithIndexes) target;
+        if (withIndexes.getIndexes().size() != 1) {
+            return null;
+        }
+        NameLink link = nr.attrNameLink();
+        if (link == null || link.getTyp() instanceof WurstTypeArray) {
+            return null;
+        }
+
+        WurstType receiverType = link.getTyp();
+        WurstType indexType = withIndexes.getIndexes().get(0).attrTyp();
+        if (nr.getParent() instanceof StmtSet && ((StmtSet) nr.getParent()).getUpdatedExpr() == nr) {
+            StmtSet set = (StmtSet) nr.getParent();
+            FuncLink f = AttrFuncDef.getIndexSetOperator(nr, receiverType, indexType, set.getRight().attrTyp());
+            return f == null ? null : f.getDef();
+        }
+
+        FuncLink f = AttrFuncDef.getIndexGetOperator(nr, receiverType, indexType);
+        return f == null ? null : f.getDef();
+    }
+
     private List<? extends Location> typeDefinitionFor(Element e) {
         if (e instanceof TypeExpr) {
             TypeExpr typeExpr = (TypeExpr) e;

de.peeeq.wurstscript/src/main/java/de/peeeq/wurstscript/attributes/AttrExprType.java

@@ -101,9 +101,23 @@ public static WurstType calculate(ExprVarArrayAccess term) {
         WurstType varDefType = varDef.getTyp();
         if (varDefType instanceof WurstTypeArray) {
             return ((WurstTypeArray) varDefType).getBaseType();
-        } else {
-            term.addError("Variable " + varDef.getName() + " is of type " + varDefType + ", should be an array variable.");
         }
+        if (term.getIndexes().size() == 1) {
+            WurstType indexType = term.getIndexes().get(0).attrTyp();
+            FuncLink getOverload = AttrFuncDef.getIndexGetOperator(term, varDefType, indexType);
+            if (getOverload != null) {
+                return getOverload.getReturnType();
+            }
+            if (isWriteAccess(term)) {
+                FuncLink setOverload = AttrFuncDef.getIndexSetOperatorByIndex(term, varDefType, indexType);
+                if (setOverload != null) {
+                    return setOverload.getParameterType(1);
+                }
+                return WurstTypeUnknown.instance();
+            }
+        }
+        term.addError("Variable " + varDef.getName() + " is of type " + varDefType +
+                ", should be an array variable or define operator overloading function " + AttrFuncDef.overloadingIndexGet + ".");
         return WurstTypeUnknown.instance();
     }
 
@@ -412,7 +426,21 @@ public static WurstType calculate(ExprMemberArrayVarDot term) {
             WurstTypeArray ar = (WurstTypeArray) typ;
             return ar.getBaseType();
         }
-        term.addError("Variable " + term.getVarName() + " is not an array.");
+        if (term.getIndexes().size() == 1) {
+            WurstType indexType = term.getIndexes().get(0).attrTyp();
+            FuncLink getOverload = AttrFuncDef.getIndexGetOperator(term, typ, indexType);
+            if (getOverload != null) {
+                return getOverload.getReturnType();
+            }
+            if (isWriteAccess(term)) {
+                FuncLink setOverload = AttrFuncDef.getIndexSetOperatorByIndex(term, typ, indexType);
+                if (setOverload != null) {
+                    return setOverload.getParameterType(1);
+                }
+                return WurstTypeUnknown.instance();
+            }
+        }
+        term.addError("Variable " + term.getVarName() + " is not an array and has no " + AttrFuncDef.overloadingIndexGet + " overload.");
         return typ;
     }
 
@@ -592,4 +620,12 @@ public static WurstType calculate(ExprArrayLength exprArrayLength) {
         exprArrayLength.addError(".length is only valid on arrays.");
         return de.peeeq.wurstscript.types.WurstTypeUnknown.instance();
     }
+
+    private static boolean isWriteAccess(NameRef node) {
+        if (node.getParent() instanceof StmtSet) {
+            StmtSet stmtSet = (StmtSet) node.getParent();
+            return stmtSet.getUpdatedExpr() == node;
+        }
+        return false;
+    }
 }

de.peeeq.wurstscript/src/main/java/de/peeeq/wurstscript/attributes/AttrFuncDef.java

@@ -34,6 +34,8 @@ public class AttrFuncDef {
     public final static String overloadingMinus = "op_minus";
     public final static String overloadingMult = "op_mult";
     public final static String overloadingDiv = "op_divReal";
+    public final static String overloadingIndexGet = "op_index";
+    public final static String overloadingIndexSet = "op_indexAssign";
 
     public static FuncLink calculate(final ExprFuncRef node) {
 
@@ -206,6 +208,76 @@ private static boolean isConstructorThisCall(ExprFunctionCall node) {
         return searchMemberFunc(left, left.attrTyp(), funcName, Collections.singletonList(right.attrTyp()));
     }
 
+    public static @Nullable FuncLink getIndexGetOperator(Expr node, WurstType receiverType, WurstType indexType) {
+        List<WurstType> argTypes = Collections.singletonList(indexType);
+        FuncLink candidate = searchMemberFunc(node, receiverType, overloadingIndexGet, argTypes);
+        if (candidate == null || !matchesArguments(node, candidate, argTypes)) {
+            return null;
+        }
+        return candidate;
+    }
+
+    public static @Nullable FuncLink getIndexSetOperator(Expr node, WurstType receiverType, WurstType indexType, WurstType valueType) {
+        List<WurstType> argTypes = Lists.newArrayList(indexType, valueType);
+        FuncLink candidate = searchMemberFunc(node, receiverType, overloadingIndexSet, argTypes);
+        if (candidate == null || !matchesArguments(node, candidate, argTypes)) {
+            return null;
+        }
+        return candidate;
+    }
+
+    public static @Nullable FuncLink getIndexSetOperatorByIndex(Expr node, WurstType receiverType, WurstType indexType) {
+        Collection<FuncLink> funcs1 = node.lookupMemberFuncs(receiverType, overloadingIndexSet);
+        if (funcs1.isEmpty()) {
+            return null;
+        }
+        List<FuncLink> funcs = filterInvisible(overloadingIndexSet, node, funcs1);
+        funcs = filterByReceiverType(node, overloadingIndexSet, funcs);
+        List<FuncLink> byParamCount = Lists.newArrayList();
+        for (FuncLink f : funcs) {
+            if (f.getParameterTypes().size() == 2) {
+                byParamCount.add(f);
+            }
+        }
+        if (byParamCount.isEmpty()) {
+            return null;
+        }
+        List<FuncLink> byIndexType = Lists.newArrayList();
+        for (FuncLink f : byParamCount) {
+            VariableBinding mapping = f.getVariableBinding();
+            WurstType expectedIndexType = f.getParameterType(0);
+            VariableBinding m2 = indexType.matchAgainstSupertype(expectedIndexType, node, mapping, VariablePosition.RIGHT);
+            if (m2 != null) {
+                byIndexType.add(f);
+            }
+        }
+        if (byIndexType.isEmpty()) {
+            return null;
+        }
+        if (byIndexType.size() == 1) {
+            return byIndexType.get(0);
+        }
+        // ambiguous write-only expected type context: pick deterministic first
+        return byIndexType.get(0);
+    }
+
+    private static boolean matchesArguments(Element node, FuncLink f, List<WurstType> argumentTypes) {
+        if (f.getParameterTypes().size() != argumentTypes.size()) {
+            return false;
+        }
+        VariableBinding mapping = f.getVariableBinding();
+        for (int i = 0; i < argumentTypes.size(); i++) {
+            WurstType at = argumentTypes.get(i);
+            WurstType pt = f.getParameterType(i);
+            VariableBinding m2 = at.matchAgainstSupertype(pt, node, mapping, VariablePosition.RIGHT);
+            if (m2 == null) {
+                return false;
+            }
+            mapping = m2;
+        }
+        return true;
+    }
+
 
     /**
      * checks if operator is a native operator like for 1+2