Shortened the code a bit