Refactor based on code review suggestions