Made little bit of refactoring