Some more refactoring... Not sure how I like this last bit