Beautified code a bit