'val' now also works in foreach loops, on both javac and ecj / eclipse.

5 files changed, 180 insertions, 129 deletions

diff --git a/src/core/lombok/eclipse/handlers/HandleVal.java b/src/core/lombok/eclipse/handlers/HandleVal.java
index 6fa39b10..b3cfa879 100644
--- a/src/core/lombok/eclipse/handlers/HandleVal.java
+++ b/src/core/lombok/eclipse/handlers/HandleVal.java
@@ -26,6 +26,7 @@ import lombok.eclipse.EclipseASTVisitor;
 import lombok.eclipse.EclipseNode;
 
 import org.eclipse.jdt.internal.compiler.ast.ArrayInitializer;
+import org.eclipse.jdt.internal.compiler.ast.ForeachStatement;
 import org.eclipse.jdt.internal.compiler.ast.LocalDeclaration;
 import org.eclipse.jdt.internal.compiler.ast.SingleTypeReference;
 import org.mangosdk.spi.ProviderFor;
@@ -41,7 +42,14 @@ public class HandleVal extends EclipseASTAdapter {
 			if (token == null || token.length != 3) return;
 			else if (token[0] != 'v' || token[1] != 'a' || token[2] != 'l') return;
 			
-			if (local.initialization == null) {
+			boolean variableOfForEach = false;
+			
+			if (localNode.directUp().get() instanceof ForeachStatement) {
+				ForeachStatement fs = (ForeachStatement) localNode.directUp().get();
+				variableOfForEach = fs.elementVariable == local;
+			}
+			
+			if (local.initialization == null && !variableOfForEach) {
 				localNode.addError("'val' on a local variable requires an initializer expression");
 			}
 			
diff --git a/src/core/lombok/javac/JavacResolution.java b/src/core/lombok/javac/JavacResolution.java
index 14de1ff8..e0eb436d 100644
--- a/src/core/lombok/javac/JavacResolution.java
+++ b/src/core/lombok/javac/JavacResolution.java
@@ -12,11 +12,13 @@ import javax.tools.DiagnosticListener;
 
 import com.sun.tools.javac.code.BoundKind;
 import com.sun.tools.javac.code.Symbol.TypeSymbol;
+import com.sun.tools.javac.code.Symtab;
 import com.sun.tools.javac.code.Type.ArrayType;
 import com.sun.tools.javac.code.Type.CapturedType;
 import com.sun.tools.javac.code.Type.ClassType;
 import com.sun.tools.javac.code.Type;
 import com.sun.tools.javac.code.TypeTags;
+import com.sun.tools.javac.code.Types;
 import com.sun.tools.javac.comp.Attr;
 import com.sun.tools.javac.comp.AttrContext;
 import com.sun.tools.javac.comp.Enter;
@@ -234,64 +236,11 @@ public class JavacResolution {
 			if (copyAt != null) return;
 			copyAt = tree;
 		}
+		
+		@Override public void visitTree(JCTree that) {
+		}
 	}
 	
-//	/**
-//	 * The {@code Env} object primarily tracks legal symbol names. i.e. its the lexical scope. To build it, we need to go from the top and drill down to the current node,
-//	 * updating the {@code Env} object at each step. This TreeVisitor does that. Requires {@code enterTrees} to be called first (this is done before processors run, normally).
-//	 */
-//	private static final class EnvChainer extends JCTree.Visitor {
-//		private Env<AttrContext> env = null;
-//		private Enter enter;
-//		private MemberEnter memberEnter;
-//		private Attr attr;
-//		private JCTree target;
-//		private boolean blocksAreInitializers;
-//		
-//		EnvChainer(Context context) {
-//			this.enter = Enter.instance(context);
-//			this.memberEnter = MemberEnter.instance(context);
-//			this.attr = Attr.instance(context);
-//		}
-//		
-//		Env<AttrContext> get() {
-//			return env;
-//		}
-//		
-//		@Override public void visitTopLevel(JCCompilationUnit tree) {
-//			env = enter.getTopLevelEnv(tree);
-//		}
-//		
-//		@Override public void visitClassDef(JCClassDecl tree) {
-//			// The commented out one leaves the 'lint' field unset, which causes NPEs during attrib. So, we use the other one.
-//			//env = enter.classEnv((JCClassDecl) tree, env);
-//			env = enter.getClassEnv(tree.sym);
-//			blocksAreInitializers = true;
-//		}
-//		
-//		@Override public void visitMethodDef(JCMethodDecl tree) {
-//			env = memberEnter.getMethodEnv(tree, env);
-//			blocksAreInitializers = false;
-//			if (tree.body != null) visitBlock(tree.body);
-//		}
-//		
-//		@Override public void visitBlock(JCBlock tree) {
-//			if (blocksAreInitializers) attr.attribStat(tree, env);
-//			for (JCStatement stat : tree.stats) {
-//				if (stat == target) return;
-//				attr.attribStat(stat, env);
-//			}
-//		}
-//		
-//		@Override public void visitTree(JCTree tree) {
-//			// Do nothing
-//		}
-//		
-//		public void setTarget(JCTree target) {
-//			this.target = target;
-//		}
-//	};
-	
 	public Map<JCTree, JCTree> resolve(JavacNode node) {
 		ArrayDeque<JCTree> stack = new ArrayDeque<JCTree>();
 		
@@ -325,58 +274,26 @@ public class JavacResolution {
 		else throw new IllegalStateException("Called with something that isn't a block, method decl, or variable decl");
 	}
 	
-//	public void resolveUpTo(JavacNode statementNode) {
-//		ArrayDeque<JCTree> stack = new ArrayDeque<JCTree>();
-//		
-//		{
-//			JavacNode n = statementNode;
-//			while (n != null) {
-//				stack.push(n.get());
-//				n = n.up();
-//			}
-//		}
-//		
-//		logDisabler.disableLoggers();
-//		try {
-//			JCTree tree = stack.isEmpty() ? null : stack.pop();
-//			while (!stack.isEmpty()) {
-//				JCTree target = stack.pop();
-//				envChainer.setTarget(target);
-//				tree.accept(envChainer);
-//				tree = target;
-//			}
-//			if (tree != null) {
-//				envChainer.setTarget(null);
-//				tree.accept(envChainer);
-//			}
-//			
-////			System.out.println("ATTRIBSTAT: " + attr.attribStat(statementNode.get(), envChainer.get()));
-////			if (statementNode.get() instanceof JCVariableDecl) {
-////				System.out.println("Force-tribbing expr");
-////				JCExpression init = ((JCVariableDecl)statementNode.get()).init;
-////				System.out.println("ATTRIBEXPR: " + attr.attribExpr(init, envChainer.get(), Type.noType));
-////				System.out.println("TYPE: " + ((JCVariableDecl)statementNode.get()).init.type);
-////			}
-//		} finally {
-//			logDisabler.enableLoggers();
-//		}
-//	}
-//	
-//	public void resolveExpr(JCExpression expr) {
-//		logDisabler.disableLoggers();
-//		try {
-//			attr.attribExpr(expr, envChainer.get(), Type.noType);
-//		} finally {
-//			logDisabler.enableLoggers();
-//		}
-//	}
-	
 	public static class TypeNotConvertibleException extends Exception {
 		public TypeNotConvertibleException(String msg) {
 			super(msg);
 		}
 	}
 	
+	public static Type ifTypeIsIterableToComponent(Type type, JavacAST ast) {
+		Types types = Types.instance(ast.getContext());
+		Symtab syms = Symtab.instance(ast.getContext());
+		Type boundType = types.upperBound(type);
+		Type elemTypeIfArray = types.elemtype(boundType);
+		if (elemTypeIfArray != null) return elemTypeIfArray;
+		
+		Type base = types.asSuper(boundType, syms.iterableType.tsym);
+		if (base == null) return syms.objectType;
+		
+		List<Type> iterableParams = base.allparams();
+		return iterableParams.isEmpty() ? syms.objectType : types.upperBound(iterableParams.head);
+	}
+	
 	public static JCExpression typeToJCTree(Type type, TreeMaker maker, JavacAST ast) throws TypeNotConvertibleException {
 		return typeToJCTree(type, maker, ast, false);
 	}
diff --git a/src/core/lombok/javac/handlers/HandleVal.java b/src/core/lombok/javac/handlers/HandleVal.java
index 7161a7c3..5da6fec0 100644
--- a/src/core/lombok/javac/handlers/HandleVal.java
+++ b/src/core/lombok/javac/handlers/HandleVal.java
@@ -64,7 +64,6 @@ public class HandleVal extends JavacASTAdapter {
 			}
 			
 			local.mods.flags |= Flags.FINAL;
-			JCExpression oldVarType = local.vartype;
 			local.vartype = JavacResolution.createJavaLangObject(localNode.getTreeMaker(), localNode.getAst());
 			
 			Type type;
@@ -89,25 +88,26 @@ public class HandleVal extends JavacASTAdapter {
 					JCExpression replacement;
 					
 					if (rhsOfEnhancedForLoop != null) {
-						localNode.addError("For now 'val' cannot be used in enhanced for loops. Fixing this is a high priority lombok issue though!");
-						return;
-						// TODO: Fix enhanced for loops - then uncomment a bunch of lines in test/transform/resource/*/ValInFor.java
+						Type componentType = JavacResolution.ifTypeIsIterableToComponent(type, localNode.getAst());
+						if (componentType == null) replacement = JavacResolution.createJavaLangObject(localNode.getTreeMaker(), localNode.getAst());
+						else replacement = JavacResolution.typeToJCTree(componentType, localNode.getTreeMaker(), localNode.getAst());
+					} else {
+						replacement = JavacResolution.typeToJCTree(type, localNode.getTreeMaker(), localNode.getAst());
 					}
-					replacement = JavacResolution.typeToJCTree(type, localNode.getTreeMaker(), localNode.getAst());
+					
 					if (replacement != null) {
 						local.vartype = replacement;
 						localNode.getAst().setChanged();
 					}
-					else local.vartype = oldVarType;
+					else local.vartype = JavacResolution.createJavaLangObject(localNode.getTreeMaker(), localNode.getAst());;
 				} catch (JavacResolution.TypeNotConvertibleException e) {
 					localNode.addError("Cannot use 'val' here because initializer expression does not have a representable type: " + e.getMessage());
-					local.vartype = oldVarType;
+					local.vartype = JavacResolution.createJavaLangObject(localNode.getTreeMaker(), localNode.getAst());;
 				}
 			} catch (RuntimeException e) {
-				local.vartype = oldVarType;
+				local.vartype = JavacResolution.createJavaLangObject(localNode.getTreeMaker(), localNode.getAst());;
 				throw e;
 			}
 		}
