[jdk20] Adds support for updated JCEnhancedForLoop

3 files changed, 96 insertions, 3 deletions

diff --git a/src/core/lombok/javac/JavacAST.java b/src/core/lombok/javac/JavacAST.java
index 0919c7d4..0e8d03d5 100644
--- a/src/core/lombok/javac/JavacAST.java
+++ b/src/core/lombok/javac/JavacAST.java
@@ -36,7 +36,6 @@ import javax.annotation.processing.Messager;
 import javax.tools.Diagnostic;
 import javax.tools.JavaFileObject;
 
-import com.sun.tools.javac.util.JCDiagnostic;
 import lombok.core.AST;
 import lombok.core.CleanupRegistry;
 import lombok.core.CleanupTask;
@@ -53,6 +52,7 @@ import com.sun.tools.javac.tree.JCTree.JCBlock;
 import com.sun.tools.javac.tree.JCTree.JCCatch;
 import com.sun.tools.javac.tree.JCTree.JCClassDecl;
 import com.sun.tools.javac.tree.JCTree.JCCompilationUnit;
+import com.sun.tools.javac.tree.JCTree.JCEnhancedForLoop;
 import com.sun.tools.javac.tree.JCTree.JCExpression;
 import com.sun.tools.javac.tree.JCTree.JCFieldAccess;
 import com.sun.tools.javac.tree.JCTree.JCIdent;
@@ -63,6 +63,7 @@ import com.sun.tools.javac.tree.JCTree.JCVariableDecl;
 import com.sun.tools.javac.tree.JCTree.JCWildcard;
 import com.sun.tools.javac.tree.TreeMaker;
 import com.sun.tools.javac.util.Context;
+import com.sun.tools.javac.util.JCDiagnostic;
 import com.sun.tools.javac.util.JCDiagnostic.DiagnosticPosition;
 import com.sun.tools.javac.util.Log;
 import com.sun.tools.javac.util.Name;
@@ -415,6 +416,18 @@ public class JavacAST extends AST<JavacAST, JavacNode, JCTree> {
 		JCANNOTATEDTYPE_FIELDS_INITIALIZED = true;
 	}
 	
