Fixed premature removal of imports

3 files changed, 27 insertions, 24 deletions

diff --git a/src/core/lombok/core/AST.java b/src/core/lombok/core/AST.java
index e769ce34..07d6ec6c 100644
--- a/src/core/lombok/core/AST.java
+++ b/src/core/lombok/core/AST.java
@@ -30,6 +30,7 @@ import java.lang.reflect.ParameterizedType;
 import java.lang.reflect.Type;
 import java.util.ArrayList;
 import java.util.Collection;
+import java.util.Collections;
 import java.util.HashMap;
 import java.util.IdentityHashMap;
 import java.util.List;
@@ -52,12 +53,16 @@ public abstract class AST<A extends AST<A, L, N>, L extends LombokNode<A, L, N>,
 	
 	private L top;
 	private final String fileName;
+	private final String packageDeclaration;
+	private final Collection<String> imports;
 	Map<N, Void> identityDetector = new IdentityHashMap<N, Void>();
 	private Map<N, L> nodeMap = new IdentityHashMap<N, L>();
 	private boolean changed = false;
 	
-	protected AST(String fileName) {
+	protected AST(String fileName, String packageDeclaration, Collection<String> imports) {
 		this.fileName = fileName == null ? "(unknown).java" : fileName;
+		this.packageDeclaration = packageDeclaration;
+		this.imports = Collections.unmodifiableCollection(new ArrayList<String>(imports));
 	}
 	
 	protected void setChanged() {
@@ -82,14 +87,18 @@ public abstract class AST<A extends AST<A, L, N>, L extends LombokNode<A, L, N>,
 	 * 
 	 * Example: "java.util".
 	 */
-	public abstract String getPackageDeclaration();
+	public final String getPackageDeclaration() {
+		return packageDeclaration;
+	}
 	
 	/**
 	 * Return the contents of each non-static import statement on this AST's top (Compilation Unit) node.
 	 * 
 	 * Example: "java.util.IOException".
 	 */
-	public abstract Collection<String> getImportStatements();
+	public final Collection<String> getImportStatements() {
+		return imports;
+	}
 	
 	/**
 	 * Puts the given node in the map so that javac/Eclipse's own internal AST object can be translated to
diff --git a/src/core/lombok/eclipse/EclipseAST.java b/src/core/lombok/eclipse/EclipseAST.java
index 7f436ddf..7ed83bfe 100644
--- a/src/core/lombok/eclipse/EclipseAST.java
+++ b/src/core/lombok/eclipse/EclipseAST.java
@@ -56,30 +56,25 @@ public class EclipseAST extends AST<EclipseAST, EclipseNode, ASTNode> {
 	 * @param ast The compilation unit, which serves as the top level node in the tree to be built.
 	 */
 	public EclipseAST(CompilationUnitDeclaration ast) {
-		super(toFileName(ast));
+		super(toFileName(ast), packageDeclaration(ast), imports(ast));
 		this.compilationUnitDeclaration = ast;
 		setTop(buildCompilationUnit(ast));
 		this.completeParse = isComplete(ast);
 		clearChanged();
 	}
 	
-	/** {@inheritDoc} */
-	@Override public String getPackageDeclaration() {
-		CompilationUnitDeclaration cud = (CompilationUnitDeclaration) top().get();
+	private static String packageDeclaration(CompilationUnitDeclaration cud) {
 		ImportReference pkg = cud.currentPackage;
 		return pkg == null ? null : Eclipse.toQualifiedName(pkg.getImportName());
 	}
 	
-	/** {@inheritDoc} */
-	@Override public Collection<String> getImportStatements() {
+	private static Collection<String> imports(CompilationUnitDeclaration cud) { 
 		List<String> imports = new ArrayList<String>();
-		CompilationUnitDeclaration cud = (CompilationUnitDeclaration) top().get();
 		if (cud.imports == null) return imports;
 		for (ImportReference imp : cud.imports) {
 			if (imp == null) continue;
 			imports.add(Eclipse.toQualifiedName(imp.getImportName()));
 		}
-		
 		return imports;
 	}
 	
diff --git a/src/core/lombok/javac/JavacAST.java b/src/core/lombok/javac/JavacAST.java
index e231f1d8..33c167f9 100644
--- a/src/core/lombok/javac/JavacAST.java
+++ b/src/core/lombok/javac/JavacAST.java
@@ -71,7 +71,7 @@ public class JavacAST extends AST<JavacAST, JavacNode, JCTree> {
 	 * @param top The compilation unit, which serves as the top level node in the tree to be built.
 	 */
 	public JavacAST(Messager messager, Context context, JCCompilationUnit top) {
-		super(top.sourcefile == null ? null : top.sourcefile.toString());
+		super(sourceName(top), packageDeclaration(top), imports(top));
 		setTop(buildCompilationUnit(top));
 		this.context = context;
 		this.messager = messager;
@@ -81,30 +81,29 @@ public class JavacAST extends AST<JavacAST, JavacNode, JCTree> {
 		this.symtab = Symtab.instance(context);
 		clearChanged();
 	}
-	
-	public Context getContext() {
-		return context;
+
+	private static String sourceName(JCCompilationUnit cu) {
+		return cu.sourcefile == null ? null : cu.sourcefile.toString();
 	}
 	
-	/** {@inheritDoc} */
-	@Override public String getPackageDeclaration() {
-		JCCompilationUnit unit = (JCCompilationUnit)top().get();
-		return unit.pid instanceof JCFieldAccess ? unit.pid.toString() : null;
+	private static String packageDeclaration(JCCompilationUnit cu) {
+		return cu.pid instanceof JCFieldAccess ? cu.pid.toString() : null;
 	}
 	
-	/** {@inheritDoc} */
-	@Override public Collection<String> getImportStatements() {
+	private static Collection<String> imports(JCCompilationUnit cu) {
 		List<String> imports = new ArrayList<String>();
-		JCCompilationUnit unit = (JCCompilationUnit)top().get();
-		for (JCTree def : unit.defs) {
+		for (JCTree def : cu.defs) {
 			if (def instanceof JCImport) {
 				imports.add(((JCImport)def).qualid.toString());
 			}
 		}
-		
 		return imports;
 	}
 	
+	public Context getContext() {
+		return context;
+	}
+	
 	/**
 	 * Runs through the entire AST, starting at the compilation unit, calling the provided visitor's visit methods
 	 * for each node, depth first.