TreeSet not always correctly sorted for "Uncomparable" exceptions

src/main/java/org/springframework/classify/SubclassClassifier.java

@@ -15,12 +15,8 @@
  */
 package org.springframework.classify;
 
-import java.io.Serializable;
-import java.util.Comparator;
 import java.util.HashMap;
 import java.util.Map;
-import java.util.Set;
-import java.util.TreeSet;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentMap;
 
@@ -113,17 +109,21 @@ public class SubclassClassifier<T, C> implements Classifier<T, C> {
 		}
 
 		// check for subclasses
-		Set<Class<? extends T>> classes = new TreeSet<Class<? extends T>>(new ClassComparator());
-		classes.addAll(classified.keySet());
-		for (Class<? extends T> cls : classes) {
-			if (cls.isAssignableFrom(exceptionClass)) {
-				C value = classified.get(cls);
-				this.classified.put(exceptionClass, value);
-				return value;
-			}
+		C value = null;
+		for (Class<?> cls = exceptionClass; !cls.equals(Object.class) && value == null; cls = cls.getSuperclass()) {
+			value = classified.get(cls);
 		}
 
-		return defaultValue;
+		if (value == null) {
+			value = defaultValue;
+		}
+
+		//ConcurrentHashMap doesn't allow nulls
+		if (value != null) {
+			this.classified.put(exceptionClass, value);
+		}
+
+		return value;
 	}
 
 	/**
@@ -138,24 +138,4 @@ public class SubclassClassifier<T, C> implements Classifier<T, C> {
 	protected Map<Class<? extends T>, C> getClassified() {
 		return classified;
 	}
-
-	/**
-	 * Comparator for classes to order by inheritance.
-	 *
-	 * @author Dave Syer
-	 *
-	 */
-	private static class ClassComparator implements Comparator<Class<?>>, Serializable {
-		/**
-		 * @return 1 if arg0 is assignable from arg1, -1 otherwise
-		 * @see java.util.Comparator#compare(java.lang.Object, java.lang.Object)
-		 */
-		public int compare(Class<?> arg0, Class<?> arg1) {
-			if (arg0.isAssignableFrom(arg1)) {
-				return 1;
-			}
-			return -1;
-		}
-	}
-
 }

src/test/java/org/springframework/classify/SubclassExceptionClassifierTests.java

@@ -16,14 +16,17 @@
 
 package org.springframework.classify;
 
-import static org.junit.Assert.assertEquals;
-import static org.junit.Assert.assertNull;
+import org.junit.Test;
 
+import java.io.FileNotFoundException;
+import java.net.ConnectException;
+import java.net.SocketException;
 import java.util.Collections;
 import java.util.HashMap;
+import java.util.NoSuchElementException;
 
-import org.junit.Test;
-import org.springframework.classify.SubclassClassifier;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNull;
 
 public class SubclassExceptionClassifierTests {
 
@@ -83,4 +86,26 @@ public class SubclassExceptionClassifierTests {
 		});
 		assertEquals("spam", classifier.classify(new IllegalStateException("Foo")));
 	}
+
+	@SuppressWarnings("serial")
+	@Test
+	public void testClassifyAncestorMatch2() {
+		classifier = new SubclassClassifier<Throwable, String>();
+		classifier.setTypeMap(new HashMap<Class<? extends Throwable>, String>() {
+			{
+				put(SocketException.class, "1");
+				put(FileNotFoundException.class, "buz");
+				put(NoSuchElementException.class, "buz");
+				put(ArrayIndexOutOfBoundsException.class, "buz");
+				put(IllegalArgumentException.class, "bar");
+				put(RuntimeException.class, "spam");
+				put(ConnectException.class, "2");
+			}
+		});
+		assertEquals("2", classifier.classify(new SubConnectException()));
+	}
+
+	public static class SubConnectException extends ConnectException {
+
+	}
 }