Refactor of code