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 a0b3edf6591..c12ad5473f9 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 @@ public final class SFor extends AStatement { 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 @@ public final class SFor extends AStatement { 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 @@ public final class SFor extends AStatement { 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 @@ public final class SFor extends AStatement { 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 fc2fffb6441..e397013e06f 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 class BasicStatementTests extends ScriptTestCase { } 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" +