Improve code based on PR comments.