Refactor code according to code review remarks