Cleaning up code somewhat