Adjust code according to pull request feedback