Problem with compareTo and TreeSet

I'm having a problem with TreeSets removing a unit from a game I'm working on. I'm making a tower defense game and the path is broken up into different blocks of a set length. The blocks know the units within it and the next block on the path. When the unit leaves the bounds of the block, the block removes it from its list and adds it to the next block.

I'm using a TreeSet to keep track of the order of the units in the block, that way I can tell which one is farthest along the path. The units have a position field that keeps track of how far along the path they are, the higher the position the farther they are.

On some of my blocks I've noticed that when it tries to remove a unit from its TreeSet, the remove returns false. I've used some breakpoints and I can see that the unit is actually in the TreeSet, so I think my problem is the compareTo method for my attacking units.

Here's my code for compareTo:

public int compareTo(Object other) {
    if (other != null && AttackingUnit.class.isAssignableFrom(other.getClass())) {
        AttackingUnit o = (AttackingUnit) other;
        int amount = position - o.position;
        if (amount != 0) {
            return amount;
        } else if (amount == 0 && this == o) {
            return 0;
        }
    }
    return 1;
}

One of the blocks where I'm noticing the problem is a block that the units enter the top, take a corner halfway through and exit the right side. The block has two ArrayLists, one for units going from top to bottom (enPath) and one for the units from left to right (exPath). Heres the code where I'm having the problem:

for (int i = 0; i < exPath.size(); i++) {
    AttackingUnit unit = exPath.get(i);
    unit.stepX();
    if (unit.getX() > rightX) {
        nextBlock.addUnit(unit);
        units.remove(unit);
        exPath.remove(unit);
        i--;
    }
}

The unit is in exPath and in units (the TreeSet), but units.开发者_开发百科remove(unit) is returning false. Any ideas on how I can change compareTo on AttackingUnit to fix this?

Your compareTo method is strange. First, you should not have any elements in your list which are not of the right type, and also no null elements, since this will give problems anyway. So you could simply throw exceptions in these cases instead of returning 1.

Second, as Steve already noted, the this == o check is not right - this violates the symmetry of your relation, giving you cases where your elements are not found. This gives this simpler version:

public int compareTo(Object other) {
    AttackingUnit o = (AttackingUnit) other;
    int amount = position - o.position;
    return amount;
}

Third, make sure the positions (i.e. the results of your comparisons) do not change while a unit is in your TreeSet. If the positions must change, first remove the element from the set, change the position and then add it again.

One problem I see is in your compareTo is:

} else if (amount == 0 && this == o) {

You should be ORing them (or get rid of the this == o check). As it is now, two different AttackingUnit instances with a same position will return as 1 (first one larger). This will definitely give inconsistent ordering in the tree set.

By the way, you can replace:

if (other != null && AttackingUnit.class.isAssignableFrom(other.getClass()))

with

if (other instanceof AttackingUnit)

which is easier to read.