记TreeSet在添加pojo对象遇到ClassCastException及解决方案

场景

在实际业务中有保证集合元素唯一且有序的需要，稍加思索后决定使用TreeSet，TreeSet底层实现为TreeMap而TreeMap底层实现为二叉树，从而可以确保TreeSet元素有序且唯一。

模拟异常（java.lang.ClassCastException）

// 女朋友实体
public class Girlfriend {
    
    private String name;

    private int age;

    public String getName() {
        return name;
    }

    public void setName(String name) {
        this.name = name;
    }

    public int getAge() {
        return age;
    }

    public void setAge(int age) {
        this.age = age;
    }
}


public static void main(String[] args)  {
    baseClassAdd();
    projectClassAdd();
}

static void baseClassAdd(){
    TreeSet<String> set = new TreeSet<>();
    set.add("stringValue");
}

static void projectClassAdd(){
    Set<Girlfriend> gfSet = new TreeSet<Girlfriend>();
    Girlfriend gf1 = new Girlfriend();
    gf1.setAge(20);
    gf1.setName("王冰冰");
    gfSet.add(gf1);
}

当执行main方法时日志报错如下，类型转换异常Girlfriend无法转换为Comparable，根据代码执行顺序可以看出同样的TreeSet在添加基本数据类型时并不会报错，接下来我们带着问题简单看下TreeSet的源码

Exception in thread "main" java.lang.ClassCastException: class com.company.Girlfriend cannot be cast to class java.lang.Comparable (com.company.Girlfriend is in unnamed module of loader 'app'; java.lang.Comparable is in module java.base of loader 'bootstrap')
at java.base/java.util.TreeMap.compare(TreeMap.java:1291)
at java.base/java.util.TreeMap.put(TreeMap.java:536)
at java.base/java.util.TreeSet.add(TreeSet.java:255)
at com.company.Main.projectClassAdd(Main.java:25)
at com.company.Main.main(Main.java:12)

简读TreeSet和TreeMap源码

先看一个泳道图，简单说明了TreeSet在add添加元素时发生了什么，TreeSet底层实现为TreeMap，所以在add元素相当于TreeMap的put元素，而TreeMap在put中又调用compare方法，下面贴部分TreeSet和TreeMap的源码帮助我们分析问题所在

//TreeSet部分源码
public class TreeSet<E> extends AbstractSet<E> implements NavigableSet<E>, Cloneable, Serializable {
    private transient NavigableMap<E, Object> m;

    TreeSet(NavigableMap<E, Object> m) {
        this.m = m;
    }
	
    //空参构造
    public TreeSet() {
        this((NavigableMap)(new TreeMap()));
    }

    //参数为Comparator的构造方法
    public TreeSet(Comparator<? super E> comparator) {
        this((NavigableMap)(new TreeMap(comparator)));
    }

}

//TreeMap部门源码
public class TreeMap<K, V> extends AbstractMap<K, V> implements NavigableMap<K, V>, Cloneable, Serializable {
    
    //一般使用TreeSet时，我们习惯直接new TreeSet()，所以comparator大部分时候都为null
    private final Comparator<? super K> comparator;
    
    //参数为Comparator的构造方法
    public TreeMap(Comparator<? super K> comparator) {
        this.comparator = comparator;
    }
    
    final int compare(Object k1, Object k2) {
        //当comparator为null时，调用Comparable接口的compareTo，这就是为什么在TreeSet在添加对象时会报错的原因
        //当comparator不为null时，使用comparator调用compare
        return this.comparator == null ? ((Comparable)k1).compareTo(k2) : this.comparator.compare(k1, k2);
    }
    
    public V put(K key, V value) {
        TreeMap.Entry<K, V> t = this.root;
        if (t == null) {
            //调用compare方法
            this.compare(key, key);
            this.root = new TreeMap.Entry(key, value, (TreeMap.Entry)null);
            this.size = 1;
            ++this.modCount;
            return null;
        } else {
            Comparator<? super K> cpr = this.comparator;
            int cmp;
            TreeMap.Entry parent;
            if (cpr != null) {
                do {
                    parent = t;
                	//调用compare方法
                    cmp = cpr.compare(key, t.key);
                    if (cmp < 0) {
                        t = t.left;
                    } else {
                        if (cmp <= 0) {
                            return t.setValue(value);
                        }

                        t = t.right;
                    }
                } while(t != null);
            } else {
                if (key == null) {
                    throw new NullPointerException();
                }
                Comparable k = (Comparable)key;
                do {
                    parent = t;
                    //当cpr为null时，通过Comparable调用compareTo方法
                    cmp = k.compareTo(t.key);
                    if (cmp < 0) {
                        t = t.left;
                    } else {
                        if (cmp <= 0) {
                            return t.setValue(value);
                        }

                        t = t.right;
                    }
                } while(t != null);
            }

            TreeMap.Entry<K, V> e = new TreeMap.Entry(key, value, parent);
            if (cmp < 0) {
                parent.left = e;
            } else {
                parent.right = e;
            }

            this.fixAfterInsertion(e);
            ++this.size;
            ++this.modCount;
            return null;
        }
    }

}

//String部分源码，已经实现了Comparable接口
public final class String implements Serializable, Comparable<String>, CharSequence {

}

结论

通过以上的部分源码和注释，相信大家已经有了答案，诸如String之类的基本数据类型已经实现了Comparable方法，所以在TreeSet在add基本数据类型值时并不会报错，而当add我们的pojo对象时由于我们没有实现Comparable接口，在TreeMap的compare方法中((Comparable)k1).compareTo(k2) 时就会抛出ClassCastException，所以在TreeSet中add pojo时请务必实现Comparable接口

同时观察TreeSet和TreeMap的构造方法，都存在参数为Comparator的构造方法，这就提供了另一个解决ClassCastException的方案，就是使用Comparator构造器，匿名内部类重写compare方法

Comparator<Girlfriend> cpr = new Comparator<Girlfriend>() {
    @Override
    public int compare(Girlfriend o, Girlfriend t1) {
        return 0;
    }
};
TreeSet<Girlfriend> set = new TreeSet<>(cpr);
set.add(new Girlfriend());