+	
+	private static Field JCENHANCEDFORLOOP_VARORRECORDPATTERN_FIELD = Permit.permissiveGetField(JCEnhancedForLoop.class, "varOrRecordPattern");
+	private static JCTree getVarOrRecordPattern(JCEnhancedForLoop loop) {
+		if (JCENHANCEDFORLOOP_VARORRECORDPATTERN_FIELD == null) {
+			return loop.var;
+		}
+		try {
+			return (JCTree) JCENHANCEDFORLOOP_VARORRECORDPATTERN_FIELD.get(loop);
+		} catch (Exception ignore) {}
+		return null;
+	}
+	
 	private JavacNode buildTry(JCTry tryNode) {
 		if (setAndGetAsHandled(tryNode)) return null;
 		List<JavacNode> childNodes = new ArrayList<JavacNode>();
@@ -471,6 +484,7 @@ public class JavacAST extends AST<JavacAST, JavacNode, JCTree> {
 		if (statement instanceof JCVariableDecl) return buildLocalVar((JCVariableDecl) statement, Kind.LOCAL);
 		if (statement instanceof JCTry) return buildTry((JCTry) statement);
 		if (statement.getClass().getName().equals("com.sun.tools.javac.tree.JCTree$JCLambda")) return buildLambda(statement);
+		if (statement instanceof JCEnhancedForLoop) return buildEnhancedForLoop((JCEnhancedForLoop) statement);
 		if (setAndGetAsHandled(statement)) return null;
 		
 		return drill(statement);
@@ -500,6 +514,17 @@ public class JavacAST extends AST<JavacAST, JavacNode, JCTree> {
 		return getBodyMethods.get(c);
 	}
 	
+	private JavacNode buildEnhancedForLoop(JCEnhancedForLoop loop) {
+		if (setAndGetAsHandled(loop)) return null;
+		
+		List<JavacNode> childNodes = new ArrayList<JavacNode>();
+		addIfNotNull(childNodes, buildTree(loop.expr, Kind.STATEMENT));
+		addIfNotNull(childNodes, buildStatement(loop.body));
+		addIfNotNull(childNodes, buildTree(getVarOrRecordPattern(loop), Kind.STATEMENT));
+		
+		return putInMap(new JavacNode(this, loop, childNodes, Kind.STATEMENT));
+	}
+	
 	private JavacNode drill(JCTree statement) {
 		try {
 			List<JavacNode> childNodes = new ArrayList<JavacNode>();
diff --git a/src/core/lombok/javac/handlers/HandleVal.java b/src/core/lombok/javac/handlers/HandleVal.java
index d4fb1027..8f4b1a86 100644
--- a/src/core/lombok/javac/handlers/HandleVal.java
+++ b/src/core/lombok/javac/handlers/HandleVal.java
@@ -24,15 +24,19 @@ package lombok.javac.handlers;
 import static lombok.core.handlers.HandlerUtil.handleFlagUsage;
 import static lombok.javac.handlers.HandleDelegate.HANDLE_DELEGATE_PRIORITY;
 import static lombok.javac.handlers.JavacHandlerUtil.*;
+
+import java.lang.reflect.Field;
+
 import lombok.ConfigurationKeys;
 import lombok.val;
-import lombok.core.HandlerPriority;
 import lombok.var;
+import lombok.core.HandlerPriority;
 import lombok.javac.JavacASTAdapter;
 import lombok.javac.JavacASTVisitor;
 import lombok.javac.JavacNode;
 import lombok.javac.JavacResolution;
 import lombok.javac.ResolutionResetNeeded;
+import lombok.permit.Permit;
 import lombok.spi.Provides;
 
 import com.sun.tools.javac.code.Flags;
@@ -87,7 +91,8 @@ public class HandleVal extends JavacASTAdapter {
 		if (local.init == null) {
 			if (parentRaw instanceof JCEnhancedForLoop) {
 				JCEnhancedForLoop efl = (JCEnhancedForLoop) parentRaw;
-				if (efl.var == local) rhsOfEnhancedForLoop = efl.expr;
+				JCTree var = EnhancedForLoopReflect.getVarOrRecordPattern(efl);
+				if (var == local) rhsOfEnhancedForLoop = efl.expr;
 			}
 		}
 		
@@ -191,4 +196,19 @@ public class HandleVal extends JavacASTAdapter {
 			recursiveSetGeneratedBy(local.vartype, typeNode);
 		}
 	}
+	
+	private static class EnhancedForLoopReflect {
+		private static final Field varOrRecordPattern = Permit.permissiveGetField(JCEnhancedForLoop.class, "varOrRecordPattern");
+		
+		private static JCTree getVarOrRecordPattern(JCEnhancedForLoop loop) {
+			if (varOrRecordPattern == null) {
+				return loop.var;
+			}
+			
+			try {
+				return (JCTree) varOrRecordPattern.get(loop);
+			} catch (Exception ignore) {}
+			return null;
+		}
+	}
 }
diff --git a/src/core/lombok/javac/handlers/JavacHandlerUtil.java b/src/core/lombok/javac/handlers/JavacHandlerUtil.java
index 9d153a72..593a1aa8 100644
--- a/src/core/lombok/javac/handlers/JavacHandlerUtil.java
+++ b/src/core/lombok/javac/handlers/JavacHandlerUtil.java
@@ -512,6 +512,13 @@ public class JavacHandlerUtil {
 		case LOCAL:
 			JCVariableDecl variable = (JCVariableDecl) parentNode.get();
 			variable.mods.annotations = filterList(variable.mods.annotations, annotation.get());
+			
+			if ((variable.mods.flags & GENERATED_MEMBER) != 0) {
+				JavacNode typeNode = upToTypeNode(annotation);
+				if (isRecord(typeNode)) {
+					RecordComponentReflect.deleteAnnotation((JCClassDecl) typeNode.get(), variable, annotation.get());
+				}
+			}
 			break;
 		case METHOD:
 			JCMethodDecl method = (JCMethodDecl) parentNode.get();
@@ -559,6 +566,47 @@ public class JavacHandlerUtil {
 		return newAnnotations.toList();
 	}
 	
+	private static List<JCAnnotation> filterListByPos(List<JCAnnotation> annotations, JCTree jcTree) {
+		ListBuffer<JCAnnotation> newAnnotations = new ListBuffer<JCAnnotation>();
+		for (JCAnnotation ann : annotations) {
+			if (jcTree.pos != ann.pos) newAnnotations.append(ann);
+		}
+		return newAnnotations.toList();
+	}
+	
+	
+	static class RecordComponentReflect {
+		
+		private static final Field astField;
+		private static final Method findRecordComponentToRemove;
+		
+		static {
+			Field a = null;
+			Method m = null;
+			try {
+				Class<?> forName = Class.forName("com.sun.tools.javac.code.Symbol$RecordComponent");
+				a = Permit.permissiveGetField(forName, "ast");
+				m = Permit.permissiveGetMethod(ClassSymbol.class, "findRecordComponentToRemove", JCVariableDecl.class);
+			} catch (Throwable e) {
+				// Ignore
+			}
+			astField = a;
+			findRecordComponentToRemove = m;
+		}
+		
+		static void deleteAnnotation(JCClassDecl record, JCVariableDecl component, JCTree annotation) {
+			if (astField == null || findRecordComponentToRemove == null) return;
+			
+			try {
+				Object toRemove = Permit.invokeSneaky(findRecordComponentToRemove, record.sym, component);
+				JCVariableDecl variable = (JCVariableDecl) Permit.get(astField, toRemove);
+				variable.mods.annotations =  filterListByPos(variable.mods.annotations, annotation);
+			} catch (Throwable e) {
+				// Ignore
+			}
+		}
+	}
+	
 	/** Serves as return value for the methods that check for the existence of fields and methods. */
 	public enum MemberExistsResult {
 		NOT_EXISTS, EXISTS_BY_LOMBOK, EXISTS_BY_USER;