Skip to content

Commit 876fbed

Browse files
eric-milleseric-milles
committed
GROOVY-10683: support for (index, item in list) looping
1 parent ee7de65 commit 876fbed

13 files changed

Lines changed: 164 additions & 52 deletions

File tree

‎src/antlr/GroovyParser.g4‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -691,7 +691,11 @@ forControl
691691
;
692692

693693
enhancedForControl
694-
: variableModifiersOpt type? identifier (COLON | IN) expression
694+
: (indexVariable COMMA)? variableModifiersOpt type? identifier (COLON | IN) expression
695+
;
696+
697+
indexVariable
698+
: (BuiltInPrimitiveType | DEF | VAR)? identifier
695699
;
696700

697701
originalForControl

‎src/main/java/org/apache/groovy/parser/antlr4/AstBuilder.java‎

Lines changed: 26 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -147,6 +147,7 @@
147147
import java.util.Map;
148148
import java.util.Objects;
149149
import java.util.Optional;
150+
import java.util.function.Function;
150151
import java.util.stream.Collectors;
151152

152153
import static groovy.lang.Tuple.tuple;
@@ -469,15 +470,15 @@ public Statement visitLoopStmtAlt(final LoopStmtAltContext ctx) {
469470

470471
@Override
471472
public ForStatement visitForStmtAlt(final ForStmtAltContext ctx) {
472-
var controlElements = this.visitForControl(ctx.forControl());
473+
Function<Statement, ForStatement> maker = this.visitForControl(ctx.forControl());
473474

474-
Statement loopBlock = this.unpackStatement((Statement) this.visit(ctx.statement()));
475+
Statement loopBody = this.unpackStatement((Statement) this.visit(ctx.statement()));
475476

476-
return configureAST(new ForStatement(controlElements.getV1(), controlElements.getV2(), loopBlock), ctx);
477+
return configureAST(maker.apply(loopBody), ctx);
477478
}
478479

479480
@Override
480-
public Tuple2<Parameter, Expression> visitForControl(final ForControlContext ctx) {
481+
public Function<Statement, ForStatement> visitForControl(final ForControlContext ctx) {
481482
// e.g. `for (var e : x) { }` and `for (e in x) { }`
482483
EnhancedForControlContext enhancedCtx = ctx.enhancedForControl();
483484
if (asBoolean(enhancedCtx)) {
@@ -494,22 +495,37 @@ public Tuple2<Parameter, Expression> visitForControl(final ForControlContext ctx
494495
}
495496

496497
@Override
497-
public Tuple2<Parameter, Expression> visitEnhancedForControl(final EnhancedForControlContext ctx) {
498-
var parameter = configureAST(new Parameter(this.visitType(ctx.type()), this.visitIdentifier(ctx.identifier())), ctx.identifier());
499-
new ModifierManager(this, this.visitVariableModifiersOpt(ctx.variableModifiersOpt())).processParameter(parameter);
498+
public Function<Statement, ForStatement> visitEnhancedForControl(final EnhancedForControlContext ctx) {
499+
var indexParameter = asBoolean(ctx.indexVariable()) ? this.visitIndexVariable(ctx.indexVariable()) : null;
500500

501-
return tuple(parameter, (Expression) this.visit(ctx.expression()));
501+
var valueParameter = configureAST(new Parameter(this.visitType(ctx.type()), this.visitIdentifier(ctx.identifier())), ctx.identifier());
502+
ModifierManager modifierManager = new ModifierManager(this, this.visitVariableModifiersOpt(ctx.variableModifiersOpt()));
503+
modifierManager.processParameter(valueParameter);
504+
505+
return (body) -> new ForStatement(indexParameter, valueParameter, (Expression) this.visit(ctx.expression()), body);
506+
}
507+
508+
@Override
509+
public Parameter visitIndexVariable(final IndexVariableContext ctx) {
510+
var primitiveType = ctx.BuiltInPrimitiveType();
511+
if (primitiveType != null && primitiveType.getText().length() != 3) {
512+
throw this.createParsingFailedException(primitiveType.getText() + " is not allowed here", ctx);
513+
}
514+
515+
var indexParameter = configureAST(new Parameter(ClassHelper.int_TYPE, this.visitIdentifier(ctx.identifier())), ctx.identifier());
516+
indexParameter.setModifiers(Opcodes.ACC_FINAL);
517+
return indexParameter;
502518
}
503519

504520
@Override
505-
public Tuple2<Parameter, Expression> visitOriginalForControl(final OriginalForControlContext ctx) {
521+
public Function<Statement, ForStatement> visitOriginalForControl(final OriginalForControlContext ctx) {
506522
ClosureListExpression closureListExpression = new ClosureListExpression();
507523
closureListExpression.addExpression(this.visitForInit(ctx.forInit()));
508524
closureListExpression.addExpression(Optional.ofNullable(ctx.expression())
509525
.map(e -> (Expression) this.visit(e)).orElse(EmptyExpression.INSTANCE));
510526
closureListExpression.addExpression(this.visitForUpdate(ctx.forUpdate()));
511527

512-
return tuple(ForStatement.FOR_LOOP_DUMMY, closureListExpression);
528+
return (body) -> new ForStatement(ForStatement.FOR_LOOP_DUMMY, closureListExpression, body);
513529
}
514530

515531
@Override

‎src/main/java/org/codehaus/groovy/ast/stmt/ForStatement.java‎

Lines changed: 34 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -25,40 +25,66 @@
2525
import org.codehaus.groovy.ast.VariableScope;
2626
import org.codehaus.groovy.ast.expr.Expression;
2727

28+
import static java.util.Objects.requireNonNull;
29+
2830
/**
2931
* Represents a for loop in Groovy.
3032
*/
3133
public class ForStatement extends Statement implements LoopingStatement {
3234

3335
public static final Parameter FOR_LOOP_DUMMY = new Parameter(ClassHelper.OBJECT_TYPE, "forLoopDummyParameter");
3436

35-
private final Parameter variable;
37+
private final Parameter indexVariable, valueVariable;
3638
private Expression collectionExpression;
3739
private Statement loopBlock;
3840

39-
public ForStatement(final Parameter variable, final Expression collectionExpression, final Statement loopBlock) {
40-
this.variable = variable;
41+
/**
42+
* @since 5.0.0
43+
*/
44+
public ForStatement(final Parameter indexVariable, final Parameter valueVariable, final Expression collectionExpression, final Statement loopBlock) {
45+
this.indexVariable = indexVariable; // optional component
46+
this.valueVariable = requireNonNull(valueVariable);
4147
setCollectionExpression(collectionExpression);
4248
setLoopBlock(loopBlock);
4349
}
4450

51+
public ForStatement(final Parameter variable, final Expression collectionExpression, final Statement loopBlock) {
52+
this(null, variable, collectionExpression, loopBlock);
53+
}
54+
4555
public void setCollectionExpression(final Expression collectionExpression) {
46-
this.collectionExpression = collectionExpression;
56+
this.collectionExpression = requireNonNull(collectionExpression);
4757
}
4858

4959
@Override
5060
public void setLoopBlock(final Statement loopBlock) {
51-
this.loopBlock = loopBlock;
61+
this.loopBlock = requireNonNull(loopBlock);
5262
}
5363

5464
//--------------------------------------------------------------------------
5565

66+
/**
67+
* @since 5.0.0
68+
*/
69+
public Parameter getIndexVariable() {
70+
return indexVariable;
71+
}
72+
73+
/**
74+
* @since 5.0.0
75+
*/
76+
public Parameter getValueVariable() {
77+
return valueVariable != FOR_LOOP_DUMMY ? valueVariable : null;
78+
}
79+
80+
@Deprecated(since = "5.0.0")
5681
public Parameter getVariable() {
57-
return variable;
82+
return valueVariable;
5883
}
5984

85+
@Deprecated(since = "5.0.0")
6086
public ClassNode getVariableType() {
61-
return variable.getType();
87+
return valueVariable.getType();
6288
}
6389

6490
public Expression getCollectionExpression() {
@@ -79,7 +105,7 @@ public VariableScope getVariableScope() {
79105
}
80106

81107
public void setVariableScope(final VariableScope scope) {
82-
this.scope = scope;
108+
this.scope = scope;
83109
}
84110

85111
//--------------------------------------------------------------------------

‎src/main/java/org/codehaus/groovy/classgen/VariableScopeVisitor.java‎

Lines changed: 8 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -69,7 +69,9 @@
6969
import java.util.LinkedList;
7070
import java.util.Map;
7171
import java.util.Objects;
72+
import java.util.Optional;
7273
import java.util.function.BiConsumer;
74+
import java.util.function.Consumer;
7375
import java.util.function.Supplier;
7476

7577
import static java.lang.reflect.Modifier.isFinal;
@@ -551,9 +553,12 @@ public void visitExpressionStatement(final ExpressionStatement statement) {
551553
public void visitForLoop(final ForStatement statement) {
552554
pushState();
553555
statement.setVariableScope(currentScope);
554-
Parameter parameter = statement.getVariable();
555-
parameter.setInStaticContext(currentScope.isInStaticContext());
556-
if (parameter != ForStatement.FOR_LOOP_DUMMY) declare(parameter, statement);
556+
Consumer<Parameter> define = (parameter) -> {
557+
parameter.setInStaticContext(currentScope.isInStaticContext());
558+
declare(parameter, statement);
559+
};
560+
Optional.ofNullable(statement.getIndexVariable()).ifPresent(define);
561+
Optional.ofNullable(statement.getValueVariable()).ifPresent(define);
557562
super.visitForLoop(statement);
558563
popState();
559564
}

‎src/main/java/org/codehaus/groovy/classgen/VerifierCodeVisitor.java‎

Lines changed: 9 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121
import org.codehaus.groovy.ast.ASTNode;
2222
import org.codehaus.groovy.ast.ClassNode;
2323
import org.codehaus.groovy.ast.CodeVisitorSupport;
24+
import org.codehaus.groovy.ast.Variable;
2425
import org.codehaus.groovy.ast.expr.ConstructorCallExpression;
2526
import org.codehaus.groovy.ast.expr.Expression;
2627
import org.codehaus.groovy.ast.expr.FieldExpression;
@@ -30,6 +31,8 @@
3031
import org.codehaus.groovy.ast.stmt.ForStatement;
3132
import org.codehaus.groovy.syntax.RuntimeParserException;
3233

34+
import java.util.Optional;
35+
3336
/**
3437
* Performs various checks on code inside methods and constructors
3538
* including checking for valid field, variables names etc. that
@@ -44,9 +47,12 @@ public VerifierCodeVisitor(ClassNode classNode) {
4447
}
4548

4649
@Override
47-
public void visitForLoop(ForStatement expression) {
48-
assertValidIdentifier(expression.getVariable().getName(), "for loop variable name", expression);
49-
super.visitForLoop(expression);
50+
public void visitForLoop(ForStatement statement) {
51+
Optional.ofNullable(statement.getIndexVariable()).map(Variable::getName)
52+
.ifPresent(name -> assertValidIdentifier(name, "for loop index variable name", statement));
53+
Optional.ofNullable(statement.getValueVariable()).map(Variable::getName)
54+
.ifPresent(name -> assertValidIdentifier(name, "for loop value variable name", statement));
55+
super.visitForLoop(statement);
5056
}
5157

5258
@Override

‎src/main/java/org/codehaus/groovy/classgen/asm/StatementWriter.java‎

Lines changed: 11 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -93,7 +93,7 @@ public void writeBlockStatement(final BlockStatement block) {
9393
}
9494

9595
public void writeForStatement(final ForStatement statement) {
96-
if (statement.getVariable() == ForStatement.FOR_LOOP_DUMMY) {
96+
if (statement.getCollectionExpression() instanceof ClosureListExpression) {
9797
writeForLoopWithClosureList(statement);
9898
} else {
9999
writeForInLoop(statement);
@@ -122,8 +122,12 @@ protected void writeForInLoopControlAndBlock(ForStatement statement) {
122122
MethodVisitor mv = controller.getMethodVisitor();
123123
OperandStack operandStack = controller.getOperandStack();
124124

125-
// declare the loop counter
126-
BytecodeVariable variable = compileStack.defineVariable(statement.getVariable(), false);
125+
// declare the loop index and value variables
126+
BytecodeVariable indexVariable = Optional.ofNullable(statement.getIndexVariable()).map(iv -> {
127+
mv.visitInsn(ICONST_M1); // initialize to -1 so increment can pair with next()
128+
return compileStack.defineVariable(iv, true);
129+
}).orElse(null);
130+
BytecodeVariable valueVariable = compileStack.defineVariable(statement.getValueVariable(), false);
127131

128132
// get the iterator and generate the loop control
129133
int iterator = compileStack.defineTemporaryVariable("iterator", ClassHelper.Iterator_TYPE, true);
@@ -141,7 +145,10 @@ protected void writeForInLoopControlAndBlock(ForStatement statement) {
141145
mv.visitVarInsn(ALOAD, iterator);
142146
writeIteratorNext(mv);
143147
operandStack.push(ClassHelper.OBJECT_TYPE);
144-
operandStack.storeVar(variable);
148+
operandStack.storeVar(valueVariable);
149+
if (indexVariable != null) {
150+
mv.visitIincInsn(indexVariable.getIndex(), 1);
151+
}
145152

146153
// generate the loop body
147154
statement.getLoopBlock().visit(controller.getAcg());

‎src/main/java/org/codehaus/groovy/classgen/asm/sc/StaticTypesStatementWriter.java‎

Lines changed: 21 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,7 @@
3838

3939
import java.util.Enumeration;
4040
import java.util.Objects;
41+
import java.util.Optional;
4142

4243
import static org.objectweb.asm.Opcodes.AALOAD;
4344
import static org.objectweb.asm.Opcodes.ALOAD;
@@ -50,6 +51,7 @@
5051
import static org.objectweb.asm.Opcodes.GOTO;
5152
import static org.objectweb.asm.Opcodes.IALOAD;
5253
import static org.objectweb.asm.Opcodes.ICONST_0;
54+
import static org.objectweb.asm.Opcodes.ICONST_M1;
5355
import static org.objectweb.asm.Opcodes.IFEQ;
5456
import static org.objectweb.asm.Opcodes.IFNULL;
5557
import static org.objectweb.asm.Opcodes.IF_ICMPGE;
@@ -95,9 +97,8 @@ protected void writeForInLoop(final ForStatement loop) {
9597
ClassNode collectionType = controller.getTypeChooser().resolveType(collectionExpression, controller.getClassNode());
9698

9799
int mark = operandStack.getStackLength();
98-
Parameter loopVariable = loop.getVariable();
99-
if (collectionType.isArray() && loopVariable.getType().equals(collectionType.getComponentType())) {
100-
writeOptimizedForEachLoop(loop, loopVariable, collectionExpression, collectionType);
100+
if (collectionType.isArray() && collectionType.getComponentType().equals(loop.getVariableType())) {
101+
writeOptimizedForEachLoop(loop, collectionExpression, collectionType);
101102
} else if (GeneralUtils.isOrImplements(collectionType, ENUMERATION_CLASSNODE)) {
102103
writeEnumerationBasedForEachLoop(loop, collectionExpression, collectionType);
103104
} else {
@@ -107,13 +108,16 @@ protected void writeForInLoop(final ForStatement loop) {
107108
compileStack.pop();
108109
}
109110

110-
private void writeOptimizedForEachLoop(final ForStatement loop, final Parameter loopVariable, final Expression arrayExpression, final ClassNode arrayType) {
111+
private void writeOptimizedForEachLoop(final ForStatement loop, final Expression arrayExpression, final ClassNode arrayType) {
111112
CompileStack compileStack = controller.getCompileStack();
112113
OperandStack operandStack = controller.getOperandStack();
113114
MethodVisitor mv = controller.getMethodVisitor();
114115
AsmClassGenerator acg = controller.getAcg();
115116

116-
BytecodeVariable variable = compileStack.defineVariable(loopVariable, arrayType.getComponentType(), false);
117+
BytecodeVariable indexVariable = Optional.ofNullable(loop.getIndexVariable()).map(iv -> {
118+
mv.visitInsn(ICONST_M1); return compileStack.defineVariable(iv, true);
119+
}).orElse(null);
120+
BytecodeVariable valueVariable = compileStack.defineVariable(loop.getValueVariable(), arrayType.getComponentType(), false);
117121
Label continueLabel = compileStack.getContinueLabel();
118122
Label breakLabel = compileStack.getBreakLabel();
119123

@@ -141,10 +145,13 @@ private void writeOptimizedForEachLoop(final ForStatement loop, final Parameter
141145
mv.visitJumpInsn(IF_ICMPGE, breakLabel);
142146

143147
// get array element
144-
loadFromArray(mv, operandStack, variable, array, loopIdx);
148+
loadFromArray(mv, operandStack, valueVariable, array, loopIdx);
145149

146150
// $idx += 1
147151
mv.visitIincInsn(loopIdx, 1);
152+
if (indexVariable != null) {
153+
mv.visitIincInsn(indexVariable.getIndex(), 1);
154+
}
148155

149156
// loop body
150157
loop.getLoopBlock().visit(acg);
@@ -190,7 +197,10 @@ private void writeEnumerationBasedForEachLoop(final ForStatement loop, final Exp
190197
OperandStack operandStack = controller.getOperandStack();
191198
MethodVisitor mv = controller.getMethodVisitor();
192199

193-
BytecodeVariable variable = compileStack.defineVariable(loop.getVariable(), false);
200+
BytecodeVariable indexVariable = Optional.ofNullable(loop.getIndexVariable()).map(iv -> {
201+
mv.visitInsn(ICONST_M1); return compileStack.defineVariable(iv, true);
202+
}).orElse(null);
203+
BytecodeVariable valueVariable = compileStack.defineVariable(loop.getValueVariable(), false);
194204
Label continueLabel = compileStack.getContinueLabel();
195205
Label breakLabel = compileStack.getBreakLabel();
196206

@@ -210,7 +220,10 @@ private void writeEnumerationBasedForEachLoop(final ForStatement loop, final Exp
210220
mv.visitVarInsn(ALOAD, enumeration);
211221
ENUMERATION_NEXT_METHOD.call(mv);
212222
operandStack.push(ClassHelper.OBJECT_TYPE);
213-
operandStack.storeVar(variable);
223+
operandStack.storeVar(valueVariable);
224+
if (indexVariable != null) {
225+
mv.visitIincInsn(indexVariable.getIndex(), 1);
226+
}
214227

215228
loop.getLoopBlock().visit(controller.getAcg());
216229
mv.visitJumpInsn(GOTO, continueLabel);

‎src/main/java/org/codehaus/groovy/transform/CategoryASTTransformation.java‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -221,7 +221,11 @@ public void visitExpressionStatement(final ExpressionStatement statement) {
221221
public void visitForLoop(final ForStatement statement) {
222222
Expression exp = statement.getCollectionExpression();
223223
exp.visit(this);
224-
Parameter loopParam = statement.getVariable();
224+
Parameter loopParam = statement.getIndexVariable();
225+
if (loopParam != null) {
226+
varStack.getLast().add(loopParam.getName());
227+
}
228+
loopParam = statement.getValueVariable();
225229
if (loopParam != null) {
226230
varStack.getLast().add(loopParam.getName());
227231
}

‎src/main/java/org/codehaus/groovy/transform/sc/StaticCompilationVisitor.java‎

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -257,9 +257,10 @@ public void visitForLoop(final ForStatement statement) {
257257
super.visitForLoop(statement);
258258
var collectionExpression = statement.getCollectionExpression();
259259
if (!(collectionExpression instanceof ClosureListExpression)) {
260-
if (statement.getVariable().isDynamicTyped()) { // GROOVY-8169
261-
ClassNode inferredType = getType(statement.getVariable());
262-
statement.getVariable().setType(inferredType); // GROOVY-5640, GROOVY-5641
260+
var valueVariable = statement.getValueVariable();
261+
if (valueVariable.isDynamicTyped()) { // GROOVY-8169
262+
ClassNode inferredType = getType(valueVariable);
263+
valueVariable.setType(inferredType); // GROOVY-5640, GROOVY-5641
263264
}
264265
}
265266
}

‎src/spec/test/SemanticsTest.groovy‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -245,6 +245,11 @@ final class SemanticsTest {
245245
list.add(c)
246246
}
247247
assert list == ['a', 'b', 'c']
248+
249+
// iterate with index
250+
for ( int i, k in map.keySet() ) {
251+
assert map.get(k) == i + 1
252+
}
248253
// end::groovy_for_loop_example[]
249254
}
250255

0 commit comments

Comments
 (0)