clean the code slightly