Refactor code

...because I have nothing better to do