refactoring code

little code refactor