Code changes after code review