Cleaning up some existing code