Cleaning up code based on final comments in #43