multi-index array expressions and related fixes

compiler/backend/src/org/jetbrains/jet/codegen/FunctionCodegen.java

@@ -81,9 +81,10 @@ public class FunctionCodegen {
                 ExpressionCodegen codegen = new ExpressionCodegen(mv, frameMap, jvmSignature.getReturnType(), context, state);
 
                 Type[] argTypes = jvmSignature.getArgumentTypes();
+                int add = functionDescriptor.getReceiverParameter().exists() ? state.getTypeMapper().mapType(functionDescriptor.getReceiverParameter().getType()).getSize() : 0;
                 for (int i = 0; i < paramDescrs.size(); i++) {
                     ValueParameterDescriptor parameter = paramDescrs.get(i);
-                    frameMap.enter(parameter, argTypes[i].getSize());
+                    frameMap.enter(parameter, argTypes[i+add].getSize());
                 }
 
                 for (final TypeParameterDescriptor typeParameterDescriptor : typeParameters) {

compiler/backend/src/org/jetbrains/jet/codegen/JetTypeMapper.java

@@ -547,6 +547,9 @@ public class JetTypeMapper {
     }
 
     public CallableMethod mapToCallableMethod(FunctionDescriptor functionDescriptor, OwnerKind kind) {
+        if(functionDescriptor == null)
+            return null;
+
         final DeclarationDescriptor functionParent = functionDescriptor.getContainingDeclaration();
         final List<Type> valueParameterTypes = new ArrayList<Type>();
         Method descriptor = mapSignature(functionDescriptor.getOriginal(), valueParameterTypes, kind);

compiler/backend/src/org/jetbrains/jet/codegen/StackValue.java

@@ -38,7 +38,7 @@ public abstract class StackValue {
         throw new UnsupportedOperationException("cannot store to value " + this);
     }
 
-    public void dupReceiver(InstructionAdapter v, int below) {
+    public void dupReceiver(InstructionAdapter v) {
     }
 
     public void condJump(Label label, boolean jumpIfFalse, InstructionAdapter v) {
@@ -84,8 +84,8 @@ public abstract class StackValue {
         return new ArrayElement(type, unbox);
     }
 
-    public static StackValue collectionElement(Type type, CallableMethod getter, CallableMethod setter) {
-        return new CollectionElement(type, getter, setter);
+    public static StackValue collectionElement(Type type, CallableMethod getter, CallableMethod setter, FrameMap frame) {
+        return new CollectionElement(type, getter, setter, frame);
     }
 
     public static StackValue field(Type type, String owner, String name, boolean isStatic) {
@@ -456,24 +456,21 @@ public abstract class StackValue {
         }
 
         @Override
-        public void dupReceiver(InstructionAdapter v, int below) {
-            if (below == 1) {
-                v.dup2X1();
-            }
-            else {
-                v.dup2();   // array and index
-            }
+        public void dupReceiver(InstructionAdapter v) {
+            v.dup2();   // array and index
         }
     }
 
     private static class CollectionElement extends StackValue {
         private final CallableMethod getter;
         private final CallableMethod setter;
+        private final FrameMap frame;
 
-        public CollectionElement(Type type, CallableMethod getter, CallableMethod setter) {
+        public CollectionElement(Type type, CallableMethod getter, CallableMethod setter, FrameMap frame) {
             super(type);
             this.getter = getter;
             this.setter = setter;
+            this.frame = frame;
         }
 
         @Override
@@ -494,14 +491,71 @@ public abstract class StackValue {
         }
 
         @Override
-        public void dupReceiver(InstructionAdapter v, int below) {
-            if (below == 1) {
-                v.dup2X1();
+        public void dupReceiver(InstructionAdapter v) {
+            int size = calcSize();
+            
+            if(size == 2) {
+                v.dup2();   // collection and index
             }
             else {
-                v.dup2();   // collection and index
+                Method signature = getter.getSignature();
+                Type[] argumentTypes = signature.getArgumentTypes();
+
+                int firstIndex = frame.enterTemp();
+                int lastIndex = firstIndex;
+                frame.leaveTemp();
+                
+                for(int i = argumentTypes.length-1; i >= 0; i--) {
+                    int sz = argumentTypes[i].getSize();
+                    frame.enterTemp(sz);
+                    lastIndex += sz;
+                    v.store(lastIndex-sz, argumentTypes[i]);
+                }
+
+                if(getter.getInvokeOpcode() != Opcodes.INVOKESTATIC) {
+                    frame.enterTemp();
+                    lastIndex++;
+                    v.store(lastIndex-1, JetTypeMapper.TYPE_OBJECT);
+                }
+
+                firstIndex = lastIndex;
+                int curIndex = lastIndex;
+                if(getter.getInvokeOpcode() != Opcodes.INVOKESTATIC) {
+                    v.load(curIndex-1, JetTypeMapper.TYPE_OBJECT);
+                    curIndex--;
+                }
+
+                for(int i = 0; i != argumentTypes.length; i++) {
+                    int sz = argumentTypes[i].getSize();
+                    v.load(curIndex-sz, argumentTypes[i]);
+                    curIndex -= sz;
+                }
+
+                curIndex = firstIndex;
+                if(getter.getInvokeOpcode() != Opcodes.INVOKESTATIC) {
+                    v.load(curIndex-1, JetTypeMapper.TYPE_OBJECT);
+                    curIndex--;
+                }
+
+                for(int i = 0; i != argumentTypes.length; i++) {
+                    int sz = argumentTypes[i].getSize();
+                    v.load(curIndex-sz, argumentTypes[i]);
+                    curIndex -= sz;
+                }
+
+                frame.leaveTemp(size);
             }
         }
+
+        private int calcSize() {
+            assert getter != null;
+            int size = getter.getInvokeOpcode() == Opcodes.INVOKESTATIC ? 0 : 1;
+            Method signature = getter.getSignature();
+            Type[] argumentTypes = signature.getArgumentTypes();
+            for(int i = 0; i != argumentTypes.length; ++i)
+                size += argumentTypes[i].getSize();
+            return size;
+        }
     }
 
 
@@ -523,14 +577,9 @@ public abstract class StackValue {
         }
 
         @Override
-        public void dupReceiver(InstructionAdapter v, int below) {
+        public void dupReceiver(InstructionAdapter v) {
             if (!isStatic) {
-                if (below == 1) {
-                    v.dupX1();
-                }
-                else {
-                    v.dup();
-                }
+                v.dup();
             }
         }
 
@@ -557,7 +606,7 @@ public abstract class StackValue {
         }
 
         @Override
-        public void dupReceiver(InstructionAdapter v, int below) {
+        public void dupReceiver(InstructionAdapter v) {
         }
 
         @Override
@@ -607,14 +656,9 @@ public abstract class StackValue {
         }
 
         @Override
-        public void dupReceiver(InstructionAdapter v, int below) {
+        public void dupReceiver(InstructionAdapter v) {
             if (!isStatic) {
-                if (below == 1) {
-                    v.dupX1();
-                }
-                else {
-                    v.dup();
-                }
+                v.dup();
             }
         }
     }
@@ -716,13 +760,8 @@ public abstract class StackValue {
         }
 
         @Override
-        public void dupReceiver(InstructionAdapter v, int below) {
-            if (below == 1) {
-                v.dupX1();
-            }
-            else {
-                v.dup();
-            }
+        public void dupReceiver(InstructionAdapter v) {
+            v.dup();
         }
 
         @Override

compiler/backend/src/org/jetbrains/jet/codegen/intrinsics/Increment.java

@@ -32,7 +32,7 @@ public class Increment implements IntrinsicMethod {
             }
         }
         StackValue value = codegen.genQualified(receiver, operand);
-        value.dupReceiver(v, 0);
+        value. dupReceiver(v);
         value.put(expectedType, v);
         if (expectedType == Type.LONG_TYPE) {
             v.lconst(myDelta);

compiler/tests/org/jetbrains/jet/codegen/ArrayGenTest.java

@@ -107,4 +107,63 @@ public class ArrayGenTest extends CodegenTestCase {
         Method foo = generateFunction();
         assertTrue((Integer)foo.invoke(null) == 12);
     }
+
+    public void testCollectionAssignGetMultiIndex () throws Exception {
+            loadText("import java.util.ArrayList\n" +
+                         "fun box() : String { val s = ArrayList<String>(1); s.add(\"\"); s [1, -1] = \"5\"; s[2, -2] += \"7\"; return s[2,-2] }\n" +
+                         "fun ArrayList<String>.get(index1: Int, index2 : Int) = this[index1+index2]\n" +
+                         "fun ArrayList<String>.set(index1: Int, index2 : Int, elem: String) { this[index1+index2] = elem }\n");
+            System.out.println(generateToText());
+            Method foo = generateFunction("box");
+        assertTrue(foo.invoke(null).equals("57"));
+        }
+
+    public void testArrayGetAssignMultiIndex () throws Exception {
+        loadText(
+                     "fun box() : String? { val s = Array<String>(1,{ \"\" }); s [1, -1] = \"5\"; s[2, -2] += \"7\"; return s[-3,3] }\n" +
+                     "fun Array<String>.get(index1: Int, index2 : Int) = this[index1+index2]\n" +
+                     "fun Array<String>.set(index1: Int, index2 : Int, elem: String) { this[index1+index2] = elem\n }");
+        System.out.println(generateToText());
+        Method foo = generateFunction("box");
+        assertTrue(foo.invoke(null).equals("57"));
+    }
+
+    public void testCollectionGetMultiIndex () throws Exception {
+            loadText("import java.util.ArrayList\n" +
+                         "fun box() : String { val s = ArrayList<String>(1); s.add(\"\"); s [1, -1] = \"5\"; return s[2, -2] }\n" +
+                         "fun ArrayList<String>.get(index1: Int, index2 : Int) = this[index1+index2]\n" +
+                         "fun ArrayList<String>.set(index1: Int, index2 : Int, elem: String) { this[index1+index2] = elem }\n");
+            System.out.println(generateToText());
+            Method foo = generateFunction("box");
+            assertTrue(foo.invoke(null).equals("5"));
+        }
+
+    public void testArrayGetMultiIndex () throws Exception {
+        loadText(
+                     "fun box() : String? { val s = Array<String>(1,{ \"\" }); s [1, -1] = \"5\"; return s[-2, 2] }\n" +
+                     "fun Array<String>.get(index1: Int, index2 : Int) = this[index1+index2]\n" +
+                     "fun Array<String>.set(index1: Int, index2 : Int, elem: String) { this[index1+index2] = elem\n }");
+        System.out.println(generateToText());
+        Method foo = generateFunction("box");
+        assertTrue(foo.invoke(null).equals("5"));
+    }
+
+    public void testMap () throws Exception {
+        loadText(
+                     "fun box() : Int? { val s = java.util.HashMap<String,Int?>(); s[\"239\"] = 239; return s[\"239\"] }\n" +
+                     "fun java.util.HashMap<String,Int?>.set(index: String, elem: Int?) { this.put(index, elem) }");
+        System.out.println(generateToText());
+        Method foo = generateFunction("box");
+        assertTrue((Integer)foo.invoke(null) == 239);
+    }
+
+    public void testLongDouble () throws Exception {
+        loadText(
+                     "fun box() : Int { var l = IntArray(1); l[0.lng] = 4; l[0.lng] += 6; return l[0.lng];}\n" +
+                     "fun IntArray.set(index: Long, elem: Int) { this[index.int] = elem }\n" +
+                     "fun IntArray.get(index: Long) = this[index.int]");
+        System.out.println(generateToText());
+        Method foo = generateFunction("box");
+        assertTrue((Integer)foo.invoke(null) == 10);
+    }
 }