Refactoring somewhat