If I understand your problem correctly it looks like a good candidate for a Visitor design pattern.
public interface IElementVisitor {
public void visit(Root root);
public void visit(Node node);
public void visit(Leaf leaf);
}
public interface IElement {
public void accept(IElementVisitor visitor);
}
Now the visitable objects Root, Node and Leaf will implement IElement (hate the reference to a tree data structure here btw)
public class Root implements IElement {
// The guts of the class
public void accept(IElementVisitor visitor) {
visitor.visit(this);
}
}
The code is identical for the Node and Leaf classes. So the impact to your data models is a single line of code. In the CustomTreeModel you would do something like this:
public Object getChild(Object parent, int index) {
IElement element = (IElement) parent;
element.visit( new IElementVisitor() {
@Override
public void visit(Root root) {
// whatever needs doing if this is a Root
}
@Override
public void visit(Node node) {
// whatever needs doing if this is a Node
}
@Override
public void visit(Leaf leaf) {
// whatever needs doing if this is a Leaf
}
} );
// and so on for the other methods
There will probably be a little messing around with the final modifier in this example but here is what I love about Visitors and why I think it's a good option for you:
- Separation of concerns amongst a set of related classes. Your Root, Node and Leaf classes are related but need to do different things in this particular circumstance. This will defer the things that need doing to them to the classes that use them.
- Syntax errors instead of run time errors when adding new classes. Lets say in a few months time you want to add a fourth level. All you need to do is add another method to the IElementVisitor and every visitation will light up red. In your posted code it will result in IllegalArgumentException that need hunting down.
- Reusability. Are there any other places where the algorithm is independent of the data structure? Don't add behaviour to the data, visit it and keep the controller in charge of behaviour. This is the heart of the Open/Closed principal as it allows a set of related classes (Root, Leaf and Node) to have behaviour assigned to them without any changes to their source code.