Touch of refactoring... little less code