From 9201d63504f2db8357ac0873faff7b94a96d26a6 Mon Sep 17 00:00:00 2001
From: Jack Conradson <jack@elastic.co>
Date: Fri, 2 Feb 2018 16:28:56 -0800
Subject: [PATCH] Fixes a null pointer exception in certain cases of for loop
 usage in Painless.  The initializer and afterthought were not having their
 types appropriately cast which is necessary with expressions which in turn
 caused values to be popped off the stack that were null.

---
 .../main/java/org/elasticsearch/painless/node/SFor.java  | 9 ++++++++-
 .../org/elasticsearch/painless/BasicStatementTests.java  | 4 ++++
 2 files changed, 12 insertions(+), 1 deletion(-)

diff --git a/modules/lang-painless/src/main/java/org/elasticsearch/painless/node/SFor.java b/modules/lang-painless/src/main/java/org/elasticsearch/painless/node/SFor.java
index a0b3edf659103..c12ad5473f922 100644
--- a/modules/lang-painless/src/main/java/org/elasticsearch/painless/node/SFor.java
+++ b/modules/lang-painless/src/main/java/org/elasticsearch/painless/node/SFor.java
@@ -76,7 +76,7 @@ void analyze(Locals locals) {
         locals = Locals.newLocalScope(locals);
 
         if (initializer != null) {
-            if (initializer instanceof AStatement) {
+            if (initializer instanceof SDeclBlock) {
                 initializer.analyze(locals);
             } else if (initializer instanceof AExpression) {
                 AExpression initializer = (AExpression)this.initializer;
@@ -87,6 +87,9 @@ void analyze(Locals locals) {
                 if (!initializer.statement) {
                     throw createError(new IllegalArgumentException("Not a statement."));
                 }
+
+                initializer.expected = initializer.actual;
+                this.initializer = initializer.cast(locals);
             } else {
                 throw createError(new IllegalStateException("Illegal tree structure."));
             }
@@ -119,6 +122,9 @@ void analyze(Locals locals) {
             if (!afterthought.statement) {
                 throw createError(new IllegalArgumentException("Not a statement."));
             }
+
+            afterthought.expected = afterthought.actual;
+            afterthought = afterthought.cast(locals);
         }
 
         if (block != null) {
@@ -197,6 +203,7 @@ void write(MethodWriter writer, Globals globals) {
         if (afterthought != null) {
             writer.mark(begin);
             afterthought.write(writer, globals);
+            writer.writePop(MethodWriter.getType(afterthought.expected).getSize());
         }
 
         if (afterthought != null || !allEscape) {
diff --git a/modules/lang-painless/src/test/java/org/elasticsearch/painless/BasicStatementTests.java b/modules/lang-painless/src/test/java/org/elasticsearch/painless/BasicStatementTests.java
index fc2fffb6441a4..e397013e06f20 100644
--- a/modules/lang-painless/src/test/java/org/elasticsearch/painless/BasicStatementTests.java
+++ b/modules/lang-painless/src/test/java/org/elasticsearch/painless/BasicStatementTests.java
@@ -108,8 +108,12 @@ public void testDoWhileStatement() {
     }
 
     public void testForStatement() {
+        assertEquals(6, exec("int x, y; for (x = 0; x < 4; ++x) {y += x;} return y;"));
         assertEquals("aaaaaa", exec("String c = \"a\"; for (int x = 0; x < 5; ++x) c += \"a\"; return c;"));
 
+        assertEquals(6, exec("double test() { return 0.0; }" +
+            "int x, y; for (test(); x < 4; test()) {y += x; ++x;} return y;"));
+
         Object value = exec(
                 " int[][] b = new int[5][5];  \n" +
                 " for (int x = 0; x < 5; ++x) {     \n" +