Improve some silly code written at 1.30am