Did a little more refactoring