refactors code after pull request comments