Cleaned up the code a bit, no functional changes