From 65aca4f71c9b88a41e2c5d6bd7fc9cceb7c0d99d Mon Sep 17 00:00:00 2001
From: Uwe Schindler <uwe@thetaphi.de>
Date: Sun, 15 May 2016 22:58:11 +0200
Subject: [PATCH] painless: make compound statement like a[1]++ work with def.
 There was also a bug with the size of LDefArray: fixed

---
 .../java/org/elasticsearch/painless/node/ADefLink.java    | 2 ++
 .../main/java/org/elasticsearch/painless/node/EChain.java | 6 ++++++
 .../java/org/elasticsearch/painless/node/LDefArray.java   | 2 +-
 .../test/java/org/elasticsearch/painless/ArrayTests.java  | 8 ++++++++
 4 files changed, 17 insertions(+), 1 deletion(-)

diff --git a/modules/lang-painless/src/main/java/org/elasticsearch/painless/node/ADefLink.java b/modules/lang-painless/src/main/java/org/elasticsearch/painless/node/ADefLink.java
index ad0beb23cbc..cf4351a63a7 100644
--- a/modules/lang-painless/src/main/java/org/elasticsearch/painless/node/ADefLink.java
+++ b/modules/lang-painless/src/main/java/org/elasticsearch/painless/node/ADefLink.java
@@ -23,6 +23,8 @@ import org.elasticsearch.painless.Definition.Type;
 
 /**
  * The superclass for all LDef* (link) nodes that store or return a DEF. (Internal only.)
+ * For this node it is allowed to change {@link ALink#after} from outside, by default
+ * {@code after} is {@code DEF}.
  */
 abstract class ADefLink extends ALink {
     
diff --git a/modules/lang-painless/src/main/java/org/elasticsearch/painless/node/EChain.java b/modules/lang-painless/src/main/java/org/elasticsearch/painless/node/EChain.java
index 463c9a5a4ff..3dd64975757 100644
--- a/modules/lang-painless/src/main/java/org/elasticsearch/painless/node/EChain.java
+++ b/modules/lang-painless/src/main/java/org/elasticsearch/painless/node/EChain.java
@@ -215,6 +215,12 @@ public final class EChain extends AExpression {
         there = AnalyzerCaster.getLegalCast(definition, location, last.after, promote, false);
         back = AnalyzerCaster.getLegalCast(definition, location, promote, last.after, true);
 
+        if (last instanceof ADefLink) {
+            final ADefLink lastDef = (ADefLink) last;
+            // Unfortunately, we don't know the real type because we load from DEF and store to DEF!
+            lastDef.storeValueType = last.after;
+        }
+
         statement = true;
         actual = read ? last.after : definition.voidType;
     }
diff --git a/modules/lang-painless/src/main/java/org/elasticsearch/painless/node/LDefArray.java b/modules/lang-painless/src/main/java/org/elasticsearch/painless/node/LDefArray.java
index d3d129d3ba9..df1c5ffac0a 100644
--- a/modules/lang-painless/src/main/java/org/elasticsearch/painless/node/LDefArray.java
+++ b/modules/lang-painless/src/main/java/org/elasticsearch/painless/node/LDefArray.java
@@ -36,7 +36,7 @@ final class LDefArray extends ADefLink {
     AExpression index;
 
     LDefArray(final int line, final String location, final AExpression index) {
-        super(line, location, 0);
+        super(line, location, 2);
 
         this.index = index;
     }
diff --git a/modules/lang-painless/src/test/java/org/elasticsearch/painless/ArrayTests.java b/modules/lang-painless/src/test/java/org/elasticsearch/painless/ArrayTests.java
index 51ee659e501..aa873e45d36 100644
--- a/modules/lang-painless/src/test/java/org/elasticsearch/painless/ArrayTests.java
+++ b/modules/lang-painless/src/test/java/org/elasticsearch/painless/ArrayTests.java
@@ -54,6 +54,14 @@ public class ArrayTests extends ScriptTestCase {
         assertEquals(5, exec("def x = new def[4]; x[0] = 5; return x[0];"));
     }
     
+    public void testArrayCompoundInt() {
+        assertEquals(6, exec("int[] x = new int[5]; x[0] = 5; x[0]++; return x[0];"));
+    }
+    
+    public void testArrayCompoundDef() {
+        assertEquals(6, exec("def x = new int[5]; x[0] = 5; x[0]++; return x[0];"));
+    }
+    
     public void testForLoop() {
         assertEquals(999*1000/2, exec("def a = new int[1000]; for (int x = 0; x < a.length; x++) { a[x] = x; } "+
             "int total = 0; for (int x = 0; x < a.length; x++) { total += a[x]; } return total;"));