Bit of code clean up and maybe slight optimisation