[javac] defer removal of `import lombok.x;` nodes until delombok

Our test harness setup for JDK15 and up doesn't work with annotations on record 'fields' (`record X(@RightHere String foo)`).

3 files changed, 24 insertions, 10 deletions

diff --git a/src/core/lombok/javac/JavacAugments.java b/src/core/lombok/javac/JavacAugments.java
index bc23131b..e5d1fc43 100644
--- a/src/core/lombok/javac/JavacAugments.java
+++ b/src/core/lombok/javac/JavacAugments.java
@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2014 The Project Lombok Authors.
+ * Copyright (C) 2014-2021 The Project Lombok Authors.
  * 
  * Permission is hereby granted, free of charge, to any person obtaining a copy
  * of this software and associated documentation files (the "Software"), to deal
@@ -24,6 +24,7 @@ package lombok.javac;
 import lombok.core.FieldAugment;
 
 import com.sun.tools.javac.tree.JCTree;
+import com.sun.tools.javac.tree.JCTree.JCImport;
 
 public final class JavacAugments {
 	private JavacAugments() {
@@ -32,4 +33,5 @@ public final class JavacAugments {
 	
 	public static final FieldAugment<JCTree, Boolean> JCTree_handled = FieldAugment.augment(JCTree.class, boolean.class, "lombok$handled");
 	public static final FieldAugment<JCTree, JCTree> JCTree_generatedNode = FieldAugment.circularSafeAugment(JCTree.class, JCTree.class, "lombok$generatedNode");
+	public static final FieldAugment<JCImport, Boolean> JCImport_deletable = FieldAugment.circularSafeAugment(JCImport.class, Boolean.class, "lombok$deletable");
 }
diff --git a/src/core/lombok/javac/handlers/JavacHandlerUtil.java b/src/core/lombok/javac/handlers/JavacHandlerUtil.java
index b4f2acc4..ac00fb1f 100644
--- a/src/core/lombok/javac/handlers/JavacHandlerUtil.java
+++ b/src/core/lombok/javac/handlers/JavacHandlerUtil.java
@@ -103,6 +103,7 @@ import lombok.delombok.LombokOptionsFactory;
 import lombok.experimental.Accessors;
 import lombok.experimental.Tolerate;
 import lombok.javac.Javac;
+import lombok.javac.JavacAugments;
 import lombok.javac.JavacNode;
 import lombok.javac.JavacTreeMaker;
 import lombok.permit.Permit;
@@ -537,21 +538,18 @@ public class JavacHandlerUtil {
 	public static void deleteImportFromCompilationUnit(JavacNode node, String name) {
 		if (inNetbeansEditor(node)) return;
 		if (!node.shouldDeleteLombokAnnotations()) return;
-		ListBuffer<JCTree> newDefs = new ListBuffer<JCTree>();
 		
 		JCCompilationUnit unit = (JCCompilationUnit) node.top().get();
 		
 		for (JCTree def : unit.defs) {
-			boolean delete = false;
-			if (def instanceof JCImport) {
-				JCImport imp0rt = (JCImport)def;
-				delete = (!imp0rt.staticImport && imp0rt.qualid.toString().equals(name));
-			}
-			if (!delete) newDefs.append(def);
+			if (!(def instanceof JCImport)) continue;
+			JCImport imp0rt = (JCImport) def;
+			if (imp0rt.staticImport) continue;
+			if (!imp0rt.qualid.toString().equals(name)) continue;
+			JavacAugments.JCImport_deletable.set(imp0rt, true);
 		}
-		unit.defs = newDefs.toList();
 	}
-
+	
 	private static List<JCAnnotation> filterList(List<JCAnnotation> annotations, JCTree jcTree) {
 		ListBuffer<JCAnnotation> newAnnotations = new ListBuffer<JCAnnotation>();
 		for (JCAnnotation ann : annotations) {
diff --git a/src/delombok/lombok/delombok/Delombok.java b/src/delombok/lombok/delombok/Delombok.java
index 10ec5015..8d39f447 100755
--- a/src/delombok/lombok/delombok/Delombok.java
+++ b/src/delombok/lombok/delombok/Delombok.java
@@ -60,6 +60,7 @@ import javax.tools.JavaFileObject;
 import lombok.Lombok;
 import lombok.javac.CommentCatcher;
 import lombok.javac.Javac;
+import lombok.javac.JavacAugments;
 import lombok.javac.LombokOptions;
 import lombok.javac.apt.LombokProcessor;
 import lombok.permit.Permit;
@@ -69,8 +70,11 @@ import com.sun.tools.javac.comp.Todo;
 import com.sun.tools.javac.file.BaseFileManager;
 import com.sun.tools.javac.main.Arguments;
 import com.sun.tools.javac.main.JavaCompiler;
+import com.sun.tools.javac.tree.JCTree;
 import com.sun.tools.javac.tree.JCTree.JCCompilationUnit;
+import com.sun.tools.javac.tree.JCTree.JCImport;
 import com.sun.tools.javac.util.Context;
+import com.sun.tools.javac.util.ListBuffer;
 import com.zwitserloot.cmdreader.CmdReader;
 import com.zwitserloot.cmdreader.Description;
 import com.zwitserloot.cmdreader.Excludes;
@@ -795,6 +799,16 @@ public class Delombok {
 				if (verbose) feedback.printf("File: %s [%s]\n", unit.sourcefile.getName(), "unchanged (skipped)");
 				continue;
 			}
+			ListBuffer<JCTree> newDefs = new ListBuffer<JCTree>();
+			for (JCTree def : unit.defs) {
+				if (def instanceof JCImport) {
+					Boolean b = JavacAugments.JCImport_deletable.get((JCImport) def);
+					if (b == null || !b.booleanValue()) newDefs.append(def);
+				} else {
+					newDefs.append(def);
+				}
+			}
+			unit.defs = newDefs.toList();
 			if (verbose) feedback.printf("File: %s [%s%s]\n", unit.sourcefile.getName(), result.isChanged() ? "delomboked" : "unchanged", force && !options.isChanged(unit) ? " (forced)" : "");
 			Writer rawWriter;
 			if (presetWriter != null) rawWriter = createUnicodeEscapeWriter(presetWriter);