Code refactoring..