Improve / clean some code