Code cleanup in response to PR comments