Refactored the type resolver, and made the 'find locally named classes that shadow' feature way better; previously it would only find very few locally named classes (your own class and any parent outers, that was it, no siblings).

6 files changed, 111 insertions, 63 deletions

diff --git a/src/core/lombok/core/TypeLibrary.java b/src/core/lombok/core/TypeLibrary.java
index cabffdde..8a03b479 100644
--- a/src/core/lombok/core/TypeLibrary.java
+++ b/src/core/lombok/core/TypeLibrary.java
@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2009 The Project Lombok Authors.
+ * Copyright (C) 2009-2011 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
@@ -21,24 +21,49 @@
  */
 package lombok.core;
 
+import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.Collection;
 import java.util.Collections;
 import java.util.HashMap;
-import java.util.HashSet;
+import java.util.List;
 import java.util.Map;
-import java.util.Set;
 
 /**
  * Library of types, which can be used to look up potential matching types.
  * 
- * For example, if you put 'foo.Spork' and 'bar.Spork' into the library, and then ask for
- * all compatible types given the type 'Spork', you'll get both of them, but you'll only
- * get the one if you ask for compatible types given 'foo.Spork'.
- * 
- * Useful to 'guess' if a given annotation AST node matches an annotation handler's target annotation.
+ * For example, if you put {@code foo.Spork} and {@code bar.Spork} into the library, and then ask for
+ * all compatible types given the type {@code Spork}, you'll get both of them, but you'll only
+ * get the one if you ask for compatible types given {@code foo.Spork}.
+ * <p>
+ * When adding {@code foo.Spork}, that FQN (Fully Qualified Name) will be returned as an option for any of these queries:
+ * <ul><li>foo.Spork</li><li>Spork</li><li>foo.*</li></ul>
  */
 public class TypeLibrary {
-	private final Map<String, Set<String>> simpleToQualifiedMap = new HashMap<String, Set<String>>();
+	private final Map<String, List<String>> keyToFqnMap;
+	private final String singletonValue;
+	private final List<String> singletonKeys;
+	
+	public TypeLibrary() {
+		keyToFqnMap = new HashMap<String, List<String>>();
+		singletonKeys = null;
+		singletonValue = null;
+	}
+	
+	private TypeLibrary(String fqnSingleton) {
+		keyToFqnMap = null;
+		singletonValue = fqnSingleton;
+		int idx = fqnSingleton.lastIndexOf('.');
+		if (idx == -1) {
+			singletonKeys = Collections.singletonList(fqnSingleton);
+		} else {
+			singletonKeys = Arrays.asList(fqnSingleton, fqnSingleton.substring(idx + 1), fqnSingleton.substring(0, idx) + ".*");
+		}
+	}
+	
+	public static TypeLibrary createLibraryForSingleType(String fqnSingleton) {
+		return new TypeLibrary(fqnSingleton);
+	}
 	
 	/**
 	 * Add a type to the library.
@@ -46,6 +71,7 @@ public class TypeLibrary {
 	 * @param fullyQualifiedTypeName the FQN type name, such as 'java.lang.String'.
 	 */
 	public void addType(String fullyQualifiedTypeName) {
+		if (keyToFqnMap == null) throw new IllegalStateException("SingleType library");
 		int idx = fullyQualifiedTypeName.lastIndexOf('.');
 		if (idx == -1) throw new IllegalArgumentException(
 				"Only fully qualified types are allowed (and stuff in the default package is not palatable to us either!)");
@@ -53,27 +79,33 @@ public class TypeLibrary {
 		final String simpleName = fullyQualifiedTypeName.substring(idx +1);
 		final String packageName = fullyQualifiedTypeName.substring(0, idx);
 		
-		if (simpleToQualifiedMap.put(fullyQualifiedTypeName, Collections.singleton(fullyQualifiedTypeName)) != null) return;
+		if (keyToFqnMap.put(fullyQualifiedTypeName, Collections.singletonList(fullyQualifiedTypeName)) != null) return;
 		
 		addToMap(simpleName, fullyQualifiedTypeName);
 		addToMap(packageName + ".*", fullyQualifiedTypeName);
 	}
 	
 	private TypeLibrary addToMap(String keyName, String fullyQualifiedTypeName) {
-		Set<String> existing = simpleToQualifiedMap.get(keyName);
-		Set<String> set = (existing == null) ? new HashSet<String>() : new HashSet<String>(existing);
-		set.add(fullyQualifiedTypeName);
-		simpleToQualifiedMap.put(keyName, Collections.unmodifiableSet(set));
+		List<String> list = keyToFqnMap.get(keyName);
+		if (list == null) {
+			list = new ArrayList<String>();
+			keyToFqnMap.put(keyName, list);
+		}
+		
+		list.add(fullyQualifiedTypeName);
 		return this;
 	}
 	
 	/**
 	 * Returns all items in the type library that may be a match to the provided type.
 	 * 
-	 * @param typeReference something like 'String' or even 'java.lang.String'.
+	 * @param typeReference something like 'String', 'java.lang.String', or 'java.lang.*'.
+	 * @return A list of Fully Qualified Names for all types in the library that fit the reference.
 	 */
 	public Collection<String> findCompatible(String typeReference) {
-		Set<String> result = simpleToQualifiedMap.get(typeReference);
-		return result == null ? Collections.<String>emptySet() : result;
+		if (singletonKeys != null) return singletonKeys.contains(typeReference) ? Collections.singletonList(singletonValue) : Collections.<String>emptyList();
+		
+		List<String> result = keyToFqnMap.get(typeReference);
+		return result == null ? Collections.<String>emptyList() : Collections.unmodifiableList(result);
 	}
 }
diff --git a/src/core/lombok/core/TypeResolver.java b/src/core/lombok/core/TypeResolver.java
index ecc042a0..2fa24fd2 100644
--- a/src/core/lombok/core/TypeResolver.java
+++ b/src/core/lombok/core/TypeResolver.java
@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2009 The Project Lombok Authors.
+ * Copyright (C) 2009-2011 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
@@ -30,17 +30,17 @@ import lombok.core.AST.Kind;
 
 /**
  * Capable of resolving a simple type name such as 'String' into 'java.lang.String'.
+ * <p>
+ * NB: This resolver gives wrong answers when there's a class in the local package with the same name as a class in a star-import,
+ *  and this importer also can't find inner types from superclasses/interfaces.
  */
 public class TypeResolver {
-	private final TypeLibrary library;
 	private Collection<String> imports;
 	
 	/**
-	 * Creates a new TypeResolver that can be used to resolve types in a given library, encountered in
-	 * a source file with the provided package and import statements.
+	 * Creates a new TypeResolver that can be used to resolve types in a source file with the given package and import statements.
 	 */
-	public TypeResolver(TypeLibrary library, String packageString, Collection<String> importStrings) {
-		this.library = library;
+	public TypeResolver(String packageString, Collection<String> importStrings) {
 		this.imports = makeImportList(packageString, importStrings);
 	}
 	
@@ -51,41 +51,76 @@ public class TypeResolver {
 		return imports;
 	}
 	
+	public boolean typeMatches(LombokNode<?, ?, ?> context, String fqn, String typeRef) {
+		return !findTypeMatches(context, TypeLibrary.createLibraryForSingleType(fqn), typeRef).isEmpty();
+	}
+	
 	/**
 	 * Finds type matches for the stated type reference. The provided context is scanned for local type names
 	 * that shadow type names listed in import statements. If such a shadowing occurs, no matches are returned
 	 * for any shadowed types, as you would expect.
 	 */
-	public Collection<String> findTypeMatches(LombokNode<?, ?, ?> context, String typeRef) {
+	public Collection<String> findTypeMatches(LombokNode<?, ?, ?> context, TypeLibrary library, String typeRef) {
+		// When asking if 'Foo' could possibly  be referring to 'bar.Baz', the answer is obviously no.
 		Collection<String> potentialMatches = library.findCompatible(typeRef);
 		if (potentialMatches.isEmpty()) return Collections.emptyList();
 		
+		// If input type appears to be fully qualified, we found a winner.
 		int idx = typeRef.indexOf('.');
 		if (idx > -1) return potentialMatches;
-		String simpleName = typeRef.substring(idx+1);
 		
-		//If there's an import statement that explicitly imports a 'Getter' that isn't any of our potentials, return no matches.
-		if (nameConflictInImportList(simpleName, potentialMatches)) return Collections.emptyList();
+		// If there's an import statement that explicitly imports a 'Getter' that isn't any of our potentials, return no matches,
+		// because if you want to know if 'Foo' could refer to 'bar.Foo' when 'baz.Foo' is explicitly imported, the answer is no.
+		if (nameConflictInImportList(typeRef, potentialMatches)) return Collections.emptyList();
 		
-		//Check if any of our potentials is even imported in the first place. If not: no matches.
-		potentialMatches = eliminateImpossibleMatches(potentialMatches);
+		// Check if any of our potentials is even imported in the first place. If not: no matches.
+		// Note that (ourPackage.*) is added to the imports.
+		potentialMatches = eliminateImpossibleMatches(potentialMatches, library);
 		if (potentialMatches.isEmpty()) return Collections.emptyList();
 		
-		//Find a lexically accessible type of the same simple name in the same Compilation Unit. If it exists: no matches.
+		// Now the hard part - inner classes or method local classes in our own scope.
+		// For method locals, this refers to any statements that are 'above' the type reference with the same name.
+		// For inners, this refers to siblings of us or any parent node that are type declarations.
 		LombokNode<?, ?, ?> n = context;
-		while (n != null) {
+		
+		mainLoop:
+		while (n != null && n.getKind() != Kind.COMPILATION_UNIT) {
+			if (n.getKind() == Kind.TYPE && typeRef.equals(n.getName())) {
+				// Our own class or one of our outer classes is named 'typeRef' so that's what 'typeRef' is referring to, not one of our type library classes.
+				return Collections.emptyList();
+			}
+			
+			if (n.getKind() == Kind.STATEMENT || n.getKind() == Kind.LOCAL) {
+				LombokNode<?, ?, ?> newN = n.directUp();
+				if (newN == null) break mainLoop;
+				
+				if (newN.getKind() == Kind.STATEMENT || newN.getKind() == Kind.INITIALIZER || newN.getKind() == Kind.METHOD) {
+					for (LombokNode<?, ?, ?> child : newN.down()) {
+						// We found a method local with the same name above our code. That's the one 'typeRef' is referring to, not
+						// anything in the type library we're trying to find, so, no matches.
+						if (child.getKind() == Kind.TYPE && typeRef.equals(child.getName())) return Collections.emptyList();
+						if (child == n) break;
+					}
+				}
+				n = newN;
+				continue mainLoop;
+			}
+			
 			if (n.getKind() == Kind.TYPE) {
-				String name = n.getName();
-				if (name != null && name.equals(simpleName)) return Collections.emptyList();
+				for (LombokNode<?, ?, ?> child : n.down()) {
+					// Inner class that's visible to us has 'typeRef' as name, so that's the one being referred to, not one of our type library classes.
+					if (child.getKind() == Kind.TYPE && typeRef.equals(child.getName())) return Collections.emptyList();
+				}
 			}
-			n = n.up();
+			
+			n = n.directUp();
 		}
 		
-		// The potential matches we found by comparing the import statements is our matching set. Return it.
+		// No class in this source file is a match, therefore the potential matches found via the import statements must be it. Return those.
 		return potentialMatches;
 	}
 	
-	private Collection<String> eliminateImpossibleMatches(Collection<String> potentialMatches) {
+	private Collection<String> eliminateImpossibleMatches(Collection<String> potentialMatches, TypeLibrary library) {
 		Set<String> results = new HashSet<String>();
 		
 		for (String importedType : imports) {
diff --git a/src/core/lombok/eclipse/HandlerLibrary.java b/src/core/lombok/eclipse/HandlerLibrary.java
index 197d7a7e..5c75c071 100644
--- a/src/core/lombok/eclipse/HandlerLibrary.java
+++ b/src/core/lombok/eclipse/HandlerLibrary.java
@@ -173,11 +173,11 @@ public class HandlerLibrary {
 		String pkgName = annotationNode.getPackageDeclaration();
 		Collection<String> imports = annotationNode.getImportStatements();
 		
-		TypeResolver resolver = new TypeResolver(typeLibrary, pkgName, imports);
+		TypeResolver resolver = new TypeResolver(pkgName, imports);
 		TypeReference rawType = annotation.type;
 		if (rawType == null) return;
 		
-		for (String fqn : resolver.findTypeMatches(annotationNode, toQualifiedName(annotation.type.getTypeName()))) {
+		for (String fqn : resolver.findTypeMatches(annotationNode, typeLibrary, toQualifiedName(annotation.type.getTypeName()))) {
 			boolean isPrintAST = fqn.equals(PrintAST.class.getName());
 			if (isPrintAST == skipPrintAst) continue;
 			AnnotationHandlerContainer<?> container = annotationHandlers.get(fqn);
diff --git a/src/core/lombok/eclipse/handlers/EclipseHandlerUtil.java b/src/core/lombok/eclipse/handlers/EclipseHandlerUtil.java
index 5130c7de..c6d7eb41 100644
--- a/src/core/lombok/eclipse/handlers/EclipseHandlerUtil.java
+++ b/src/core/lombok/eclipse/handlers/EclipseHandlerUtil.java
@@ -30,7 +30,6 @@ import java.lang.reflect.Method;
 import java.lang.reflect.Modifier;
 import java.util.ArrayList;
 import java.util.Arrays;
-import java.util.Collection;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
@@ -46,7 +45,6 @@ import lombok.core.AnnotationValues.AnnotationValue;
 import lombok.eclipse.EclipseAST;
 import lombok.eclipse.EclipseNode;
 import lombok.core.TransformationsUtil;
-import lombok.core.TypeLibrary;
 import lombok.core.TypeResolver;
 
 import org.eclipse.core.runtime.ILog;
@@ -260,16 +258,9 @@ public class EclipseHandlerUtil {
 		if (!lastPartA.equals(lastPartB)) return false;
 		String typeName = toQualifiedName(typeRef.getTypeName());
 		
-		TypeLibrary library = new TypeLibrary();
-		library.addType(type.getName());
-		TypeResolver resolver = new TypeResolver(library, node.getPackageDeclaration(), node.getImportStatements());
-		Collection<String> typeMatches = resolver.findTypeMatches(node, typeName);
+		TypeResolver resolver = new TypeResolver(node.getPackageDeclaration(), node.getImportStatements());
+		return resolver.typeMatches(node, type.getName(), typeName);
 		
-		for (String match : typeMatches) {
-			if (match.equals(type.getName())) return true;
-		}
-		
-		return false;
 	}
 	
 	public static Annotation copyAnnotation(Annotation annotation, ASTNode source) {
diff --git a/src/core/lombok/javac/HandlerLibrary.java b/src/core/lombok/javac/HandlerLibrary.java
index 17684a7e..cba318e2 100644
--- a/src/core/lombok/javac/HandlerLibrary.java
+++ b/src/core/lombok/javac/HandlerLibrary.java
@@ -171,9 +171,9 @@ public class HandlerLibrary {
 	 * @param annotation 'node.get()' - convenience parameter.
 	 */
 	public void handleAnnotation(JCCompilationUnit unit, JavacNode node, JCAnnotation annotation) {
-		TypeResolver resolver = new TypeResolver(typeLibrary, node.getPackageDeclaration(), node.getImportStatements());
+		TypeResolver resolver = new TypeResolver(node.getPackageDeclaration(), node.getImportStatements());
 		String rawType = annotation.annotationType.toString();
-		for (String fqn : resolver.findTypeMatches(node, rawType)) {
+		for (String fqn : resolver.findTypeMatches(node, typeLibrary, rawType)) {
 			boolean isPrintAST = fqn.equals(PrintAST.class.getName());
 			if (isPrintAST && phase != 2) continue;
 			if (!isPrintAST && phase == 2) continue;
diff --git a/src/core/lombok/javac/handlers/JavacHandlerUtil.java b/src/core/lombok/javac/handlers/JavacHandlerUtil.java
index fc9435d8..821415d2 100644
--- a/src/core/lombok/javac/handlers/JavacHandlerUtil.java
+++ b/src/core/lombok/javac/handlers/JavacHandlerUtil.java
@@ -27,7 +27,6 @@ import java.lang.annotation.Annotation;
 import java.lang.reflect.Method;
 import java.lang.reflect.Modifier;
 import java.util.ArrayList;
-import java.util.Collection;
 import java.util.HashMap;
 import java.util.Map;
 import java.util.WeakHashMap;
@@ -37,7 +36,6 @@ import lombok.AccessLevel;
 import lombok.Data;
 import lombok.Getter;
 import lombok.core.AnnotationValues;
-import lombok.core.TypeLibrary;
 import lombok.core.TypeResolver;
 import lombok.core.AST.Kind;
 import lombok.core.AnnotationValues.AnnotationValue;
@@ -137,16 +135,8 @@ public class JavacHandlerUtil {
 	public static boolean typeMatches(Class<?> type, JavacNode node, JCTree typeNode) {
 		String typeName = typeNode.toString();
 		
-		TypeLibrary library = new TypeLibrary();
-		library.addType(type.getName());
-		TypeResolver resolver = new TypeResolver(library, node.getPackageDeclaration(), node.getImportStatements());
-		Collection<String> typeMatches = resolver.findTypeMatches(node, typeName);
-		
-		for (String match : typeMatches) {
-			if (match.equals(type.getName())) return true;
-		}
-		
-		return false;
+		TypeResolver resolver = new TypeResolver(node.getPackageDeclaration(), node.getImportStatements());
+		return resolver.typeMatches(node, type.getName(), typeName);
 	}
 	
 	/**