Implement suggested changes from code review