Code refactoring
Small refactoring