Implement code changes from PR discussion