Is there a better and cleaner way to write this two functions?

Question

So I am struggling to write this code in a clean way. I do not know if there is a better way.

private function1(Collection<D> collection) {
    for (Iterator it = collection.iterator(); it.hasNext(); ) {
        Object object = it.next();
        switch (object .getClass().getSimpleName()) {
            case "A": do some work with class A
                break;
            case "B": do some work with class B
                break; 
            case "C": do some work with class C
                break;  
        }
    }
}

So I get a collection which I iterate. But the collection can be three different classes in the collection since class A,B,C are of the parent class D. I think my code is not clean and I am looking for a good way to write this. Here is another example that I have which is a bit different.

private function2(Collection<A,B,C> collection) {
    for (Iterator it = collection.iterator(); it.hasNext(); ) {
        Object object = it.next();
        switch (object .getClass().getSimpleName()) {
            case "A": do some work with class A
                break;
            case "B": do some work with class B
                break; 
            case "C": do some work with class C
                break;  
        }
    }
}

In this function I can get either a collection of class A of class B or of class C. I though about making three separate functions for each class. But than I would have code duplicates. But I don't know if it would be actually better to split function2 into function2A, function2B, function2C.

Is there a better and cleaner way to implement those two functions. Classes A,B,C and D are from a framework. So I am unable to edit them.

Answer 1

You can create an interface that A, B, and C can implement from. In your interface you can declare a method that will 'do the work' and A,B, and C can have their own implementations.

interface MyInterface{
  void doTheWork();
}

Then do this for each class

class A implements MyInterface { 
    void doTheWork() { 
        ...
    } 
}

Now you can just loop through the collection without having to check what class the instance is of

---

EDIT - If you cannot modify A,B, or C than you can use a map with the class as the key and a lambda as the value.

//pseudocode
Map<Class<?>,Runnable> lambdaClassMap = new HashMap<>();
lambdaClassMap.put(A.class, () -> doAWork());
lambdaClassMap.put(B.class, () -> doBWork());
lambdaClassMap.put(C.class, () -> doCWork());

for(...) {
    Object object = it.next();
    lambdaClassMap.get(object.getClass()).run();
}

---

EDIT #2 - So since you want to access your object, instead of using Runnable you can use Function.

// change the function do what you want
// this one accepts the A object and calls an imaginary method 'getSize()'
Function<A, Integer> funcA = x -> x.getSize();

Map<Class<?>,Function> functionClassMap = new HashMap<>();
functionClassMap.put(A.class, funcA));

Now, you can use the apply() method on your function to access your object directly when you are looping.

Integer size = functionClassMap.get(object.getClass()).apply(object);