(refactor) Minor refactoring