Address code review suggestions