Code review refactor