Cleaned up the code slightly.