Code clean up!