Coming up with better design
Refactor for nicer design