Fixed previous refactoring