-		// TODO search for val decls in either kind of for.
 	}
 }
diff --git a/src/eclipseAgent/lombok/eclipse/agent/EclipsePatcher.java b/src/eclipseAgent/lombok/eclipse/agent/EclipsePatcher.java
index 0c9f1ccc..b5132ca4 100644
--- a/src/eclipseAgent/lombok/eclipse/agent/EclipsePatcher.java
+++ b/src/eclipseAgent/lombok/eclipse/agent/EclipsePatcher.java
@@ -267,6 +267,7 @@ public class EclipsePatcher extends Agent {
 	
 	private static void patchHandleVal(ScriptManager sm, boolean ecj) {
 		final String LOCALDECLARATION_SIG = "org.eclipse.jdt.internal.compiler.ast.LocalDeclaration";
+		final String FOREACHSTATEMENT_SIG = "org.eclipse.jdt.internal.compiler.ast.ForeachStatement";
 		final String EXPRESSION_SIG = "org.eclipse.jdt.internal.compiler.ast.Expression";
 		final String BLOCKSCOPE_SIG = "org.eclipse.jdt.internal.compiler.lookup.BlockScope";
 		final String PARSER_SIG = "org.eclipse.jdt.internal.compiler.parser.Parser";
@@ -278,6 +279,24 @@ public class EclipsePatcher extends Agent {
 				.decisionMethod(new Hook("lombok.eclipse.agent.PatchFixes", "handleValForLocalDeclaration", "boolean", LOCALDECLARATION_SIG, BLOCKSCOPE_SIG))
 				.build());
 		
+		sm.addScript(ScriptBuilder.replaceMethodCall()
+				.target(new MethodTarget(LOCALDECLARATION_SIG, "resolve", "void", BLOCKSCOPE_SIG))
+				.methodToReplace(new Hook(EXPRESSION_SIG, "resolveType", TYPEBINDING_SIG, BLOCKSCOPE_SIG))
+				.replacementMethod(new Hook("lombok.eclipse.agent.PatchFixes", "skipResolveInitializerIfAlreadyCalled", TYPEBINDING_SIG, EXPRESSION_SIG, BLOCKSCOPE_SIG))
+				.build());
+		
+		sm.addScript(ScriptBuilder.exitEarly()
+				.target(new MethodTarget(FOREACHSTATEMENT_SIG, "resolve", "void", BLOCKSCOPE_SIG))
+				.request(StackRequest.THIS, StackRequest.PARAM1)
+				.decisionMethod(new Hook("lombok.eclipse.agent.PatchFixes", "handleValForForEach", "boolean", FOREACHSTATEMENT_SIG, BLOCKSCOPE_SIG))
+				.build());
+		
+		sm.addScript(ScriptBuilder.replaceMethodCall()
+				.target(new MethodTarget(FOREACHSTATEMENT_SIG, "resolve", "void", BLOCKSCOPE_SIG))
+				.methodToReplace(new Hook(EXPRESSION_SIG, "resolveType", TYPEBINDING_SIG, BLOCKSCOPE_SIG))
+				.replacementMethod(new Hook("lombok.eclipse.agent.PatchFixes", "skipResolveInitializerIfAlreadyCalled", TYPEBINDING_SIG, EXPRESSION_SIG, BLOCKSCOPE_SIG))
+				.build());
+		
 		if (!ecj) {
 			sm.addScript(ScriptBuilder.addField()
 					.fieldName("$initCopy")
@@ -287,17 +306,25 @@ public class EclipsePatcher extends Agent {
 					.targetClass("org.eclipse.jdt.internal.compiler.ast.LocalDeclaration")
 					.build());
 			
+			sm.addScript(ScriptBuilder.addField()
+					.fieldName("$iterableCopy")
+					.fieldType("Lorg/eclipse/jdt/internal/compiler/ast/ASTNode;")
+					.setPublic()
+					.setTransient()
+					.targetClass("org.eclipse.jdt.internal.compiler.ast.LocalDeclaration")
+					.build());
+			
 			sm.addScript(ScriptBuilder.wrapReturnValue()
 					.target(new MethodTarget(PARSER_SIG, "consumeExitVariableWithInitialization", "void"))
 					.request(StackRequest.THIS)
 					.wrapMethod(new Hook("lombok.eclipse.agent.PatchFixes", "copyInitializationOfLocalDeclarationForVal", "void", PARSER_SIG))
 					.build());
+			
+			sm.addScript(ScriptBuilder.wrapReturnValue()
+					.target(new MethodTarget(PARSER_SIG, "consumeEnhancedForStatementHeader", "void"))
+					.request(StackRequest.THIS)
+					.wrapMethod(new Hook("lombok.eclipse.agent.PatchFixes", "copyInitializationOfForEachIterable", "void", PARSER_SIG))
+					.build());
 		}
-		
-		sm.addScript(ScriptBuilder.replaceMethodCall()
-				.target(new MethodTarget(LOCALDECLARATION_SIG, "resolve", "void", BLOCKSCOPE_SIG))
-				.methodToReplace(new Hook(EXPRESSION_SIG, "resolveType", TYPEBINDING_SIG, BLOCKSCOPE_SIG))
-				.replacementMethod(new Hook("lombok.eclipse.agent.PatchFixes", "skipResolveInitializerIfAlreadyCalled", TYPEBINDING_SIG, EXPRESSION_SIG, BLOCKSCOPE_SIG))
-				.build());
 	}
 }
diff --git a/src/eclipseAgent/lombok/eclipse/agent/PatchFixes.java b/src/eclipseAgent/lombok/eclipse/agent/PatchFixes.java
index 14687e32..32335ecc 100644
--- a/src/eclipseAgent/lombok/eclipse/agent/PatchFixes.java
+++ b/src/eclipseAgent/lombok/eclipse/agent/PatchFixes.java
@@ -38,14 +38,20 @@ import org.eclipse.jdt.core.dom.SimpleName;
 import org.eclipse.jdt.internal.compiler.ast.ASTNode;
 import org.eclipse.jdt.internal.compiler.ast.AbstractVariableDeclaration;
 import org.eclipse.jdt.internal.compiler.ast.Expression;
+import org.eclipse.jdt.internal.compiler.ast.ForeachStatement;
 import org.eclipse.jdt.internal.compiler.ast.LocalDeclaration;
 import org.eclipse.jdt.internal.compiler.ast.QualifiedTypeReference;
 import org.eclipse.jdt.internal.compiler.ast.SingleTypeReference;
 import org.eclipse.jdt.internal.compiler.ast.TypeReference;
 import org.eclipse.jdt.internal.compiler.classfmt.ClassFileConstants;
+import org.eclipse.jdt.internal.compiler.lookup.ArrayBinding;
+import org.eclipse.jdt.internal.compiler.lookup.Binding;
 import org.eclipse.jdt.internal.compiler.lookup.BlockScope;
+import org.eclipse.jdt.internal.compiler.lookup.ParameterizedTypeBinding;
+import org.eclipse.jdt.internal.compiler.lookup.ReferenceBinding;
 import org.eclipse.jdt.internal.compiler.lookup.TypeBinding;
 import org.eclipse.jdt.internal.compiler.lookup.TypeConstants;
+import org.eclipse.jdt.internal.compiler.lookup.TypeIds;
 import org.eclipse.jdt.internal.compiler.parser.Parser;
 
 public class PatchFixes {
@@ -140,18 +146,54 @@ public class PatchFixes {
 		return new BufferedOutputStream(PostCompiler.wrapOutputStream(out, fileName, DiagnosticsReceiver.CONSOLE));
 	}
 	
+	private static Field astStackField, astPtrField;
+	
+	static {
+		try {
+			astStackField = Parser.class.getDeclaredField("astStack");
+			astStackField.setAccessible(true);
+			astPtrField = Parser.class.getDeclaredField("astPtr");
+			astPtrField.setAccessible(true);
+		} catch (Exception e) {
+			// Most likely we're in ecj or some other plugin usage of the eclipse compiler. No need for this.
+		}
+	}
+	
+	public static void copyInitializationOfForEachIterable(Parser parser) {
+		ASTNode[] astStack;
+		int astPtr;
+		try {
+			astStack = (ASTNode[]) astStackField.get(parser);
+			astPtr = (Integer)astPtrField.get(parser);
+		} catch (Exception e) {
+			// Most likely we're in ecj or some other plugin usage of the eclipse compiler. No need for this.
+			return;
+		}
+		
+		ForeachStatement foreachDecl = (ForeachStatement) astStack[astPtr];
+		ASTNode init = foreachDecl.collection;
+		if (init == null) return;
+		if (foreachDecl.elementVariable != null && foreachDecl.elementVariable.type instanceof SingleTypeReference) {
+			SingleTypeReference ref = (SingleTypeReference) foreachDecl.elementVariable.type;
+			if (ref.token == null || ref.token.length != 3 || ref.token[0] != 'v' || ref.token[1] != 'a' || ref.token[2] != 'l') return;
+		} else return;
+		
+		try {
+			if (iterableCopyField != null) iterableCopyField.set(foreachDecl.elementVariable, init);
+		} catch (Exception e) {
+			// In ecj mode this field isn't there and we don't need the copy anyway, so, we ignore the exception.
+		}
+	}
+	
 	public static void copyInitializationOfLocalDeclarationForVal(Parser parser) {
 		ASTNode[] astStack;
 		int astPtr;
 		try {
-			Field astStackF = Parser.class.getDeclaredField("astStack");
-			astStackF.setAccessible(true);
-			astStack = (ASTNode[]) astStackF.get(parser);
-			Field astPtrF = Parser.class.getDeclaredField("astPtr");
-			astPtrF.setAccessible(true);
-			astPtr = (Integer)astPtrF.get(parser);
+			astStack = (ASTNode[]) astStackField.get(parser);
+			astPtr = (Integer)astPtrField.get(parser);
 		} catch (Exception e) {
-			throw new RuntimeException(e);
+			// Most likely we're in ecj or some other plugin usage of the eclipse compiler. No need for this.
+			return;
 		}
 		AbstractVariableDeclaration variableDecl = (AbstractVariableDeclaration) astStack[astPtr];
 		if (!(variableDecl instanceof LocalDeclaration)) return;
@@ -165,22 +207,69 @@ public class PatchFixes {
 		try {
 			if (initCopyField != null) initCopyField.set(variableDecl, init);
 		} catch (Exception e) {
-			e.printStackTrace(System.out);
 			// In ecj mode this field isn't there and we don't need the copy anyway, so, we ignore the exception.
 		}
 	}
 	
-	private static Field initCopyField;
+	private static Field initCopyField, iterableCopyField;
 	
 	static {
 		try {
 			initCopyField = LocalDeclaration.class.getDeclaredField("$initCopy");
+			iterableCopyField = LocalDeclaration.class.getDeclaredField("$iterableCopy");
 		} catch (Throwable t) {
 			 //ignore - no $initCopy exists when running in ecj.
 		}
 	}
 	
+	public static boolean handleValForForEach(ForeachStatement forEach, BlockScope scope) {
+		if (forEach.elementVariable != null && forEach.elementVariable.type instanceof SingleTypeReference) {
+			char[] token = ((SingleTypeReference)forEach.elementVariable.type).token;
+			if (token == null || token.length != 3) return false;
+			else if (token[0] != 'v' || token[1] != 'a' || token[2] != 'l') return false;
+		} else return false;
+		
+		TypeBinding component = getForEachComponentType(forEach.collection, scope);
+		TypeReference replacement = Eclipse.makeType(component, forEach.elementVariable.type, false);
+		
+		forEach.elementVariable.modifiers |= ClassFileConstants.AccFinal;
+		forEach.elementVariable.type = replacement != null ? replacement :
+				new QualifiedTypeReference(TypeConstants.JAVA_LANG_OBJECT, Eclipse.poss(forEach.elementVariable.type, 3));
+		
+		return false;
+	}
+
+	private static TypeBinding getForEachComponentType(Expression collection, BlockScope scope) {
+		if (collection != null) {
+			TypeBinding resolved = collection.resolveType(scope);
+			if (resolved.isArrayType()) {
+				resolved = ((ArrayBinding) resolved).elementsType();
+				return resolved;
+			} else if (resolved instanceof ReferenceBinding) {
+				ReferenceBinding iterableType = ((ReferenceBinding)resolved).findSuperTypeOriginatingFrom(TypeIds.T_JavaLangIterable, false);
+				
+				TypeBinding[] arguments = null;
+				if (iterableType != null) switch (iterableType.kind()) {
+					case Binding.GENERIC_TYPE : // for (T t : Iterable<T>) - in case used inside Iterable itself
+						arguments = iterableType.typeVariables();
+						break;
+					case Binding.PARAMETERIZED_TYPE : // for(E e : Iterable<E>)
+						arguments = ((ParameterizedTypeBinding)iterableType).arguments;
+						break;
+				}
+				
+				if (arguments != null && arguments.length == 1) {
+					return arguments[0];
+				}
+			}
+		}
+		
+		return null;
+	}
+	
 	public static boolean handleValForLocalDeclaration(LocalDeclaration local, BlockScope scope) {
+		boolean decomponent = false;
+		
 		if (local.type instanceof SingleTypeReference) {
 			char[] token = ((SingleTypeReference)local.type).token;
 			if (token == null || token.length != 3) return false;
@@ -192,13 +281,23 @@ public class PatchFixes {
 			try {
 				init = (Expression) initCopyField.get(local);
 			} catch (Exception e) {
-				throw new RuntimeException(e);
+			}
+		}
+		
+		if (init == null && iterableCopyField != null) {
+			try {
+				init = (Expression) iterableCopyField.get(local);
+				decomponent = true;
+			} catch (Exception e) {
 			}
 		}
 		
 		TypeReference replacement = null;
+		if (init != null && decomponent) {
+		}
+		
 		if (init != null) {
-			TypeBinding resolved = init.resolveType(scope);
+			TypeBinding resolved = decomponent ? getForEachComponentType(init, scope) : init.resolveType(scope);
 			if (resolved != null) {
 				replacement = Eclipse.makeType(resolved, local.type, false);
 			}